Deploy MOFScreen-Agent FastAPI backend

Dockerfile

@@ -13,6 +13,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 
 COPY requirements.txt /app/requirements.txt
 RUN python -m pip install --upgrade pip && \
+    python -m pip install --no-cache-dir --index-url https://download.pytorch.org/whl/cpu torch==2.6.0 && \
     python -m pip install --no-cache-dir -r /app/requirements.txt
 
 COPY agent /app/agent
@@ -21,6 +22,7 @@ COPY configs /app/configs
 COPY data /app/data
 COPY tools /app/tools
 COPY trained_toxicity_models /app/trained_toxicity_models
+COPY coprinet /app/coprinet
 
 EXPOSE 7860
 

coprinet/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 rsanchezgarc
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

coprinet/data/models/trained_inStock/lightning_logs/version_0/checkpoints/epoch=238-step=1295857.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ef10cad50629f328a2302d0301c50efbf2666d091692a4382beef8a18dfab152
+size 25121165

coprinet/pricePrediction/ArgParser_base.py

@@ -0,0 +1,112 @@
+import abc
+import argparse
+import inspect
+from argparse import ArgumentParser
+
+from pricePrediction import config
+
+
+class ArgParseable(abc.ABC):
+
+    @property
+    @abc.abstractmethod
+    def DESIRED_PARAMS_TO_ASK(self): #TOOD: Properties always apply to instances, not classes. to be modified
+        raise NotImplementedError()
+
+    @classmethod
+    def fromTypeStrToType(cls, strType: str):
+        if strType == "str":
+            return str
+        elif strType == "float":
+            return float
+        elif strType == "int":
+            return int
+        elif strType == "bool":
+            return bool
+        else:
+            raise ValueError("Parser not recognized type: "+str(strType))
+
+    @classmethod
+    def fromDefaultStrToValue(cls, argname:str, typeFun:type, defaultStr: str): #Currently not used
+        '''
+        TO BE USED if  defaults in the docstring E.g.    :param int batch_size: Batch size. defaults to %(config.BATCH_SIZE)s
+
+        :param argname:
+        :param typeFun:
+        :param defaultStr:
+        :return:
+        '''
+
+        if defaultStr is None or defaultStr == "None":
+            return None
+        return typeFun(defaultStr)
+
+    @classmethod
+    def getArgsForParser(cls):
+        from docstring_parser import parse
+        docstring = cls.__init__.__doc__
+        config_vars = vars(config)
+        config_vars = { "config."+key: val for key,val in config_vars.items() }
+        docstring = docstring%config_vars
+        docstring = parse(docstring)
+
+        signature = inspect.signature(cls.__init__)
+        name_to_default = {
+            k: v.default
+            for k, v in signature.parameters.items()
+            if v.default is not inspect.Parameter.empty
+        }
+
+        params = []
+        for elem in docstring.params:
+            if elem.arg_name in cls.DESIRED_PARAMS_TO_ASK:
+                typeFun = cls.fromTypeStrToType(elem.type_name)
+                # print(elem.default, elem.description)
+                ### default = cls.fromDefaultStrToValue(elem.arg_name, typeFun, elem.default)
+                default = name_to_default[elem.arg_name]
+                params.append((elem.arg_name, typeFun, default, elem.description))
+        return params
+
+    # @classmethod
+    # def _addParamToArgParse(cls, parser:ArgumentParser, paramTuple):
+    #     name, typeFun, default, help= paramTuple
+    #     if typeFun == bool:
+    #         if default is True:
+    #             name = "no_"+name
+    #             help = "Deactivate "+help
+    #         parser.add_argument("--" + name, help=help+" Default=%(default)s", action="store_true")
+    #     else:
+    #         parser.add_argument("--"+name, type=typeFun, help=help+" Default=%(default)s", default=default)
+    #     return parser
+
+    @classmethod
+    def addParamsToArgParse(cls, parser: ArgumentParser):
+        '''
+
+        :param parser: Argparser to which methods will be added
+        :return: the input parser but with added arguments
+        '''
+        for paramTuple in cls.getArgsForParser():
+            name, typeFun, default, help= paramTuple
+            if typeFun == bool:
+                if default is True:
+                    action="store_false"
+                else:
+                    action="store_true"
+                help += " Action: "+action
+                parser.add_argument("--" + name, help=help, action=action)
+            else:
+                parser.add_argument("--"+name, type=typeFun, help=help+" Default=%(default)s", default=default)
+        return parser
+
+
+
+class MyArgParser(argparse.ArgumentParser):
+    def parse_args(self, args=None, namespace=None):
+        args = super().parse_args(args=args, namespace=namespace)
+        arg_groups = {}
+        for group in self._action_groups:
+            group_dict = {a.dest: getattr(args, a.dest, None) for a in group._group_actions}
+            arg_groups[group.title] = group_dict
+
+        return arg_groups

coprinet/pricePrediction/config.py

@@ -0,0 +1,50 @@
+import os
+
+################## Common options ##################
+NUM_BUCKETS = 100
+N_CPUS = 12
+
+MULTIPROC_PREFETCH_FACTOR = 4
+MULTIPROC_TIMEOUT = 60
+PERSISTENT_WORKERS = False
+
+BATCH_SIZE = 512
+NUM_WORKERS_PER_GPU = 2
+BUFFER_N_BATCHES_FOR_PRED=50 #Process as many batches together for speed. Set it as large as possible to fit your RAM memory
+
+USE_MMOL_INSTEAD_GRAM = True
+
+RANDOM_SEED = 121
+
+################## Net  ##################
+N_LAYERS=10
+N_HIDDEN_NODE=75
+N_HIDDEN_EDGE=50
+
+LEARNING_RATE= 1e-5
+PATIENT_REDUCE_LR_PLATEAU_N_EPOCHS = 25
+COSINE_LR_SCHEDULE_N_EPOCHS = 10
+
+############# Data ##################
+
+TEST_SIZE= int(2e5)
+CSV_CHUNK_SIZE= int(1e5)
+
+DATA_DIR = os.path.abspath(os.path.join(__file__, "../.."))
+
+RAW_DATA_DIRNAME= "data/MCule" #Can be provided via command line arguments
+BUILDING_BLOCKS_FNAME= os.path.join(RAW_DATA_DIRNAME, "mcule_purchasable_building_blocks_210806.smi.gz")
+
+DATASET_DIRNAME = "data/mcule_chunks"
+DATASET_IN_MEMORY = True # Set to true to load the dataset into memory
+ENCODED_DIR = os.path.join(DATASET_DIRNAME, "encoded")
+
+RAW_DATA_FILE_SUFFIX = r"mcule_full_(train|test|val)\.csv_split_\w+\.csv$"
+
+DEFAULT_MODEL = os.path.join( DATA_DIR, "data/models/trained_inStock/lightning_logs/version_0/checkpoints/epoch=238-step=1295857.ckpt")
+# DEFAULT_MODEL = os.path.join( DATA_DIR, "data/models/trained_virtual/lightning_logs/version_0/checkpoints/epoch=172-step=3231639.ckpt")
+
+
+######################### Ignore the next lines if not running a comparison ################
+USE_FEATURES_NET = False       #set this to True to use a MLP instead of GNN
+NORMALIZE_NODE_FEATS = False   #set this to True if USE_FEATURES_NET==True

coprinet/pricePrediction/dataManager/dataManager.py

@@ -0,0 +1,255 @@
+import atexit
+import logging
+import os
+import shutil
+
+import dirsync
+import joblib
+import torch
+
+import pytorch_lightning as pl
+from pytorch_lightning import seed_everything
+
+from torch.utils.data import DataLoader, Subset
+
+from pricePrediction import config
+from pricePrediction.ArgParser_base import ArgParseable
+from pricePrediction.dataManager.dataset import Dataset_graphPrice
+from pricePrediction.preprocessData.prepareDataMol2Price import DataBuilder
+from pricePrediction.utils import EncodedDirNamesAndTemplates
+
+
+class GraphPriceDatamodule(pl.LightningDataModule, ArgParseable):
+
+
+    DESIRED_PARAMS_TO_ASK= ['encodedDir', 'batch_size', 'num_workers', 'use_lds', 'use_weights', 'force_compute_data',
+                            'do_undersample', 'augment_labels', 'copy_data_to_wdir', 'debug']
+
+    def __init__(self, encodedDir=config.ENCODED_DIR, deg_fname=None, batch_size: int = config.BATCH_SIZE,
+                 num_workers: int = config.NUM_WORKERS_PER_GPU, random_seed: int = config.RANDOM_SEED,
+                 use_log_for_target: bool = True, use_lds: bool = False,
+                 use_weights: bool = False, force_compute_data: bool = False, do_undersample: bool = False,
+                 augment_labels: bool =True, shuffle_training: bool =True, copy_data_to_wdir: str = None, debug =False):
+        '''The DataManager that will create several DataLoader(s)
+
+        :param str encodedDir: The directory where the dataset has been prepared as lmdb files.
+        :param str deg_fname: The file containing the statistics about the nodes degree. By default is automatically searched within
+                              encodedDir.
+        :param int batch_size: Batch size.
+        :param int num_workers: Number of workers as in "pytorch.DataLoader".
+        :param int random_seed: Random seed.
+        :param bool use_log_for_target: Apply log to the target (log(price)).
+        :param bool use_lds: Use label density smoothing to reduce data imbalance.
+        :param bool use_weights: Apply weights to reduce data imbalance.
+        :param bool force_compute_data: Recompute all data before using the DataManager. Takes some time.
+        :param bool do_undersample: Use undersampling to reduce data imbalance.
+        :param bool augment_labels: Add random perturbation to training labels to reduce overfitting.
+        :param bool shuffle_training: Shuffle training data.
+        :param str copy_data_to_wdir: An optional working dir where data can be copied. Useful to avoid network disk usage.
+        :param bool debug: Use only a small subset of the dataset
+        '''
+
+        seed_everything(random_seed)
+        super().__init__()
+
+        if encodedDir is None:
+            encodedDir = config.ENCODED_DIR
+        if deg_fname is None:
+            names = EncodedDirNamesAndTemplates(encodedDir)
+            deg_fname = names.DEGREES_FNAME
+
+        self.deg_fname = deg_fname
+        self.do_undersample = do_undersample
+        self.use_weights = use_weights or use_lds
+        self.use_lds = use_lds
+        self.use_log_for_target = use_log_for_target
+        self.random_seed = random_seed
+
+        self.copy_data_to_wdir = copy_data_to_wdir
+
+        self.encodedDir = encodedDir
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.force_compute_data = force_compute_data
+        self.augment_labels = augment_labels
+        self.shuffle_training = shuffle_training
+
+        self.debug = debug
+
+        ###### Delayed variables
+        self._dims = None
+        self._train_n_instances = None
+        self.setup_done = False
+        self.dataset_train = None
+        self.dataset_val = None
+        self.dataset_test = None
+
+    def _getBucketsFname(self, training_phase):
+        names = EncodedDirNamesAndTemplates(self.encodedDir)
+        return names.BUCKETS_FNAME_TEMPLATE % training_phase
+
+    @property
+    def dims(self):
+        if self._dims is None:
+            metadata_fname = EncodedDirNamesAndTemplates(self.encodedDir).DATASET_METADATA_FNAME_TEMPLATE % "train"
+            size_metadata_info = joblib.load(metadata_fname)
+            self._dims = (size_metadata_info["nodes_n_features"], size_metadata_info["edges_n_features"])
+            self._train_n_instances = size_metadata_info["total_size"]
+        return self._dims
+
+    def get_nodes_degree(self):
+        if os.path.isfile( self.deg_fname ):
+            return joblib.load(self.deg_fname)
+        else:
+            self.prepare_data()
+            return self.get_nodes_degree()
+
+    def prepare_data(self):
+        # Compute if not already done
+        if self.force_compute_data:
+            dataBuilder = DataBuilder()
+            dataBuilder.prepareDataset(encodedDir=self.encodedDir, datasetSplit="train")
+            dataBuilder.prepareDataset(encodedDir=self.encodedDir, datasetSplit="test")
+            dataBuilder.prepareDataset(encodedDir=self.encodedDir, datasetSplit="val")
+        else:
+            dataBuilder = DataBuilder()
+            if not os.path.exists(os.path.join(self.encodedDir, "train")):
+                dataBuilder.prepareDataset(encodedDir=self.encodedDir, datasetSplit="train")
+            if not os.path.exists(os.path.join(self.encodedDir, "test")):
+                dataBuilder.prepareDataset(encodedDir=self.encodedDir, datasetSplit="test")
+            if not os.path.exists(os.path.join(self.encodedDir, "val")):
+                dataBuilder.prepareDataset(encodedDir=self.encodedDir, datasetSplit="val")
+
+        # if self.copy_data_to_wdir:
+        #     tmpArgs = {}
+        #     print(self.copy_data_to_wdir, os.path.isdir(self.copy_data_to_wdir))
+        #     if not os.path.isdir(self.copy_data_to_wdir):
+        #         tmpArgs["prefix"] = self.copy_data_to_wdir
+        #         self._tmpdir = TemporaryDirectory(**tmpArgs)
+        #         def cleanTmpdir():
+        #             del self._tmpdir
+        #         atexit.register( cleanTmpdir)
+        #
+        #         self.copy_data_to_wdir = self._tmpdir.name
+
+        if self.copy_data_to_wdir:
+            if not os.path.isdir(self.copy_data_to_wdir):
+                os.makedirs(self.copy_data_to_wdir)
+                _tmpdir = self.copy_data_to_wdir
+                def cleanTmpdir():
+                    shutil.rmtree( _tmpdir)
+                atexit.register(cleanTmpdir)
+
+            print("copying training data from %s to %s" % (self.encodedDir, self.copy_data_to_wdir))
+            dirsync.sync(self.encodedDir, self.copy_data_to_wdir, 'sync', verbose=True, logger=logging.getLogger('dummy'))
+
+    def setup(self, stage=None):
+        if self.dataset_train is not None and self.dataset_val is not None and self.dataset_test is not None:
+            return
+
+        if self.copy_data_to_wdir is not None:
+            encodedDir = self.copy_data_to_wdir
+        else:
+            encodedDir = self.encodedDir
+
+        self.dataset_train = Dataset_graphPrice(encodedDir, datasetSplit="train",
+                                                use_log_for_target=self.use_log_for_target,
+                                                use_weights=self.use_weights, augment_labels = self.augment_labels,
+                                                buckets_fname=self._getBucketsFname("train"),
+                                                use_lds=self.use_lds, do_undersample=self.do_undersample)
+
+        self.dataset_val = Dataset_graphPrice(encodedDir, datasetSplit="val",
+                                              use_log_for_target=self.use_log_for_target,
+                                              use_weights=self.use_weights, augment_labels = False,
+                                              buckets_fname=self._getBucketsFname("val"),
+                                              use_lds=self.use_lds, do_undersample=False)
+
+        self.dataset_test = Dataset_graphPrice(encodedDir, datasetSplit="test",
+                                                use_log_for_target=self.use_log_for_target,
+                                                use_weights=self.use_weights, augment_labels = False,
+                                                buckets_fname=self._getBucketsFname("test"),
+                                                use_lds=self.use_lds, do_undersample=False)
+
+        atexit.register( self.dataset_train.close )
+        atexit.register( self.dataset_val.close )
+        atexit.register( self.dataset_test.close )
+
+    def _generic_dataloader(self, dataset, num_workers=None, **kwargs):
+        if num_workers is None:
+            num_workers = self.num_workers
+
+
+        if num_workers>0 and "timeout" not in kwargs:
+            kwargs["timeout"] = config.MULTIPROC_TIMEOUT
+
+        if num_workers > 0 and "prefetch_factor" not in kwargs:
+            kwargs["prefetch_factor"] = config.MULTIPROC_PREFETCH_FACTOR
+
+        if num_workers > 0 and "persistent_workers" not in kwargs:
+            kwargs["persistent_workers"] = config.PERSISTENT_WORKERS
+
+        print("Using num_workers:", num_workers)
+        #print(self.trainer.reload_dataloaders_every_epoch)
+
+        from torch_geometric.data import Batch
+        def from_list_of_examples_to_batch(examples):
+            graphs, ys = zip(* examples)
+            g_batch = Batch.from_data_list(graphs)
+            ys = torch.stack(ys)
+            return [g_batch, ys]
+        if self.debug:
+            dataset = Subset(dataset, range(21*self.batch_size))
+        return DataLoader(dataset, num_workers=num_workers, batch_size=self.batch_size,
+                          pin_memory=num_workers>0, collate_fn = from_list_of_examples_to_batch,
+                          **kwargs)
+
+    def train_dataloader(self, **kwargs):
+        if self.trainer:
+            print("Generating a train_dataloader at epoch: %d"% (  self.trainer.current_epoch  ))
+        if "persistent_workers" not in kwargs:
+            kwargs["persistent_workers"] = False
+        return self._generic_dataloader(self.dataset_train,  shuffle=True, **kwargs)
+
+    def val_dataloader(self, **kwargs):
+        return self._generic_dataloader(self.dataset_val, shuffle=False, **kwargs)
+
+    def test_dataloader(self, **kwargs):
+        return self._generic_dataloader(self.dataset_test, shuffle=False, **kwargs)
+
+    def predict_dataloader(self, dataset=None):
+
+        if dataset == "train":
+            raise NotImplementedError()
+        elif dataset == "test":
+            raise NotImplementedError()
+        elif dataset == "val":
+            raise NotImplementedError()
+        if dataset is not None:  # TODO: check type
+            raise NotImplementedError()
+        else:
+            raise NotImplementedError()
+
+
+
+if __name__ == "__main__":
+
+    from pricePrediction.ArgParser_base import ArgParseable, MyArgParser
+
+    parser = MyArgParser(prog=None, usage=None, description=None, )
+
+    group = parser.add_argument_group(title="data")
+    GraphPriceDatamodule.addParamsToArgParse(group)
+    cmd_args = parser.parse_args()
+    dm = GraphPriceDatamodule(**cmd_args.get("data", {}))
+    dm.setup()
+    allPrices = []
+    for batchG, batchY in dm.test_dataloader():
+        allPrices.extend( batchY.tolist() )
+
+    import matplotlib.pyplot as plt
+    plt.hist(allPrices, bins=200)
+    plt.show()
+
+    '''
+python -m pricePrediction.dataManager.dataManager
+    '''

coprinet/pricePrediction/dataManager/dataset.py

@@ -0,0 +1,171 @@
+import operator
+import os
+import joblib
+import lmdb
+import numpy as np
+import torch
+
+from torch.utils.data import Dataset
+
+from pricePrediction import config
+from pricePrediction.preprocessData.serializeDatapoints import getExampleId, deserializeExample
+from pricePrediction.utils import search_buckedId, EncodedDirNamesAndTemplates
+
+
+class Dataset_graphPrice(Dataset):
+    def __init__(self, encodedDir, datasetSplit, use_log_for_target=True, use_weights=False, buckets_fname=None,
+                 use_lds=False, do_undersample=False, augment_labels=False, store_in_memory=config.DATASET_IN_MEMORY,
+                 normalize_node_feats=config.NORMALIZE_NODE_FEATS):
+
+        self.use_log_for_target = use_log_for_target
+        self.use_weights = use_weights
+        self.buckets_fname = buckets_fname
+        self.use_lds = use_lds
+        self.do_undersample = do_undersample
+        self.augment_labels = augment_labels
+        self.store_in_memory = store_in_memory
+        self.normalize_node_feats = normalize_node_feats
+
+
+        assert (self.buckets_fname is not None and self.use_weights) or not self.use_weights
+        if self.do_undersample: raise NotImplementedError()
+
+        dataNames = EncodedDirNamesAndTemplates(encodedDir)
+        size_metadata_info = joblib.load( dataNames.DATASET_METADATA_FNAME_TEMPLATE % datasetSplit)
+
+        if self.normalize_node_feats:
+            self.min_x = torch.FloatTensor(size_metadata_info["feature_stats"]["min"])
+            self.max_x = torch.FloatTensor(size_metadata_info["feature_stats"]["max"])
+            self.range_x = self.max_x - self.min_x
+            assert not torch.isclose(self.range_x, torch.zeros(1)).any(), f"Error, range close to 0, {torch.where(torch.isclose(self.range_x, torch.zeros(1)))}"
+            assert not torch.isinf(self.range_x).any()
+
+        self.total_size = size_metadata_info["total_size"]
+        self.size_per_file = size_metadata_info["sizes_list"]
+        self.db_fnames = [ os.path.join( encodedDir, datasetSplit, os.path.basename(fname))
+                                        for fname in size_metadata_info["fnames_list"]]
+        self.cum_size_at_file = np.cumsum(np.array(size_metadata_info["sizes_list"], dtype=int))
+        self.last_cum_elemIdx_at_file = self.cum_size_at_file - 1
+
+        self.fileIdx_to_fileBasename = [os.path.basename(fname) for fname in self.db_fnames]
+
+        #Delayed load
+        self.db_managers= [None]*len(self.size_per_file)
+
+        self._bucket_ranges = None
+        self._n_per_bucket = None
+        self._cache = {}
+
+
+    def _idx_to_file_and_num(self, idx):
+        fileIdx = search_buckedId(idx, self.last_cum_elemIdx_at_file)
+        idx = idx - (self.cum_size_at_file[fileIdx - 1] if fileIdx - 1 >= 0 else 0)
+        return fileIdx, idx
+
+    def _init_db(self, fileIdx):
+        env = lmdb.open(self.db_fnames[fileIdx],
+            readonly=True, lock=False,
+            readahead=False, meminit=False)
+        txn = env.begin()
+        self.db_managers[fileIdx] = (env, txn)
+        return env, txn
+
+    @property
+    def bucket_ranges(self):
+        if self._bucket_ranges is None:
+            self._getBuckets()
+        return self._bucket_ranges
+
+    @property
+    def n_per_bucket(self):
+        if self._n_per_bucket is None:
+            self._getBuckets()
+        return self._n_per_bucket
+
+    def n_batches(self, batch_size):
+        return len(self)//batch_size + int(bool(len(self)%batch_size))
+
+    def _getBuckets(self):
+        buckets = joblib.load(self.buckets_fname)
+        bucket_ranges, n_per_bucket = buckets["bucket_ranges"], buckets["n_per_bucket"]
+
+        if self.use_lds:
+            print("using lds")
+            from scipy.ndimage import gaussian_filter1d
+            n_per_bucket = gaussian_filter1d(n_per_bucket, sigma=2)
+        self._bucket_ranges = bucket_ranges
+        self._n_per_bucket = n_per_bucket
+
+    def compute_weight(self, label):
+        if not self.use_log_for_target:
+            label = np.exp(label)
+        bucket_num = search_buckedId(label, self.bucket_ranges)
+        w = self.total_size / (1 + self.n_per_bucket[bucket_num])
+        return w
+
+    def prepare_undersample(self, maximum_percentile=80):
+        maximum_per_bucket = np.percentile(self.n_per_bucket, maximum_percentile)
+        prob_per_bucket = maximum_per_bucket / (1e-10 + self.n_per_bucket)
+        prob_per_bucket = np.where(prob_per_bucket < 1, prob_per_bucket, 1)
+        print("maximum_per_bucket while sampling: %d" % maximum_per_bucket)
+        new_len = sum([min(maximum_per_bucket, n_elems) for n_elems in self.n_per_bucket])
+
+        def predicate(graph_y):
+            y = graph_y[1]
+            bucket_num = search_buckedId(y, self.bucket_ranges)
+            return np.random.rand() < prob_per_bucket[bucket_num]
+
+        # dataset = dataset.select(predicate)
+        raise NotImplementedError() #This code is obsolete and requires refactoring
+
+    def normalize_x(self, x):
+        return (x-self.min_x)/self.range_x
+
+    def __getitem__(self, index):
+
+        out = self._cache.get(index, None)
+        if out is not None:
+            return out
+        else:
+            if index >= self.total_size:
+                raise IndexError()
+            fileIdx, pointIdx = self._idx_to_file_and_num(index)
+            db_manager = self.db_managers[fileIdx]
+            if db_manager is None:
+                txn= self._init_db(fileIdx)[-1]
+            else:
+                txn = db_manager[-1]
+
+            dataBytes = txn.get(getExampleId(self.fileIdx_to_fileBasename[fileIdx], pointIdx))
+            if dataBytes is None:
+                msg = "Error:\n%s %s   %s   %s"%(index, fileIdx, pointIdx, getExampleId(self.fileIdx_to_fileBasename[fileIdx], pointIdx))
+                raise Exception(msg)
+            graph, label = deserializeExample(dataBytes)
+
+            if self.normalize_node_feats:
+                graph.x = self.normalize_x(graph.x)
+
+            if self.use_log_for_target:
+                label = np.log(label)
+
+            if self.augment_labels:
+                label = label * (1+0.05*np.random.rand()-0.025)
+
+            if self.use_weights:
+                graph.w = self.compute_weight(label)
+            else:
+                graph.w = 1
+            graph.y = label
+            if self.store_in_memory:
+                self._cache[index] = (graph, label)
+            return graph, label
+
+    def __len__(self):
+        return self.total_size
+
+    def close(self):
+        for db_handler in self.db_managers:
+            if db_handler:
+                env, txn = db_handler
+                env.close()
+

coprinet/pricePrediction/evaluation/evaluation.py

@@ -0,0 +1,198 @@
+import os, io
+from datetime import datetime
+
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+from scipy.stats import pearsonr, spearmanr
+from pytorch_lightning.callbacks import Callback
+
+def plt_to_numpy(fig):
+    #it should be called after having done the plt.plot stuff, in replacement of plt.show()
+    #TO be used with torch tensorboard
+    from skimage.color import rgba2rgb
+    io_buf = io.BytesIO()
+    fig.savefig(io_buf, format='raw', dpi=fig.dpi)
+    io_buf.seek(0)
+    img_arr = np.reshape(np.frombuffer(io_buf.getvalue(), dtype=np.uint8),
+                         newshape=(int(fig.bbox.bounds[3]), int(fig.bbox.bounds[2]), -1))
+    img_arr = rgba2rgb( img_arr )
+    img_arr = img_arr.astype(np.float32)
+    io_buf.close()
+    return img_arr
+
+class InferencePlotter():
+    def __init__(self, pl_module: 'pl.LightningModule'):
+        self.pl_module = pl_module
+        self.data_loaders = {}
+
+    def get_dataloader(self, mode, n_workers=None):
+
+        if mode in self.data_loaders:
+            return self.data_loaders[mode]
+        else:
+            if mode == "train":
+                dataloader = self.pl_module.train_dataloader(num_workers=n_workers)
+            elif mode == "val":
+                dataloader = self.pl_module.val_dataloader(num_workers=n_workers)
+            elif mode == "test":
+                dataloader = self.pl_module.test_dataloader(num_workers=n_workers)
+            else:
+                raise ValueError("Error. Option mode=%s not supported" % mode)
+
+            self.data_loaders[mode] = dataloader
+            return dataloader
+
+    @classmethod
+    def do_plots(cls, ys, y_preds, label="test", tensorboard_wrapper=None, step=0):
+
+        if tensorboard_wrapper:
+            figArgs = dict(figsize=(5, 4), dpi=450)
+        else:
+            figArgs = {}
+        fig = plt.figure(**figArgs)
+        plt.hist(ys, bins=200)
+        name = label + " price hist"
+        plt.title(name)
+        img_arr = plt_to_numpy(fig)
+        if tensorboard_wrapper:
+            tensorboard_wrapper.add_image(name, img_arr.transpose([2, 0, 1]), step)
+            plt.close()
+        else:
+            plt.show()
+
+        if y_preds is not None:
+            fig = plt.figure(**figArgs)
+            plt.hist(y_preds, bins=200)
+            name = label + " pred hist"
+            plt.title(name)
+            img_arr = plt_to_numpy(fig)
+            if tensorboard_wrapper:
+                tensorboard_wrapper.add_image(name, img_arr.transpose([2, 0, 1]), step)
+                plt.close()
+            else:
+                plt.show()
+
+            r, p = pearsonr(ys, y_preds)
+            rho, p2 = spearmanr(ys, y_preds)
+            if tensorboard_wrapper:
+                tensorboard_wrapper.add_scalar(label + '_pearson-r', r, step)
+                tensorboard_wrapper.add_scalar(label + '_spearman-r', rho, step)
+
+                if label == "test":
+                    print("pearson-r", r, "p-val", p)
+                    print("spearman-r", rho, "p-val", p)
+            else:
+                print("pearson-r", r, "p-val", p)
+                print("spearman-r", rho, "p-val", p)
+
+            fig = plt.figure(**figArgs)
+            max_val = max([9, ys.max()])
+            min_val = 1
+            plt.hexbin(ys, y_preds, bins="log", gridsize=100, extent=[min_val, max_val, min_val, max_val] )
+            plt.axline((1,1), slope=1, c="r", lw=1, ls="--", alpha=0.5)
+            plt.xlabel("log(price)")
+            plt.ylabel("Pred")
+            plt.colorbar()
+            name = label + " price vs pred"
+            plt.title(name)
+            img_arr = plt_to_numpy(fig)
+
+            if tensorboard_wrapper:
+                tensorboard_wrapper.add_image(name, img_arr.transpose([2, 0, 1]), step)
+                plt.close()
+            else:
+                plt.show()
+
+            # fig = plt.figure(**figArgs)
+            # plt.boxplot( abs(y_preds - ys)/ys )
+            # name = method_name + "_relative_error_boxplot"
+            # plt.title(name)
+            # img_arr = plt_to_numpy(fig)
+            # if tensorboard_wrapper:
+            #     tensorboard_wrapper.add_image(name, img_arr.transpose([2,0,1]), step )
+            #     plt.close()
+            # else:
+            #     plt.show()
+            return r
+
+    def get_ytrue_ypred(self, mode="test"):
+        with torch.no_grad():
+            ytrue_ypred = []
+            dataLoader = self.get_dataloader(mode=mode)
+            n_elems = len(dataLoader)
+            for i, (batch, y) in enumerate(dataLoader):
+                ytrue_ypred.append((y, self.pl_module(batch.to(self.pl_module.device))))
+                if i % 10 == 0: print("%07d/%07d            " % (i,n_elems), end="\r")
+            print()
+            ytrue_train, ypred_train = zip(*ytrue_ypred)
+            ytrue = torch.cat(ytrue_train).to("cpu").numpy()
+            ypred = torch.cat(ypred_train).to("cpu").numpy()
+        return ytrue, ypred
+
+    def pred_and_plot(self, mode="test", step=0):
+        ytrue, ypred = self.get_ytrue_ypred(mode=mode)
+        r = self.do_plots(ytrue, ypred, label=mode,
+                          tensorboard_wrapper=self.pl_module.logger.experiment,
+                          step=step)
+        return r, (ytrue, ypred)
+
+
+class PlotFigsTensorboardCallback(Callback):
+
+    def __init__(self, output_dir: str = None, frequency=10, save_csv=True):
+        super().__init__()
+        self.output_dir = output_dir
+        self.frequency = frequency
+        self.current_epoch = 0
+        self.plotter = None  # InferencePlotter(pl_module=)
+        self.save_csv = save_csv
+
+    def setup(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', stage=None) -> None:
+        self.plotter = InferencePlotter(pl_module=pl_module)
+
+    # def on_init_end(self, trainer):
+    #     if self.output_dir is None:
+    #         self.output_dir = trainer.log_dir
+
+    def on_train_start(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
+        if trainer.is_global_zero:
+            if self.output_dir is None:
+                self.output_dir = trainer.log_dir
+            if self.save_csv:
+                self.csv_fname = os.path.join(self.output_dir, "logs.csv")
+                with open(self.csv_fname, "w") as f:
+                    f.write("Epoch,Date,trainL1,valL1,trainR,valR")
+
+    def _getDateHourStr(self):
+        return datetime.now().strftime("%d/%m-%H:%M")
+
+    def on_train_epoch_end(
+            self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', unused=None
+    ) -> None:
+
+        if trainer.is_global_zero:
+            self.current_epoch += 1
+            if self.current_epoch % self.frequency != 0:
+                return
+            print("Evaluating preds on epoch %d [%s]" % (self.current_epoch - 1, self._getDateHourStr()))
+            r_train, lossL1_train = self.process_data(mode="train")
+            r_val, lossL1_val = self.process_data(mode="val")
+
+            if self.save_csv:
+                with open(self.csv_fname, "a") as f:
+                    f.write("\n%d,%s,%f,%f,%1.4f,%1.4f" % (self.current_epoch - 1, self._getDateHourStr(),
+                                                           lossL1_train, lossL1_val, r_train, r_val,))
+
+    def process_data(self, mode="train"):
+        r, (ytrue, ypred) = self.plotter.pred_and_plot(mode=mode, step=self.current_epoch)
+        lossL1 = float(np.mean(np.abs(ytrue - ypred)))
+        return r, lossL1
+
+
+if __name__ == "__main__":
+    import numpy as np
+
+    ys = 1 + np.random.rand(4000) * 7
+    y_preds = ys - np.random.rand()
+    InferencePlotter.do_plots(ys, y_preds, label="test", tensorboard_wrapper=None)

coprinet/pricePrediction/nets/FDS_imbalance.py

@@ -0,0 +1,189 @@
+import logging
+import numpy as np
+from scipy.ndimage import gaussian_filter1d
+from scipy.signal.windows import triang
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+'''
+Modified from:
+
+https://github.com/YyzHarry/imbalanced-regression
+https://arxiv.org/pdf/2102.09554.pdf
+'''
+
+print = logging.info
+
+def calibrate_mean_var(matrix, m1, v1, m2, v2, clip_min=0.1, clip_max=10):
+    if torch.sum(v1) < 1e-10:
+        return matrix
+    if (v1 == 0.).any():
+        valid = (v1 != 0.)
+        factor = torch.clamp(v2[valid] / v1[valid], clip_min, clip_max)
+        matrix[:, valid] = (matrix[:, valid] - m1[valid]) * torch.sqrt(factor) + m2[valid]
+        return matrix
+
+    factor = torch.clamp(v2 / v1, clip_min, clip_max)
+    return (matrix - m1) * torch.sqrt(factor) + m2
+
+
+class FDS(nn.Module):
+
+    def __init__(self, feature_dim, max_val, min_val=1, bucket_step=0.1, start_update=0, start_smooth=1,
+                 kernel='gaussian', ks=5, sigma=2, momentum=0.9):
+        super(FDS, self).__init__()
+        self.feature_dim = feature_dim
+
+        self.bucket_start = min_val
+        self.bucket_step = bucket_step
+
+        bucket_end = min_val
+        bucket_ranges = []
+        while bucket_end < max_val:
+            next_bucket = bucket_end + bucket_step
+            bucket_ranges.append((bucket_end, next_bucket))
+            bucket_end = next_bucket
+
+        n_buckets = len( bucket_ranges )
+        self.n_buckets = n_buckets
+        self.bucket_end = bucket_end
+
+        self.kernel_window = self._get_kernel_window(kernel, ks, sigma)
+        self.half_ks = (ks - 1) // 2
+        self.momentum = momentum
+        self.start_update = start_update
+        self.start_smooth = start_smooth
+
+        self.register_buffer('epoch', torch.zeros(1).fill_(start_update))
+        self.register_buffer('running_mean', torch.zeros(n_buckets, feature_dim))
+        self.register_buffer('running_var', torch.ones(n_buckets, feature_dim))
+        self.register_buffer('running_mean_last_epoch', torch.zeros(n_buckets , feature_dim))
+        self.register_buffer('running_var_last_epoch', torch.ones(n_buckets , feature_dim))
+        self.register_buffer('smoothed_mean_last_epoch', torch.zeros(n_buckets , feature_dim))
+        self.register_buffer('smoothed_var_last_epoch', torch.ones(n_buckets , feature_dim))
+        self.register_buffer('num_samples_tracked', torch.zeros(n_buckets))
+
+    @staticmethod
+    def _get_kernel_window(kernel, ks, sigma):
+        assert kernel in ['gaussian', 'triang', 'laplace']
+        half_ks = (ks - 1) // 2
+        if kernel == 'gaussian':
+            base_kernel = [0.] * half_ks + [1.] + [0.] * half_ks
+            base_kernel = np.array(base_kernel, dtype=np.float32)
+            kernel_window = gaussian_filter1d(base_kernel, sigma=sigma) / sum(gaussian_filter1d(base_kernel, sigma=sigma))
+        elif kernel == 'triang':
+            kernel_window = triang(ks) / sum(triang(ks))
+        else:
+            laplace = lambda x: np.exp(-abs(x) / sigma) / (2. * sigma)
+            kernel_window = list(map(laplace, np.arange(-half_ks, half_ks + 1))) / sum(map(laplace, np.arange(-half_ks, half_ks + 1)))
+
+        print(f'Using FDS: [{kernel.upper()}] ({ks}/{sigma})')
+        return torch.tensor(kernel_window, dtype=torch.float32).cuda()
+
+    def _update_last_epoch_stats(self):
+        self.running_mean_last_epoch = self.running_mean
+        self.running_var_last_epoch = self.running_var
+
+        self.smoothed_mean_last_epoch = F.conv1d(
+            input=F.pad(self.running_mean_last_epoch.unsqueeze(1).permute(2, 1, 0),
+                        pad=(self.half_ks, self.half_ks), mode='reflect'),
+            weight=self.kernel_window.view(1, 1, -1), padding=0
+        ).permute(2, 1, 0).squeeze(1)
+        self.smoothed_var_last_epoch = F.conv1d(
+            input=F.pad(self.running_var_last_epoch.unsqueeze(1).permute(2, 1, 0),
+                        pad=(self.half_ks, self.half_ks), mode='reflect'),
+            weight=self.kernel_window.view(1, 1, -1), padding=0
+        ).permute(2, 1, 0).squeeze(1)
+
+    def reset(self):
+        self.running_mean.zero_()
+        self.running_var.fill_(1)
+        self.running_mean_last_epoch.zero_()
+        self.running_var_last_epoch.fill_(1)
+        self.smoothed_mean_last_epoch.zero_()
+        self.smoothed_var_last_epoch.fill_(1)
+        self.num_samples_tracked.zero_()
+
+    def update_last_epoch_stats(self, epoch):
+        if epoch == self.epoch + 1:
+            self.epoch += 1
+            self._update_last_epoch_stats()
+            print(f"Updated smoothed statistics on Epoch [{epoch}]!")
+
+    def get_bucket_id(self, ys):
+        bucket_id = (ys - self.bucket_start) // self.bucket_step
+        return bucket_id.type(torch.int32)
+
+    def update_running_stats(self, features, ys, epoch):
+        if epoch < self.epoch:
+            return
+
+        assert self.feature_dim == features.size(1), "Input feature dimension is not aligned!"
+        assert features.size(0) == ys.size(0), "Dimensions of features and labels are not aligned!"
+
+        bucket_ids = self.get_bucket_id(ys)
+        lower_mask = ys <= self.bucket_start
+        upper_mask = ys >= self.bucket_end
+        # for y in torch.unique(ys):
+        #     bucket_id = self.get_bucket_id(y)
+        for bucket_id in torch.unique(bucket_ids):
+            if bucket_id > self.n_buckets - 1 or bucket_id < self.bucket_start:
+                continue
+            elif bucket_id == 0:
+                curr_feats = features[lower_mask]
+            elif bucket_id == self.n_buckets - 1:
+                curr_feats = features[upper_mask]
+            else:
+                curr_feats = features[bucket_ids == bucket_id]
+
+            curr_num_sample = curr_feats.size(0)
+            curr_mean = torch.mean(curr_feats, 0)
+            curr_var = torch.var(curr_feats, 0, unbiased=True if curr_feats.size(0) != 1 else False)
+            self.num_samples_tracked[bucket_id] += curr_num_sample
+            factor = self.momentum if self.momentum is not None else \
+                (1 - curr_num_sample / float(self.num_samples_tracked[bucket_id]))
+            factor = 0 if epoch == self.start_update else factor
+            self.running_mean[bucket_id] = \
+                (1 - factor) * curr_mean + factor * self.running_mean[bucket_id]
+            self.running_var[bucket_id] = \
+                (1 - factor) * curr_var + factor * self.running_var[bucket_id]
+
+        print(f"Updated running statistics with Epoch [{epoch}] features!")
+
+    def smooth(self, features, ys, epoch): #TODO: rename to forward
+        if epoch < self.start_smooth:
+            return features
+
+        bucket_ids = self.get_bucket_id(ys)
+        lower_mask = ys <= self.bucket_start
+        upper_mask = ys >= self.bucket_end
+
+        # ys = ys.squeeze(1)
+        # for y in torch.unique(ys):
+        for bucket_id in torch.unique(bucket_ids):
+            if bucket_id > self.n_buckets - 1 or bucket_id < self.bucket_start:
+                continue
+            elif bucket_id == self.bucket_start:
+                features[lower_mask] = calibrate_mean_var(
+                    features[lower_mask], #TODO: indexing many times is not efficient
+                    self.running_mean_last_epoch[bucket_id],
+                    self.running_var_last_epoch[bucket_id],
+                    self.smoothed_mean_last_epoch[bucket_id],
+                    self.smoothed_var_last_epoch[bucket_id])
+            elif bucket_id == self.n_buckets - 1:
+                features[upper_mask] = calibrate_mean_var(
+                    features[upper_mask],
+                    self.running_mean_last_epoch[bucket_id],
+                    self.running_var_last_epoch[bucket_id],
+                    self.smoothed_mean_last_epoch[bucket_id],
+                    self.smoothed_var_last_epoch[bucket_id])
+            else:
+                bucket_mask = bucket_ids == bucket_id
+                features[bucket_mask] = calibrate_mean_var(
+                    features[bucket_mask],
+                    self.running_mean_last_epoch[bucket_id],
+                    self.running_var_last_epoch[bucket_id],
+                    self.smoothed_mean_last_epoch[bucket_id],
+                    self.smoothed_var_last_epoch[bucket_id])
+        return features

coprinet/pricePrediction/nets/basicArchitectures.py

@@ -0,0 +1,165 @@
+import torch
+import numpy as np
+from torch import nn
+import torch.nn.functional as F
+from torch_geometric.nn import PNAConv, BatchNorm, global_add_pool, AttentiveFP, Set2Set
+from torch.nn import ModuleList, Linear, ReLU, Sequential, Dropout
+
+
+class MLP(nn.Module):
+
+    def __init__(self,dims, n_layers, hidden_size, dropout=0 ):
+
+        super().__init__()
+
+        self.n_layers = n_layers
+        self.hidden_size = hidden_size
+        self.dims = dims
+        self.dropout = dropout
+
+
+        def block(in_size, n_hidden):
+            layers = [
+                nn.Linear(in_size, n_hidden),
+                nn.BatchNorm1d(n_hidden),
+                nn.ReLU(),
+            ]
+            if self.dropout > 0:
+                layers.append(
+                    nn.Dropout(self.dropout),
+                )
+            return layers
+
+        # Define PyTorch model
+        self.model = nn.Sequential(
+            *block(np.prod(dims), self.hidden_size)
+        )
+
+        self.latent_size = self.hidden_size
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class GNN_PNAConv(torch.nn.Module):
+    def __init__(self,
+        nodes_n_features: int,
+        edges_n_features: int,
+        deg: torch.tensor ,
+        n_layers: int = 6,
+        hidden_size_node: int = 75,
+        hidden_size_edges: int = 50,
+        towers: int = 5,
+        fcc_hidden_size: int = 50,
+        dropout: float = 0,
+        use_fds: bool = False,
+                 **args
+        ):
+        super(GNN_PNAConv, self).__init__()
+
+        self.node_emb = MLP(nodes_n_features, 1,  hidden_size_node)
+        self.edge_emb = MLP(edges_n_features, 1, hidden_size_edges)
+
+        aggregators = ['mean', 'min', 'max', 'std']
+        scalers = ['identity', 'amplification', 'attenuation']
+
+        self.convs = ModuleList()
+        self.batch_norms = ModuleList()
+        for _ in range(n_layers):
+            conv = PNAConv(in_channels=hidden_size_node, out_channels=hidden_size_node,
+                           aggregators=aggregators, scalers=scalers, deg=deg,
+                           edge_dim=hidden_size_edges, towers=towers, pre_layers=2, post_layers=2,
+                           divide_input=False)
+            self.convs.append(conv)
+            self.batch_norms.append(BatchNorm(hidden_size_node))
+        self.set2set = Set2Set(hidden_size_node, processing_steps=6)
+        fc_layers = [
+            Linear(2*hidden_size_node, hidden_size_node), BatchNorm(hidden_size_node), ReLU(),
+            Linear(hidden_size_node, fcc_hidden_size), BatchNorm(fcc_hidden_size), ReLU(),
+        ]
+        self.fcc_hidden_size = fcc_hidden_size
+        if dropout>0:
+            fc_layers += [ Dropout(p=dropout) ]
+
+        self.mlp = Sequential( *fc_layers )
+
+    def forward(self, x, edge_index, edge_attr, batch):
+        x = self.node_emb(x.squeeze())
+        edge_attr = self.edge_emb(edge_attr)
+
+        for conv, batch_norm in zip(self.convs, self.batch_norms):
+            x = F.relu(batch_norm(conv(x, edge_index, edge_attr)))
+
+        x = self.set2set(x, batch) #Set2Set #GlobalAttention #GraphMultisetTransformer
+
+        return self.mlp(x)
+
+
+class QdolarAR(torch.nn.Module):
+    def __init__(self,
+                 nodes_n_features: int,
+                 *args, **kargs
+        ):
+        super(QdolarAR, self).__init__()
+
+        n_layers: int = 5
+        hidden_size_node: int = nodes_n_features
+        fcc_hidden_size: int = 100
+        dropout: float = 0.28
+
+        self.fcc_hidden_size = fcc_hidden_size
+        self.dropout = dropout
+        self.n_layers = n_layers
+        self.hidden_size_node = hidden_size_node
+
+        assert n_layers >=3
+        fc_layers = [
+                BatchNorm(hidden_size_node),
+                Linear( hidden_size_node, fcc_hidden_size), BatchNorm(fcc_hidden_size), ReLU(),
+        ]
+        for i in range(n_layers-2):
+            fc_layers += [
+                Linear(fcc_hidden_size, fcc_hidden_size), BatchNorm(fcc_hidden_size), ReLU(),
+            ]
+        if dropout>0:
+            fc_layers += [ Dropout(p=dropout) ]
+
+        self.mlp = Sequential( *fc_layers )
+
+    def forward(self, x, edge_index, edge_attr, batch):
+        assert len(batch) == len(batch.unique())
+        assert not torch.isnan(x).any(), "Error, x contains nan"
+        return self.mlp(x)
+
+class GNN_AttentiveFP(torch.nn.Module):
+    def __init__(self,
+        nodes_n_features: int,
+        edges_n_features: int,
+        deg: torch.tensor ,
+        n_layers: int = 4,
+        hidden_size_node: int = 75,
+        hidden_size_edges: int = 50,
+        towers: int = 5,
+        fcc_hidden_size: int = 50,
+        dropout: float = 0,
+        use_fds: bool = False,
+                 **args
+        ):
+        super(GNN_AttentiveFP, self).__init__()
+
+        self.attent_net = AttentiveFP(nodes_n_features, hidden_size_node, hidden_size_node, edges_n_features,
+                                      n_layers, n_layers, dropout= dropout)
+
+        fc_layers = [
+            Linear(hidden_size_node, hidden_size_edges), ReLU(),
+            Linear(hidden_size_edges, fcc_hidden_size), ReLU(),
+        ]
+        self.fcc_hidden_size = fcc_hidden_size
+        if dropout>0:
+            fc_layers += Dropout(p=dropout)
+
+        self.mlp = Sequential( *fc_layers )
+
+    def forward(self, x, edge_index, edge_attr, batch):
+        x = self.attent_net(x, edge_index, edge_attr, batch)
+        return self.mlp(x)

coprinet/pricePrediction/nets/netsGraph.py

@@ -0,0 +1,236 @@
+from collections import OrderedDict
+from typing import Iterable
+
+import torch
+import numpy as np
+from torch import nn
+import torch.nn.functional as F
+import pytorch_lightning as pl
+
+from pricePrediction import config
+from pricePrediction.ArgParser_base import ArgParseable
+from pricePrediction.config import USE_FEATURES_NET
+from pricePrediction.nets.FDS_imbalance import FDS
+
+
+class PricePredictorModule(pl.LightningModule, ArgParseable):
+
+    DESIRED_PARAMS_TO_ASK= ['n_layers', 'hidden_size_node', 'hidden_size_edges', 'fcc_hidden_size', 'use_fds', 'lr',
+                            'dropout', 'gnn_class', 'weight_decay', 'training_loss']
+    def __init__(
+        self,
+        n_layers: int = config.N_LAYERS,
+        hidden_size_node: int = config.N_HIDDEN_NODE,
+        hidden_size_edges: int = config.N_HIDDEN_EDGE,
+        gnn_class="GNN_PNAConv",
+        towers: int = 5,
+        fcc_hidden_size: int = 25,
+        dropout: float = 0,
+        use_fds: bool = False,
+        training_loss : str= "l2", # l1, l2 or huber
+        last_layer_size: int = 1,
+
+        nodes_n_features: int = None,
+        edges_n_features: int = None,
+        deg: Iterable[float] = None,
+
+        lr: float = config.LEARNING_RATE,
+        b1: float = 0.5,
+        b2: float = 0.999,
+        weight_decay=1e-8,
+        data_hparams =None,
+        logs_only_in_epoch = False
+    ):
+
+
+        '''
+
+        :param str encodedDir: The directory where the dataset has been prepared as lmdb files. defaults to %(config.ENCODED_DIR)s
+        :param str deg_fname: The file containing the statistics about the nodes degree. By default is automatically searched within
+                              encodedDir defaults to None
+        :param int n_layers: The number of layers.
+        :param int hidden_size_node: Channels for nodes.
+        :param int hidden_size_edges: Channels for edges.
+        :param str gnn_class: Name of the model type.
+        :param int towers: Only for gnn_class="GNN_PNAConv", The number of towers of the model.
+        :param int fcc_hidden_size: Size of fully connected.
+        :param float dropout: Dropout rate.
+        :param bool use_fds: Use feature density smoothing for data imbalance
+        :param str training_loss: Training loss: "l2" "l1" or "huber"
+        :param int last_layer_size: The number of neurons for last layer. Set it to 1 for regression
+        :param int nodes_n_features: The number of features used to encode a node
+        :param int edges_n_features: The number of features used to encode an edge
+        :param Iterable deg: The degree count of the dataset
+        :param float lr: Learning rate
+        :param float b1: Adam optimizer b1
+        :param float b2: Adam optimizer b1
+        '''
+
+        super().__init__()
+        self.save_hyperparameters() #copies __init__() kwargs to self.hparams
+        deg =  self.hparams.deg
+        if isinstance(deg, list):
+            deg = torch.tensor(deg[:], dtype=torch.long)
+        else:
+            deg = deg.clone().detach()
+        self.hparams.deg = deg
+
+
+
+        if gnn_class == "GNN_PNAConv":
+            from pricePrediction.nets.basicArchitectures import GNN_PNAConv as GNN
+        elif gnn_class == "GNN_AttentiveFP":
+            from pricePrediction.nets.basicArchitectures import GNN_AttentiveFP  as GNN
+        elif gnn_class == "QdolarAR":
+            assert USE_FEATURES_NET
+            from pricePrediction.nets.basicArchitectures import QdolarAR as GNN
+
+        else:
+            raise ValueError("Error, gnn_class not supported")
+
+        self.net = GNN(** self.hparams)
+
+        self.dropoutLayer = nn.Dropout(self.hparams.dropout)
+
+        self.final_layer = nn.Linear(self.net.fcc_hidden_size, 1)
+
+        if use_fds:
+            self.fds = FDS(self.net.latent_size, max_val=10, start_update=1)
+        else:
+            self.fds = None
+        if training_loss=="l1":
+            self.trainLossF = F.l1_loss
+        elif training_loss=="l2":
+            self.trainLossF = F.mse_loss
+        elif training_loss=="huber":
+            self.trainLossF = F.smooth_l1_loss
+        else:
+            raise ValueError("Loss option not recognized")
+
+        if logs_only_in_epoch: #TODO: check if this prevents multi-gpu hang
+            self.logs_kwargs = dict(on_step=False, on_epoch=True, sync_dist=True)
+        else:
+            self.logs_kwargs = dict(on_step=True, on_epoch=False, sync_dist=False)
+
+
+    def compute_training_loss(self, y_pred, y, w=None):
+        loss = self.trainLossF(y_pred, y, reduction='none')
+        if w is not None:
+            loss *= w
+        loss = torch.mean(loss)
+        return loss
+
+    def forward(self, g):
+        return self.compute_y_pred( g, training=False)
+
+    def compute_y_pred(self, g, y=None, training=True):
+        x = self.net(g.x, g.edge_index, g.edge_attr, g.batch)
+        if training and self.fds is not None:
+            self.fds.smooth(x, y, self.current_epoch)
+        x = self.dropoutLayer(x)
+        y_pred = self.final_layer(x)
+        return y_pred.view(-1)
+
+    def resolve_batch(self, batch):
+        if isinstance(batch, list):
+            return batch[0]
+        else:
+            return batch
+
+    def training_step(self, batch, batch_idx):
+        graphs = self.resolve_batch(batch)
+        y_pred = self.compute_y_pred(graphs, training=True)
+        loss = self.compute_training_loss(y_pred, graphs.y, graphs.w)
+
+        loss_l1 = F.l1_loss(y_pred, graphs.y)
+
+        self.log('loss', loss, **self.logs_kwargs)
+        self.log('loss_l1', loss_l1, prog_bar=True, **self.logs_kwargs)
+
+        tqdm_dict = {'loss': loss.detach()}
+        output = OrderedDict({
+            'loss': loss,
+            'targets': graphs.y,
+            'preds': y_pred.detach(),
+            'progress_bar': tqdm_dict,
+            'log': tqdm_dict
+        })
+        return output
+
+    def _validation_step(self, graphs, batch_idx):
+        with torch.no_grad():
+            y_pred = self.compute_y_pred(graphs, training=False)
+            loss = self.compute_training_loss(y_pred, graphs.y, graphs.w)
+            loss_l1 = F.l1_loss(y_pred, graphs.y)
+            return y_pred, loss, loss_l1
+
+    def validation_step(self, batch, batch_idx):
+        graphs = self.resolve_batch(batch)
+        y_pred, loss, loss_l1 = self._validation_step(graphs, batch_idx)
+        # tensorboard = self.logger.experiment
+        # tensorboard.add_histogram("val_price", graphs.y)
+        # tensorboard.add_histogram("val_preds", y_pred)
+        self.log('val_loss', loss, prog_bar=True, sync_dist=True) #rank_zero_only=True
+        self.log('val_lossL1', loss_l1, prog_bar=True, sync_dist=True)
+        return loss
+
+    def training_epoch_end(self, outputs ):
+        # tensorboard = self.logger.experiment
+        # tensorboard.add_histogram("train_price", y)
+        # tensorboard.add_histogram("train_preds", torch.cat([ x["preds"] for x in  outputs]))
+        # tensorboard.add_histogram("train_price", torch.cat([ x["targets"] for x in  outputs]))
+        print("Epoch %d done"%self.current_epoch)
+
+    def test_step(self, batch, batch_idx):
+        graphs = self.resolve_batch(batch)
+        y_pred, loss, loss_l1 = self._validation_step(graphs, batch_idx)
+        tensorboard = self.logger.experiment
+        # tensorboard.add_histogram("test_price", graphs.y)
+        # tensorboard.add_histogram("test_preds", y_pred)
+        self.log('test_loss', loss, prog_bar=True, sync_dist=True)
+        self.log('test_lossL1', loss_l1, prog_bar=True, sync_dist=True)
+        return loss
+
+    def configure_optimizers(self):
+        lr = self.hparams.lr
+        b1 = self.hparams.b1
+        b2 = self.hparams.b2
+
+        opt = torch.optim.Adam(self.parameters(), lr=lr, betas=(b1, b2), weight_decay=self.hparams.weight_decay)
+
+        # for param_group in opt.param_groups:
+        #     print(param_group["lr"])
+
+        conf = {
+            'optimizer': opt,
+            'scheduler': torch.optim.lr_scheduler.ReduceLROnPlateau(opt, verbose=True, cooldown=2,
+                                                                patience=config.PATIENT_REDUCE_LR_PLATEAU_N_EPOCHS, ),
+            'monitor': 'val_loss'
+        }
+
+        return conf
+
+if __name__ == "__main__":
+    from torch_geometric.data import Batch
+
+    degree =[0, 41130, 117278, 70152, 3104]
+
+    gnn_class = "QdolarAR" #"GNN_PNAConv"# "GNN_AttentiveFP"
+
+    if gnn_class != "QdolarAR":
+        from pricePrediction.preprocessData.smilesToGraph import smiles_to_graph
+    else:
+        from pricePrediction.preprocessData.smilesToDescriptors import smiles_to_graph
+
+    g = smiles_to_graph("CCCCCCCCCCC")
+    nodes_n_features = g.x.shape[1]
+    edges_n_features = g.edge_attr.shape[1]
+
+    net = PricePredictorModule(nodes_n_features=nodes_n_features, edges_n_features=edges_n_features,
+                               deg=degree, lr=1e-3, gnn_class=gnn_class)
+
+    print( net )
+
+    graphs = [smiles_to_graph(smi) for smi in ["CCCCCO", "CCCCCNCCCCN"]]
+    print(graphs[0])
+    print( net(Batch.from_data_list(graphs)))

coprinet/pricePrediction/predict/__main__.py

@@ -0,0 +1,4 @@
+from pricePrediction.predict.predict import main
+
+if __name__ == '__main__':
+    main()

coprinet/pricePrediction/predict/predict.py

@@ -0,0 +1,137 @@
+import argparse
+import os, sys
+from itertools import islice
+import numpy as np
+import pandas as pd
+import torch
+from joblib import Parallel, delayed
+from more_itertools import chunked
+from torch.utils.data import IterableDataset, DataLoader
+from torch_geometric.data import Batch
+from tqdm import tqdm
+import pytorch_lightning as pl
+from pricePrediction.nets.netsGraph import PricePredictorModule
+from pricePrediction.config import NUM_WORKERS_PER_GPU, BATCH_SIZE, DEFAULT_MODEL, USE_FEATURES_NET, \
+    BUFFER_N_BATCHES_FOR_PRED
+
+
+class GraphPricePredictor():
+    name = "price_GNN"
+
+    def __init__(self, model_path=DEFAULT_MODEL, n_gpus = 1, n_cpus= NUM_WORKERS_PER_GPU, batch_size:int=BATCH_SIZE,
+                 **kwargs):
+        self.model_path = model_path
+        self.n_gpus = n_gpus
+        self.n_cpus = n_cpus
+        self.batch_size = batch_size
+        self.trainer = pl.Trainer(gpus=self.n_gpus, logger=False)
+        self.model = PricePredictorModule.load_from_checkpoint(
+            self.model_path,
+            batch_size=self.batch_size,
+            strict=False,
+        )
+
+        if USE_FEATURES_NET:
+            from pricePrediction.preprocessData.smilesToDescriptors import smiles_to_graph
+
+        else:
+            from pricePrediction.preprocessData.smilesToGraph import smiles_to_graph
+
+        self.smiles_to_graph = smiles_to_graph
+
+    def prepare_smi(self, idx_smi):
+        idx, smi = idx_smi
+        graph = self.smiles_to_graph(smi)
+        if graph is None:
+            return None
+        graph.input_idx = idx
+        return graph
+
+    def yieldPredictions(self, smiles_generator, buffer_n_batches=BUFFER_N_BATCHES_FOR_PRED):
+        buffer_size = buffer_n_batches * self.batch_size
+        preds_iter = map(lambda x: self.predictListOfSmiles(x), tqdm(chunked(smiles_generator, buffer_size)))
+        for preds_batch in preds_iter:
+            for pred in preds_batch:
+                yield pred
+
+    def predictListOfSmiles(self, smiles_list):
+        smiles_list = list(smiles_list)
+        graphs_list = list(filter(None.__ne__, map(self.prepare_smi,  enumerate(smiles_list) )))
+        graphs_fn = lambda : graphs_list
+        dataset = MyIterableDataset(graphs_fn, self.n_cpus)
+        dataloader = DataLoader(dataset=dataset, batch_size=self.batch_size, collate_fn=Batch.from_data_list,
+                                num_workers=self.n_cpus)
+
+        preds = self.trainer.predict(self.model, dataloader)
+        n_smiles = len(smiles_list)
+        all_preds = np.nan * np.ones(n_smiles)
+        for i, batch in enumerate(dataloader):
+            batch_preds = preds[i].to("cpu").numpy()
+            idxs = batch.input_idx.to("cpu").numpy().astype(np.int64).tolist()
+            all_preds[idxs] = batch_preds
+        return all_preds
+
+
+class MyIterableDataset(IterableDataset):
+    def __init__(self, generator_fun, num_workers):
+        super().__init__()
+        self.generator_fun = generator_fun
+        self.num_workers = num_workers
+
+    def __iter__(self):
+        worker_info = torch.utils.data.get_worker_info()
+        if worker_info is not None:
+            uid = torch.utils.data.get_worker_info().id
+            return islice( self.generator_fun(), uid, None, self.num_workers)
+        else:
+            return iter(self.generator_fun())
+
+
+def main():
+    parser = argparse.ArgumentParser(prog="CoPriNet")
+    parser.add_argument("input_csv_file", type=str, help="The input csv filename containing a smiles column")
+    parser.add_argument("-o", "--output_file", type=str, required=False, default=None, help="The output filename that will be "
+                            "identical to the input_csv_file but with one additional column for the f1")
+    parser.add_argument("--smiles_colname", type=str, required=False, default="SMILES", help="The colname for SMILES "
+                                                                             "in the input file.  Default: %(default)s")
+    parser.add_argument("--model_path", type=str, required=False, default=DEFAULT_MODEL,
+                        help="The CoPriNet model checkpoing path. Default: %(default)s")
+
+    parser.add_argument("--n_cpus", type=int, required=False, default=NUM_WORKERS_PER_GPU,
+                        help="The number of cpu workers. Default: %(default)s")
+
+    parser.add_argument("--batch_size", type=int, required=False, default=BATCH_SIZE,
+                        help="Batch size. Default: %(default)s")
+
+    parser.add_argument("--convert_to_g", action="store_true", help="Convert the f1 from $/mmol to $/g")
+
+    coprinet_colname = "CoPriNet"
+
+    args = parser.parse_args()
+
+    df = pd.read_csv(args.input_csv_file)
+    nans = df[args.smiles_colname].isna()
+    df = df[~nans]
+    smiles_list = df[args.smiles_colname]
+    predictor = GraphPricePredictor( **vars(args))
+    preds = predictor.yieldPredictions(smiles_list)
+    if args.convert_to_g:
+        from rdkit import Chem
+        def convert_pred(smi, pred):
+            mol = Chem.MolFromSmiles(smi)
+            if mol is None:
+                return np.nan
+            mw = Chem.Descriptors.ExactMolWt(mol)
+            price = np.log(  np.exp(pred)*1000/mw)
+            return price
+
+        preds = Parallel(n_jobs=args.n_cpus)(delayed(convert_pred)(smi, pred) for smi, pred in zip(smiles_list, preds))
+    if args.output_file is None:
+        for smi, pred in zip(smiles_list, preds):
+            print("%s\t%.4f" % (smi, pred))
+    else:
+        df[coprinet_colname] = list(preds)
+        df.to_csv(args.output_file, index=False)
+
+if __name__ == '__main__':
+    main()

coprinet/pricePrediction/preprocessData/prepareDataMol2Price.py

@@ -0,0 +1,185 @@
+import gzip
+import os
+import re
+import sys
+import time
+from functools import reduce
+from itertools import chain
+from multiprocessing import cpu_count
+
+import lmdb
+import psutil
+import joblib
+import torch
+from joblib import Parallel, delayed
+import numpy as np
+from pricePrediction import config
+from pricePrediction.config import USE_MMOL_INSTEAD_GRAM
+from pricePrediction.preprocessData.serializeDatapoints import getExampleId, serializeExample
+from pricePrediction.utils import tryMakedir, getBucketRanges, search_buckedId, EncodedDirNamesAndTemplates
+from .smilesToGraph import compute_nodes_degree, fromPerGramToPerMMolPrice
+
+if config.USE_FEATURES_NET:
+    from .smilesToDescriptors import smiles_to_graph
+
+else:
+    from .smilesToGraph import smiles_to_graph
+
+PER_WORKER_MEMORY_GB = 2
+
+class DataBuilder():
+
+    def __init__(self, n_cpus= config.N_CPUS):
+        if n_cpus is None:
+            mem_gib = psutil.virtual_memory().available / (1024. ** 3)
+            n_cpus = int(max(1, min(cpu_count(), mem_gib // PER_WORKER_MEMORY_GB)))
+
+        self.n_cpus = n_cpus
+
+    def processOneFileOfSmiles(self, encodedDir, fileNum, datasetSplit, fname, nrows=None):
+        print("processing %s"%fname)
+        if fname.endswith(".csv"):
+            open_fun = open
+            decode_line = lambda line: line
+        elif fname.endswith(".csv.gz"):
+            open_fun = gzip.open
+            decode_line = lambda line : line.decode('utf-8')
+        else:
+            raise ValueError("Bad file format")
+
+        cols = ['SMILES','price']
+
+        names = EncodedDirNamesAndTemplates(encodedDir)
+        outFname_base = datasetSplit + "_" + str(fileNum) + "_lmdb"
+        outFname = os.path.join(encodedDir, datasetSplit,outFname_base)
+        env = lmdb.open(outFname,  map_size=10737418240)
+        one_example_graph = smiles_to_graph("CCOCCC")
+        degs = compute_nodes_degree(None)
+        min_in_x = torch.ones(one_example_graph.x.shape[-1])*float('inf')
+        max_in_x = -torch.ones(one_example_graph.x.shape[-1])*float('inf')
+
+        num_examples = 0
+        bucket_ranges = getBucketRanges()
+        n_per_bucket = np.zeros(len(bucket_ranges), dtype= np.int64)
+
+        with open_fun(fname) as f_in:
+            header = decode_line(f_in.readline()).strip().split(",")
+            try:
+                smi_index, price_index = [ header.index(col) for col in cols]
+            except ValueError:
+                smi_index, price_index = 0,1
+            with env.begin(write=True) as sink, \
+                    gzip.open(names.SELECTED_DATAPOINTS_TEMPLATE % (datasetSplit, fileNum), "wt") as f_out:
+                f_out.write("SMILES,price\n")
+                cur_time = time.time()
+                for i, line in enumerate(f_in):
+                    lineArray = decode_line(line).strip().split(",")
+                    smi, price = lineArray[smi_index], lineArray[price_index]
+                    price = float(price)
+                    graph = smiles_to_graph(smi)
+                    if graph is None:
+                        continue
+                    #Save the original smiles-price
+                    f_out.write("%s,%s\n"%(smi, price))
+
+                    # Use the per mmol price
+                    if USE_MMOL_INSTEAD_GRAM:
+                        price = fromPerGramToPerMMolPrice(price, smi)
+                    bucketId = search_buckedId( np.log(price), bucket_ranges)
+                    n_per_bucket[bucketId] += 1
+
+                    degs += compute_nodes_degree([graph])
+                    min_in_x = torch.stack([min_in_x, torch.min(graph.x, 0)[0]]).min(0)[0]
+                    max_in_x = torch.stack([max_in_x, torch.max(graph.x, 0)[0]]).max(0)[0]
+
+                    fileId= getExampleId(outFname_base, num_examples)
+                    sink.put(fileId, serializeExample(price, graph))
+                    num_examples += 1
+                    if nrows is not None and num_examples > nrows:
+                        break
+                    if i % 10000 == 0 and fileNum % self.n_cpus == 0:
+                        new_time = time.time()
+                        print("Current iteration: %d # task: %d  (%.2f s)       " % (i, fileNum, new_time - cur_time), end="\r")
+                        cur_time = new_time
+        if fileNum % self.n_cpus == 0:
+            print()
+        return ((outFname, degs, min_in_x, max_in_x, num_examples, n_per_bucket),)
+
+
+    def getNFeatures(self):
+        one_graph = smiles_to_graph("CCCCCCO")
+        # print(one_graph)
+        return dict(nodes_n_features=one_graph["x"].shape[-1], edges_n_features=one_graph["edge_attr"].shape[-1])
+
+    def prepareDataset(self, inputDir=config.DATASET_DIRNAME, encodedDir=config.ENCODED_DIR, datasetSplit="train",
+                       nrows=None, **kwargs):
+
+        assert datasetSplit in ["train", "val", "test"]
+        print("Computing %s dataset" % datasetSplit)
+        print("Using %d workers for data preparation"%self.n_cpus)
+
+        # os.environ["OMP_NUM_THREADS"] = "1"
+        # os.environ["MKL_NUM_THREADS"] = "1"
+
+        names = EncodedDirNamesAndTemplates(encodedDir)
+
+        tryMakedir(encodedDir, remove=False)
+        tryMakedir(os.path.join(encodedDir, datasetSplit))
+        tryMakedir(names.DIR_RAW_DATA_SELECTED)
+
+
+        fnames = [os.path.join(inputDir, fname) for fname in os.listdir(inputDir) if
+                  re.match(config.RAW_DATA_FILE_SUFFIX, fname) and datasetSplit in fname]
+
+        assert len(fnames) > 0
+
+        results = Parallel(n_jobs=self.n_cpus, batch_size=1,
+                           verbose=10)(delayed(self.processOneFileOfSmiles)(encodedDir, i, datasetSplit, fname, nrows=nrows)
+                                                                        for i, fname in enumerate(fnames))
+        results = chain.from_iterable(results)
+        results = list(results)
+        # print( results )
+        fnames_list, degrees, min_in_x, max_in_x, sizes_lis, n_per_bucket = zip(*results)
+
+        degrees = reduce(lambda prev, x: prev + x, degrees).numpy().tolist()
+        min_in_x = reduce(lambda prev, x: torch.stack([prev, x]).min(0)[0], min_in_x).numpy().tolist()
+        max_in_x = reduce(lambda prev, x: torch.stack([prev, x]).max(0)[0], max_in_x).numpy().tolist()
+
+        n_per_bucket = reduce(lambda prev, x: prev + x, n_per_bucket)
+        metadata_dict = {"name":datasetSplit, "fnames_list":fnames_list, "sizes_list": sizes_lis,
+                      "total_size": sum(sizes_lis), "feature_stats":{"min":min_in_x, "max":max_in_x}}
+        metadata_dict.update(self.getNFeatures())
+        # print(metadata_dict)
+        joblib.dump(metadata_dict,
+                    names.DATASET_METADATA_FNAME_TEMPLATE % datasetSplit)
+        joblib.dump(degrees, names.DEGREES_FNAME)
+
+        # print(n_per_bucket)
+        joblib.dump({"bucket_ranges": getBucketRanges(), "n_per_bucket":n_per_bucket}, names.BUCKETS_FNAME_TEMPLATE % datasetSplit)
+        print("Dataset %s computed" % datasetSplit)
+        return fnames_list
+
+
+
+if __name__ == "__main__":
+    print( " ".join(sys.argv))
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-i", "--inputDir", type=str, default=config.DATASET_DIRNAME, help="Directory where smiles-price pairs are located")
+    parser.add_argument("-o", "--encodedDir", type=str, default=config.ENCODED_DIR)
+    parser.add_argument("-n", "--ncpus", type=int, default=config.N_CPUS)
+    parser.add_argument("--nrows", type=int, default=None, help="The number of rows to process in each file")
+
+    args = vars( parser.parse_args())
+    config.N_CPUS = args.get("ncpus", config.N_CPUS)
+
+    dataBuilder = DataBuilder(n_cpus=config.N_CPUS)
+    dataBuilder.prepareDataset(datasetSplit="train", **args)
+    dataBuilder.prepareDataset(datasetSplit="val", **args)
+    dataBuilder.prepareDataset(datasetSplit="test", **args)
+
+'''
+python -m  pricePrediction.preprocessData.prepareDataMol2Price
+
+
+'''

coprinet/pricePrediction/preprocessData/serializeDatapoints.py

@@ -0,0 +1,15 @@
+import pickle
+import torch
+
+
+def getExampleId(outFname_base, num_example):
+    return ("%s-%08d" % (outFname_base, num_example)).encode("UTF-8")
+
+def serializeExample(price, graph):
+    price = torch.tensor(float(price), dtype=torch.float32)
+    one_data_serialized = pickle.dumps( {"graph": graph, "price" :price})
+    return one_data_serialized
+
+def deserializeExample(exampleBytes):
+    example_dict = pickle.loads(exampleBytes)
+    return example_dict["graph"], example_dict["price"]

coprinet/pricePrediction/preprocessData/smilesToDescriptors.py

@@ -0,0 +1,239 @@
+import torch
+import numpy as np
+from rdkit import Chem
+from rdkit.Chem import rdMolDescriptors
+from rdkit.Chem.Descriptors import ExactMolWt
+from rdkit.ML.Descriptors.CompoundDescriptors import GetAllDescriptorNames
+from torch_geometric.data import Data as GraphData
+from rdkit.ML.Descriptors.MoleculeDescriptors import MolecularDescriptorCalculator
+
+
+
+
+calculator = {}
+def smiles_to_graph(smi):
+    mol = Chem.MolFromSmiles(smi)
+    if mol is None:
+        return None
+    if not calculator:
+        chosen_descriptors = ['BertzCT',
+ 'Chi0',
+ 'Chi0n',
+ 'Chi0v',
+ 'Chi1',
+ 'Chi1n',
+ 'Chi1v',
+ 'Chi2n',
+ 'Chi2v',
+ 'Chi3n',
+ 'Chi3v',
+ 'Chi4n',
+ 'Chi4v',
+ 'EState_VSA1',
+ 'EState_VSA10',
+ 'EState_VSA11',
+ 'EState_VSA2',
+ 'EState_VSA3',
+ 'EState_VSA4',
+ 'EState_VSA5',
+ 'EState_VSA6',
+ 'EState_VSA7',
+ 'EState_VSA8',
+ 'EState_VSA9',
+ 'ExactMolWt',
+ 'FpDensityMorgan1',
+ 'FpDensityMorgan2',
+ 'FpDensityMorgan3',
+ 'FractionCSP3',
+ 'HallKierAlpha',
+ 'HeavyAtomCount',
+ 'HeavyAtomMolWt',
+ 'Kappa1',
+ 'Kappa2',
+ 'Kappa3',
+ 'LabuteASA',
+ 'MaxAbsEStateIndex',
+ 'MaxEStateIndex',
+ 'MinAbsEStateIndex',
+ 'MinAbsPartialCharge',
+ 'MinEStateIndex',
+ 'MinPartialCharge',
+ 'MolLogP',
+ 'MolMR',
+ 'MolWt',
+ 'NHOHCount',
+ 'NOCount',
+ 'NumAliphaticCarbocycles',
+ 'NumAliphaticHeterocycles',
+ 'NumAliphaticRings',
+ 'NumAromaticCarbocycles',
+ 'NumAromaticHeterocycles',
+ 'NumAromaticRings',
+ 'NumHAcceptors',
+ 'NumHDonors',
+ 'NumHeteroatoms',
+ 'NumRotatableBonds',
+ 'NumSaturatedCarbocycles',
+ 'NumSaturatedHeterocycles',
+ 'NumSaturatedRings',
+ 'NumValenceElectrons',
+ 'PEOE_VSA1',
+ 'PEOE_VSA10',
+ 'PEOE_VSA11',
+ 'PEOE_VSA12',
+ 'PEOE_VSA13',
+ 'PEOE_VSA14',
+ 'PEOE_VSA2',
+ 'PEOE_VSA3',
+ 'PEOE_VSA4',
+ 'PEOE_VSA5',
+ 'PEOE_VSA6',
+ 'PEOE_VSA7',
+ 'PEOE_VSA8',
+ 'PEOE_VSA9',
+ 'RingCount',
+ 'SMR_VSA1',
+ 'SMR_VSA10',
+ 'SMR_VSA2',
+ 'SMR_VSA3',
+ 'SMR_VSA4',
+ 'SMR_VSA5',
+ 'SMR_VSA6',
+ 'SMR_VSA7',
+ 'SMR_VSA9',
+ 'SlogP_VSA1',
+ 'SlogP_VSA10',
+ 'SlogP_VSA11',
+ 'SlogP_VSA12',
+ 'SlogP_VSA2',
+ 'SlogP_VSA3',
+ 'SlogP_VSA4',
+ 'SlogP_VSA5',
+ 'SlogP_VSA6',
+ 'SlogP_VSA7',
+ 'SlogP_VSA8',
+ 'TPSA',
+ 'VSA_EState1',
+ 'VSA_EState10',
+ 'VSA_EState2',
+ 'VSA_EState3',
+ 'VSA_EState4',
+ 'VSA_EState5',
+ 'VSA_EState6',
+ 'VSA_EState7',
+ 'VSA_EState8',
+ 'VSA_EState9',
+ 'fr_Al_COO',
+ 'fr_Al_OH',
+ 'fr_Al_OH_noTert',
+ 'fr_ArN',
+ 'fr_Ar_COO',
+ 'fr_Ar_N',
+ 'fr_Ar_NH',
+ 'fr_Ar_OH',
+ 'fr_COO',
+ 'fr_COO2',
+ 'fr_C_O',
+ 'fr_C_O_noCOO',
+ 'fr_C_S',
+ 'fr_HOCCN',
+ 'fr_Imine',
+ 'fr_NH0',
+ 'fr_NH1',
+ 'fr_NH2',
+ 'fr_N_O',
+ 'fr_Ndealkylation1',
+ 'fr_Ndealkylation2',
+ 'fr_Nhpyrrole',
+ 'fr_SH',
+ 'fr_aldehyde',
+ 'fr_alkyl_carbamate',
+ 'fr_alkyl_halide',
+ 'fr_allylic_oxid',
+ 'fr_amide',
+ 'fr_amidine',
+ 'fr_aniline',
+ 'fr_aryl_methyl',
+ 'fr_azide',
+ 'fr_azo',
+ 'fr_barbitur',
+ 'fr_benzene',
+ 'fr_benzodiazepine',
+ 'fr_bicyclic',
+ 'fr_dihydropyridine',
+ 'fr_epoxide',
+ 'fr_ester',
+ 'fr_ether',
+ 'fr_furan',
+ 'fr_guanido',
+ 'fr_halogen',
+ 'fr_hdrzine',
+ 'fr_hdrzone',
+ 'fr_imidazole',
+ 'fr_imide',
+ 'fr_isothiocyan',
+ 'fr_ketone',
+ 'fr_ketone_Topliss',
+ 'fr_lactam',
+ 'fr_lactone',
+ 'fr_methoxy',
+ 'fr_morpholine',
+ 'fr_nitrile',
+ 'fr_nitro',
+ 'fr_nitro_arom',
+ 'fr_nitro_arom_nonortho',
+ 'fr_nitroso',
+ 'fr_oxazole',
+ 'fr_oxime',
+ 'fr_para_hydroxylation',
+ 'fr_phenol',
+ 'fr_phenol_noOrthoHbond',
+ 'fr_phos_acid',
+ 'fr_phos_ester',
+ 'fr_piperdine',
+ 'fr_piperzine',
+ 'fr_priamide',
+ 'fr_pyridine',
+ 'fr_quatN',
+ 'fr_sulfide',
+ 'fr_sulfonamd',
+ 'fr_sulfone',
+ 'fr_term_acetylene',
+ 'fr_tetrazole',
+ 'fr_thiazole',
+ 'fr_thiocyan',
+ 'fr_thiophene',
+ 'fr_unbrch_alkane',
+ 'fr_urea',
+ 'qed']
+
+        calculator["calculator"] = MolecularDescriptorCalculator(chosen_descriptors)
+
+    x = np.array(calculator["calculator"].CalcDescriptors(mol))
+    x = x.reshape((1,-1))
+    x = torch.FloatTensor(x)
+    if torch.isnan(x).any():
+        return None
+    graph = GraphData(x=x, edge_index=torch.zeros(2,0, dtype=torch.long), edge_attr=torch.zeros(0,1))
+    return graph
+
+def test(smi = "C1=CC=C(C=C1)CC(C(=O)O)N"):
+
+    mol = Chem.MolFromSmiles(smi)
+    # print(mol)
+    graph = smiles_to_graph(smi)
+    print(graph.keys)
+    print(graph.x)
+    print(graph.edge_index)
+    print( graph.edge_attr )
+    print( graph.x.shape, graph.edge_index.shape, graph.edge_attr.shape)
+    print(graph.x.max(), graph.x.min())
+    for idx in range(mol.GetNumAtoms()):
+        mol.GetAtomWithIdx(idx).SetProp('molAtomMapNumber', str(mol.GetAtomWithIdx(idx).GetIdx()))
+    from matplotlib import pyplot as plt; from rdkit.Chem import Draw; plt.imshow(Draw.MolsToGridImage([mol], molsPerRow=2)); plt.show()
+
+
+if __name__ == "__main__":
+    smis_list = ["[C@H](C)1CCCO1", "O[C@@H](N)C", "C1=CC=C(C=C1)CC(C(=O)O)N"]
+    for smi in smis_list:
+        test(smi)

coprinet/pricePrediction/preprocessData/smilesToGraph.py

@@ -0,0 +1,160 @@
+import torch
+import numpy as np
+from rdkit import Chem
+from rdkit.Chem.Descriptors import ExactMolWt
+from torch_geometric.data import Data as GraphData
+
+ALLOWED_ATOMIC_NUMS = {6, 7, 8, 9, 15, 16, 17, 35}
+MIN_NUM_ATOMS = 3
+MAX_NUM_ATOMS = 60
+
+def get_node_feats(mol: Chem.Mol, atom: Chem.Atom):
+    '''
+
+    :param mol: The mol from which the node features will be computed. Currently not used
+    :param atom: The atom from which the node features will be computed
+    :return:
+    '''
+    atom_num = atom.GetAtomicNum()
+    valence = atom.GetTotalValence()
+    charge = atom.GetFormalCharge()
+    degree = atom.GetDegree()
+    is_aromatic = atom.GetIsAromatic()
+    return [float(elem) for elem in [atom_num, valence, charge, degree, is_aromatic]]
+    # return [float(elem) for elem in [atom_num, valence, charge, is_aromatic]]
+
+# results = one_of_k_encoding_unk(
+#     atom.GetSymbol(),
+#     [
+#         'B',
+#         'C',
+#         'N',
+#         'O',
+#         'F',
+#         'Si',
+#         'P',
+#         'S',
+#         'Cl',
+#         'As',
+#         'Se',
+#         'Br',
+#         'Te',
+#         'I',
+#         'At',
+#         'other'
+#     ]) + one_of_k_encoding(atom.GetDegree(),
+#                            [0, 1, 2, 3, 4, 5]) + \
+#           [atom.GetFormalCharge(), atom.GetNumRadicalElectrons()] + \
+#           one_of_k_encoding_unk(atom.GetHybridization(), [
+#               Chem.rdchem.HybridizationType.SP, Chem.rdchem.HybridizationType.SP2,
+#               Chem.rdchem.HybridizationType.SP3, Chem.rdchem.HybridizationType.
+#                                 SP3D, Chem.rdchem.HybridizationType.SP3D2, 'other'
+#           ]) + [atom.GetIsAromatic()]
+
+def get_bond_feats(mol: Chem.Mol, bond: Chem.Bond):
+    '''
+
+    :param mol: The mol from which the edge features will be computed. Currently not used
+    :param bond: The bond from which the edge features will be computed
+    :return:
+    '''
+    if bond.GetIsAromatic():
+        bond_num = 1.5
+    else:
+        bond_num = bond.GetBondType()
+        bond_num = float(bond_num)
+    if bond_num > 3:
+        bond_num = 4
+
+    # return [bond_num ]
+    bondInfo = [bond_num, int(bond.GetIsAromatic()), int(bond.GetIsConjugated()), int(bond.IsInRing()) ]
+    # bondInfo += [["STEREONONE", "STEREOANY", "STEREOZ", "STEREOE"].index(str(bond.GetStereo()))]
+
+    return bondInfo
+
+def smiles_to_graph(smi):
+    mol = Chem.MolFromSmiles(smi)
+    if mol is None:
+        return None
+
+    # remove stereo information, such as inward and outward edges
+    Chem.RemoveStereochemistry(mol)
+
+    nodes = []
+    edges_index, edges_attr = [], []
+
+    cur_atom_id = 0
+    rdkit_atomId_to_atom_id = {}
+    for bond in mol.GetBonds():
+        atom1_idx = bond.GetBeginAtomIdx()
+        atom2_idx = bond.GetEndAtomIdx()
+        atom1 = mol.GetAtomWithIdx(atom1_idx)
+        atom2 = mol.GetAtomWithIdx(atom2_idx)
+        atomicNum1 = atom1.GetAtomicNum()
+        atomicNum2 = atom2.GetAtomicNum()
+
+        if atomicNum1 == 0 or atomicNum2 == 0:
+            continue
+        if atomicNum1 not in ALLOWED_ATOMIC_NUMS or atomicNum2 not in ALLOWED_ATOMIC_NUMS:
+            return None
+        if atom1_idx not in rdkit_atomId_to_atom_id:
+            rdkit_atomId_to_atom_id[atom1_idx] = cur_atom_id
+            nodes.append( get_node_feats(mol, atom1) )
+            cur_atom_id += 1
+        if atom2_idx not in rdkit_atomId_to_atom_id:
+            rdkit_atomId_to_atom_id[atom2_idx] = cur_atom_id
+            nodes.append(  get_node_feats(mol, atom2) )
+            cur_atom_id += 1
+        idx1, idx2 = rdkit_atomId_to_atom_id[atom1_idx], rdkit_atomId_to_atom_id[atom2_idx]
+        # print(atom1_idx, idx1, atom2_idx, idx2)
+        edges_index.append([idx1, idx2])
+        edges_index.append([idx2, idx1])
+
+        bond_feats = get_bond_feats(mol, bond)
+        edges_attr.extend([bond_feats, bond_feats])
+
+    x = np.array(nodes, dtype=np.float32)
+
+    if x.shape[0] < MIN_NUM_ATOMS or x.shape[0] > MAX_NUM_ATOMS:
+        return None
+    edges_index = np.array( edges_index).T.copy().astype( np.int64 )
+    edges_attr = np.array(edges_attr, dtype=np.float32)
+    graph_dict = dict(x=x, edge_index=edges_index, edge_attr=edges_attr)
+    graph = GraphData(**{key: torch.tensor(val) for key, val in graph_dict.items()} )
+    return graph
+
+def fromPerGramToPerMMolPrice(price, smi):
+    mol = Chem.MolFromSmiles(smi)
+    if mol is None:
+        return None
+    return price * ExactMolWt(mol) / 1000
+
+def compute_nodes_degree(graphs, max_degree=7):
+    from torch_geometric.utils import degree
+    import torch
+    deg = torch.zeros(max_degree, dtype=torch.long)
+    if graphs is None:
+        return deg
+    for data in graphs:
+        d = degree(data.edge_index[1], num_nodes=data.num_nodes, dtype=torch.long)
+        deg += torch.bincount(d, minlength=deg.numel())[:deg.numel()]
+    return deg
+
+def test(smi = "C1=CC=C(C=C1)CC(C(=O)O)N"):
+
+    mol = Chem.MolFromSmiles(smi)
+    # print(mol)
+    graph = smiles_to_graph(smi)
+    print(graph.keys)
+    print(graph.x)
+    print(graph.edge_index)
+    print( graph.edge_attr )
+    print( graph.x.shape, graph.edge_index.shape, graph.edge_attr.shape)
+    for idx in range(mol.GetNumAtoms()):
+        mol.GetAtomWithIdx(idx).SetProp('molAtomMapNumber', str(mol.GetAtomWithIdx(idx).GetIdx()))
+    from matplotlib import pyplot as plt; from rdkit.Chem import Draw; plt.imshow(Draw.MolsToGridImage([mol], molsPerRow=2)); plt.show()
+
+if __name__ == "__main__":
+    smis_list = ["[C@H](C)1CCCO1", "O[C@@H](N)C", "C1=CC=C(C=C1)CC(C(=O)O)N"]
+    for smi in smis_list:
+        test(smi)

coprinet/pricePrediction/selectFromRawData/selectFromRawDataNonVirtualAll.py

@@ -0,0 +1,41 @@
+import os, sys
+from argparse import ArgumentParser
+
+from pricePrediction.config import TEST_SIZE
+from pricePrediction.selectFromRawData.selectFromRawDataNonVirtualAll_base import RawDataSelector
+from pricePrediction import config
+
+class SelectFromRawDataNonVirtualAll(RawDataSelector):
+
+    DESIRED_PARAMS_TO_ASK= ['full_dataset_fname', 'computed_datadir', 'test_size', 'nrows']
+    def __init__(self, full_dataset_fname: str =
+             os.path.join(config.RAW_DATA_DIRNAME, "mcule_purchasable_full_prices_210616_O9Jw1D_only_READILY.csv.gz"),
+             computed_datadir: str = config.DATASET_DIRNAME,
+             test_size:int=TEST_SIZE, nrows:int=None):
+        '''
+
+        :param str full_dataset_fname: The fname for the whole Mcule csv file (csv.gz)
+        :param str bb_fname: The fname for the building blocks csv file (csv.gz)
+        :param str computed_datadir: The directory where the selected records will be saved
+        :param int test_size: The number of entries to include in the test set
+        :param int nrows: The number of rows to process in the full_dataset_fname. If None, process all rows
+        '''
+        super().__init__(full_dataset_fname, None, computed_datadir, test_size=test_size, nrows=nrows)
+
+    def operate_data(self, data):
+
+        # data = data[ ~ data.iloc[:,-1].isna() ]
+        data = data[self.COLNAMES[:3]]
+        data = data.rename(columns={self.COLNAMES[2]: "price"})
+        del data["Mcule ID"]
+        return data
+
+if __name__ == "__main__":
+    parser = ArgumentParser(description="extract from raw data train/test/val partitions")
+    SelectFromRawDataNonVirtualAll.addParamsToArgParse(parser)
+    args = parser.parse_args()
+    SelectFromRawDataNonVirtualAll(**vars(args)).compute()
+
+    '''
+python -m pricePrediction.selectFromRawData.selectFromRawDataNonVirtualAll
+    '''

coprinet/pricePrediction/selectFromRawData/selectFromRawDataNonVirtualAll_base.py

@@ -0,0 +1,76 @@
+import os
+import dask.dataframe as dd
+import pandas as pd
+
+from pricePrediction import config
+from pricePrediction.ArgParser_base import ArgParseable
+from pricePrediction.config import TEST_SIZE
+
+
+class RawDataSelector(ArgParseable):
+    COLNAMES = "Mcule ID,SMILES,price 1 (USD),amount 1 (mg),delivery time 1 (w.days),available amount 1 (mg)".split(",")
+
+    DESIRED_PARAMS_TO_ASK= ['full_dataset_fname', 'bb_fname', 'computed_datadir', 'test_size', 'nrows']
+    def __init__(self, full_dataset_fname: str, bb_fname: str= config.BUILDING_BLOCKS_FNAME,
+                 computed_datadir: str= config.DATASET_DIRNAME, test_size:int=TEST_SIZE, nrows:int=None):
+        '''
+
+        :param str full_dataset_fname: The fname for the whole Mcule csv file (csv.gz)
+        :param str bb_fname: The fname for the building blocks csv file (csv.gz)
+        :param str computed_datadir: The directory where the selected records will be saved
+        :param int test_size: The number of entries to include in the test set
+        :param int nrows: The number of rows to process in the full_dataset_fname. If None, process all rows
+        '''
+        self.full_dataset_fname = full_dataset_fname
+        self.bb_fname = bb_fname
+        self.computed_datadir = computed_datadir
+        self.test_size = test_size
+        self.nrows = nrows
+
+    def operate_data(self, data):
+        raise NotImplementedError()
+
+    def compute(self):
+        if self.full_dataset_fname.endswith(".gz"):
+            data = pd.read_csv(self.full_dataset_fname, usecols=self.COLNAMES, nrows=self.nrows)
+            data = dd.from_pandas(data, npartitions=30)
+        else:
+            data = dd.read_csv(self.full_dataset_fname, usecols=self.COLNAMES,blocksize=None)
+            if self.nrows is not None:
+                data = data.head(n=self.nrows, compute=False)
+
+        data = self.operate_data(data)
+        self.saveDataSplits(data)
+
+
+    def saveDataSplits(self, data, computed_datadir=None):
+        if computed_datadir is None:
+            computed_datadir = self.computed_datadir
+        print("Saving results to %s"%computed_datadir)
+        if isinstance(data, dd.DataFrame):
+            n_elems = data.shape[0].compute()
+        print("Total size: ", n_elems)
+        test_frac = self.test_size/n_elems
+        data_train, data_test = data.random_split([1-test_frac, test_frac])
+
+        del data
+
+        data_test = data_test.repartition(npartitions= int(n_elems*test_frac//self.test_size)+1)
+        data_test.to_csv(os.path.join(computed_datadir, "mcule_full_test.csv_split_*.csv"), index=False)
+        del data_test
+
+        n_elems = data_train.shape[0].compute()
+        val_frac = self.test_size/n_elems
+        data_train, data_val = data_train.random_split([1-val_frac, val_frac])
+        data_val= data_val.repartition(npartitions= int(n_elems*val_frac//self.test_size)+1)
+
+        data_val.to_csv(os.path.join(computed_datadir, "mcule_full_val.csv_split_*.csv"), index=False)
+        del data_val
+
+        data_train= data_train.repartition(npartitions= int(n_elems*(1-val_frac)//self.test_size)+1)
+        data_train.to_csv(os.path.join(computed_datadir, "mcule_full_train.csv_split_*.csv"), index=False)
+        with open(os.path.join(computed_datadir,"params.txt"), "w") as f:
+            f.write(str(type(self))+"\n" )
+            f.write(str(self.full_dataset_fname)+"\n" )
+            f.write(str(self.bb_fname)+"\n" )
+            f.write(str(self.computed_datadir))

coprinet/pricePrediction/selectFromRawData/selectFromRawDataOnlyVirtualAll.py

@@ -0,0 +1,55 @@
+import os, sys
+import dask.dataframe as dd
+from argparse import ArgumentParser
+
+from pricePrediction.config import TEST_SIZE
+from pricePrediction.selectFromRawData.selectFromRawDataNonVirtualAll_base import RawDataSelector
+import pandas as pd
+from pricePrediction import config
+
+class SelectFromRawDataOnlyVirtualAll(RawDataSelector):
+    DESIRED_PARAMS_TO_ASK= ['full_dataset_fname', 'bb_fname', 'computed_datadir',]
+    def __init__(self, full_dataset_fname: str =
+             os.path.join(config.RAW_DATA_DIRNAME, "mcule_purchasable_full_prices_210319_kLp4Zt_march.csv.gz"),
+             bb_fname: str = config.BUILDING_BLOCKS_FNAME,
+             computed_datadir: str = config.DATASET_DIRNAME,
+             max_size:int=10000000,
+             test_size:int=TEST_SIZE, nrows:int=None):
+        '''
+        :param str full_dataset_fname: The fname for the whole Mcule csv file (csv.gz)
+        :param str bb_fname: The fname for the building blocks csv file (csv.gz)
+        :param str computed_datadir: The directory where the selected records will be saved
+        :param int test_size: The number of entries to include in the test set
+        :param int nrows: The number of rows to process in the full_dataset_fname. If None, process all rows
+        '''
+        super().__init__(full_dataset_fname, bb_fname, computed_datadir,test_size=test_size,nrows=nrows)
+        self.max_size = max_size
+
+    def read_csv(self, csv_name):
+        data = dd.read_csv(csv_name, usecols=self.COLNAMES,blocksize=None)
+        data = data[ data[self.COLNAMES[-1]].isna() ] #To select compounds that are not in stock
+        data = data[self.COLNAMES[:3]]
+        data = data.rename(columns={self.COLNAMES[2]: "price"})
+        return data
+
+    def operate_data(self, data):
+        data = self.read_csv(self.full_dataset_fname)
+        size = data.shape[0].compute()
+        bb = pd.read_csv(self.bb_fname, sep="\t", header=None)
+        bb.columns = ["SMILES", "Mcule ID"]
+        del bb["SMILES"]
+        data = data[ (~data["Mcule ID"].isin(bb["Mcule ID"])) ]
+        size = data.shape[0].compute()
+        print("Total size: ",size)
+        sampling_size = min(size, self.max_size)
+        data = data.sample(frac= sampling_size/size)
+        print("Sampled size: ",sampling_size)
+        data = data[["SMILES", "price"]]
+        del bb
+        return data
+
+if __name__ == "__main__":
+    parser = ArgumentParser(description="extract from raw data train/test/val partitions")
+    SelectFromRawDataOnlyVirtualAll.addParamsToArgParse(parser)
+    args = parser.parse_args()
+    SelectFromRawDataOnlyVirtualAll(**vars(args)).compute()

coprinet/pricePrediction/train/trainNet.py

@@ -0,0 +1,147 @@
+import json
+import sys, os
+
+# import torch
+# torch.multiprocessing.set_start_method('spawn')
+# torch.multiprocessing.set_sharing_strategy('file_system')
+
+from pytorch_lightning import seed_everything
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint, StochasticWeightAveraging
+import pytorch_lightning as pl
+
+from pytorch_lightning.plugins import DDPPlugin
+
+from pricePrediction import config
+from pricePrediction.ArgParser_base import MyArgParser
+from pricePrediction.dataManager.dataManager import GraphPriceDatamodule
+from pricePrediction.evaluation.evaluation import PlotFigsTensorboardCallback, InferencePlotter
+from pricePrediction.nets.netsGraph import PricePredictorModule
+
+def parse_args():
+
+    parser = MyArgParser(prog=None, usage=None, description=None, )
+
+    #### Program args ####
+    parser.add_argument("-e", "--n_epochs", type=int, help="Number of epochs to use", default=1000)
+    parser.add_argument("-m", "--msg", type=str, help="A message describing the run", required=True)
+    parser.add_argument("-r", "--restore", type=str, help="Directory to trained checkpoint to continue training", default=None)
+    parser.add_argument("-c", "--config", type=str, help="Path to json file with the config arguments to create/update the "
+                                                         "model")
+    parser.add_argument("-o", "--default_root_dir", type=str, help="The directory where to save "
+                                                                          "checkpoints and logs", default=None)
+    parser.add_argument("-g", "--gpus", type=int, help="Number of gpus to use", default=1)
+    parser.add_argument("-n", "--num_nodes", type=int, help="Number of nodes to use", default=None)
+    parser.add_argument("-s", "--random_seed", type=int, help="Random seed", default=121)
+
+    parser.add_argument( "--limit_train_batches", type=float, help="Train only with a fraction of the tranining set", default=None)
+    parser.add_argument("--swa", action="store_true", default=False,
+                                help="Perform  Stochastic Weight Averaging")
+    #### Model and data args ####
+
+    group = parser.add_argument_group(title="data")
+    GraphPriceDatamodule.addParamsToArgParse(group)
+    group = parser.add_argument_group(title="model")
+    PricePredictorModule.addParamsToArgParse(group)
+
+    cmd_args = parser.parse_args()
+    del cmd_args['positional arguments']
+    cmd_args['main'] = cmd_args['optional arguments']
+    del cmd_args['optional arguments']
+    return cmd_args
+
+if __name__ == "__main__":
+    print( " ".join(sys.argv))
+    cmd_args = parse_args()
+
+    if cmd_args.get("config"):
+        print("Loading configuration file:", cmd_args.get("config"))
+        with open(cmd_args.get("config")) as f:
+            args = json.load(f)
+
+        for groupName in cmd_args:
+            args[groupName].update(cmd_args[groupName])
+    else:
+        args = cmd_args
+
+    print(args)
+
+    seed_everything(args["main"]["random_seed"])
+
+
+    if args['main']["restore"]:
+        prev_run_dir = os.path.expanduser(args['main']["restore"])
+        # hparams = yaml.load( os.path.join(prev_run_dir, "hparams.yaml"))
+        checkpointsDir = os.path.join(prev_run_dir, "checkpoints")
+        most_recent_checkpoint_fname = max([os.path.join(checkpointsDir, basename) for basename in os.listdir(checkpointsDir)],
+                                           key=os.path.getctime)
+
+        pl_model = PricePredictorModule.load_from_checkpoint(most_recent_checkpoint_fname, **args["model"])
+        data_args = pl_model.hparams.get("data_hparams")
+        data_args.update(args["data"])
+        dataModule = GraphPriceDatamodule(**data_args)
+
+        print("Previous model loaded!!")
+
+    else:
+        dataModule = GraphPriceDatamodule(**args["data"])
+
+        dataModule.prepare_data()
+        nodes_n_features, edges_n_features = dataModule.dims
+        nodes_degree = dataModule.get_nodes_degree()
+        args["model"].update( dict(
+                            nodes_n_features=nodes_n_features, edges_n_features=edges_n_features,
+                            deg=nodes_degree, data_hparams= args["data"], logs_only_in_epoch= args["main"].get("gpus", 0) > 1)
+                                   )
+
+        pl_model = PricePredictorModule( **args["model"] )
+
+    print( pl_model.hparams )
+
+    callbacks = [
+        EarlyStopping(monitor='val_loss', patience=60),
+        ModelCheckpoint(monitor='val_loss', verbose=True),
+        PlotFigsTensorboardCallback(frequency=5, save_csv=True)
+    ]
+
+    if args["main"].get("swa", False):
+        callbacks += [
+            StochasticWeightAveraging(annealing_epochs=config.COSINE_LR_SCHEDULE_N_EPOCHS, swa_lrs= pl_model.hparams.lr)
+        ]
+
+    default_root_dir = args['main'].get('default_root_dir', os.getcwd())
+    trainer_args = dict(gpus=1, max_epochs=args['main']["n_epochs"], progress_bar_refresh_rate=20,
+                        default_root_dir=default_root_dir, auto_lr_find=True,
+                        callbacks = callbacks)
+    if args["main"]["limit_train_batches"]:
+        trainer_args["limit_train_batches"] = args["main"]["limit_train_batches"]
+
+    ngpus = args['main'].get("gpus", -1)
+    if ngpus > 1:
+        trainer_args["gpus"] = ngpus
+        trainer_args["accelerator"] = "ddp"
+        trainer_args["plugins"] = DDPPlugin(find_unused_parameters=False),
+        os.environ["NCCL_NSOCKS_PERTHREAD"]="4"
+        os.environ["NCCL_SOCKET_NTHREADS"] = "2"
+    elif ngpus ==0:
+        trainer_args["gpus"] = 0
+        trainer_args["accelerator"] = "ddp_cpu"
+
+    if args['main'].get("num_nodes", None):
+        trainer_args["num_nodes"] = int(args['main'].get("num_nodes"))
+        trainer_args["accelerator"] = "ddp"
+        raise NotImplementedError()
+        # import torch.distributed as dist
+        # dist.init_process_group("gloo", rank=int(os.environ.get("NODE_RANK")), world_size=trainer_args["num_nodes"] )
+
+    try:
+        print(trainer_args)
+        trainer = pl.Trainer( **trainer_args )
+        trainer.fit(pl_model, dataModule)
+        res = trainer.test(pl_model, dataModule.test_dataloader(), verbose=True)[0]
+        r, __ = InferencePlotter(pl_module=pl_model).pred_and_plot( mode="test")
+        res["test_pearson-r"] = r
+        print(res)
+    finally:
+        if args['main'].get("num_nodes", None):
+            # dist.destroy_process_group()
+            pass

coprinet/pricePrediction/utils.py

@@ -0,0 +1,32 @@
+import os
+import shutil
+import numpy as np
+
+from pricePrediction import config
+
+def tryMakedir(dirname, remove=False):
+    if  os.path.isdir(dirname):
+        if remove:
+            # print("removing %s"%dirname)
+            shutil.rmtree(dirname)
+            os.mkdir(dirname)
+    else:
+        os.mkdir(dirname)
+
+
+def getBucketRanges(x0=1,x1=11, n_buckets=config.NUM_BUCKETS):
+    return np.linspace(x0, x1, num=n_buckets)
+
+
+def search_buckedId(y, bucket_ranges):
+    n_buckets = len(bucket_ranges)
+    return min(n_buckets - 1, np.searchsorted(bucket_ranges, y))
+
+
+class EncodedDirNamesAndTemplates(object):
+    def __init__(self, encodedDir):
+        self.DEGREES_FNAME = os.path.join(encodedDir, "graphs_degrees.joblib.pckl")
+        self.BUCKETS_FNAME_TEMPLATE = os.path.join(encodedDir, "n_per_bucket.%s.joblib.pckl")
+        self.DIR_RAW_DATA_SELECTED= os.path.join(encodedDir, "selected_smiles_price")
+        self.SELECTED_DATAPOINTS_TEMPLATE = os.path.join(self.DIR_RAW_DATA_SELECTED, "selected_smiles_price.%s.%s.csv.gz")
+        self.DATASET_METADATA_FNAME_TEMPLATE = os.path.join(encodedDir, "%s_metadata.joblib.pckl")

requirements.txt

@@ -19,3 +19,8 @@ uvicorn[standard]>=0.30.0
 python-multipart>=0.0.9
 supabase>=2.31.0
 huggingface_hub>=0.28.0
+pytorch-lightning==1.9.5
+torch-geometric==2.8.0
+torchmetrics==1.8.1
+more-itertools==10.3.0
+tqdm>=4.66.0

tools/online_screening.py

@@ -10,7 +10,7 @@ from typing import Any
 
 import numpy as np
 from pymatgen.core import Structure
-from rdkit import Chem
+from rdkit import Chem, RDConfig
 
 from tools.adsorption import predict_benzene, predict_toluene
 from tools.descriptors import compute_scm_eigenvalues
@@ -19,12 +19,7 @@ from tools.toxicity import predict_toxicity
 
 
 ROOT = Path(__file__).resolve().parent.parent
-ALL_DES_ROOT = ROOT.parent
-DOC_ROOT = ALL_DES_ROOT.parent
-SA_DIR = ALL_DES_ROOT / "0629_revise" / "Screening" / "筛选顺序" / "筛选顺序" / "2 配体可合成性"
-PRICE_ROOT = DOC_ROOT / "price_filter"
-COPRINET_ROOT = PRICE_ROOT / "CoPriNet-main"
-PYG201 = PRICE_ROOT / ".deps" / "pyg201_nodeps"
+COPRINET_ROOT = ROOT / "coprinet"
 
 ACTIVE_TOOLS = [
     "parse_cif",
@@ -150,7 +145,7 @@ def check_metals_tool(context: dict[str, Any], cif_path: str) -> dict[str, Any]:
 
 @lru_cache(maxsize=1)
 def _sa_module():
-    path = SA_DIR / "sascore" / "sascorer.py"
+    path = Path(RDConfig.RDContribDir) / "SA_Score" / "sascorer.py"
     spec = importlib.util.spec_from_file_location("mofscreen_sascorer", path)
     if spec is None or spec.loader is None:
         raise RuntimeError(f"cannot load sascorer: {path}")
@@ -210,8 +205,6 @@ def predict_aquatic_toxicity_tool(context: dict[str, Any]) -> dict[str, Any]:
 
 @lru_cache(maxsize=1)
 def _price_predictor():
-    if str(PYG201) not in sys.path:
-        sys.path.insert(0, str(PYG201))
     if str(COPRINET_ROOT) not in sys.path:
         sys.path.insert(0, str(COPRINET_ROOT))