Removing roitberg examples

import numpy as np

from flask import request, abort, Flask

import flask

webapp = Flask(__name__)

@webapp.route('/potential', methods=["POST"])

def potential():

  content = request.get_json(force=True)

  if not content or not 'X' in content:

    abort(400)

  X = np.array(content['X'])

  x0 = X[:, 1:]

  a0 = X[:, :1]

  result = webapp.model.pred_one(x0, a0)

  return flask.jsonify({'y': result.tolist()[0]}), 200

@webapp.route('/gradient', methods=["POST"])

def index():

  content = request.get_json(force=True)

  if not content or not 'X' in content:

    abort(400)

  X = np.array(content['X'])

  num_atoms = X.shape[0]

  x0 = X[:, 1:]

  a0 = X[:, :1]

  res = webapp.model.grad_one(x0, a0)

  res = res.reshape((num_atoms, 3))

  return flask.jsonify({'grad': res.tolist()}), 200

@webapp.route('/minimize', methods=["POST"])

def minimize():

  content = request.get_json(force=True)

  if not content or not 'X' in content:

    abort(400)

  X = np.array(content['X'])

  constraints = None

  if 'constraints' in content:

    constraints = content['constraints']

    print('setting constraints')

  num_atoms = X.shape[0]

  x0 = X[:, 1:]

  a0 = X[:, :1]

  res = webapp.model.minimize_structure(x0, a0, constraints)

  res = res.reshape((num_atoms, 3))

  y = webapp.model.pred_one(res, a0).tolist()[0]

  return flask.jsonify({'X': res.tolist(), 'y': y}), 200

# Written by Roman Zubatyuk and Justin S. Smith

# Modified by Yutong Zhao to make python2 compatible

import h5py

import numpy as np

import platform

import os

PY_VERSION = int(platform.python_version().split('.')[0]) > 3

class datapacker(object):

  def __init__(self, store_file, mode='w-', complib='gzip', complevel=6):

    """Wrapper to store arrays within HFD5 file

        """

    # opening file

    self.store = h5py.File(store_file, mode=mode)

    self.clib = complib

    self.clev = complevel

  def store_data(self, store_loc, **kwargs):

    """Put arrays to store

        """

    #print(store_loc)

    g = self.store.create_group(store_loc)

    for k, v, in kwargs.items():

      #print(type(v[0]))

      #print(k)

      if type(v) == list:

        if len(v) != 0:

          if type(v[0]) is np.str_ or type(v[0]) is str:

            v = [a.encode('utf8') for a in v]

      g.create_dataset(

          k, data=v, compression=self.clib, compression_opts=self.clev)

  def cleanup(self):

    """Wrapper to close HDF5 file

        """

    self.store.close()

class anidataloader(object):

  ''' Contructor '''

  def __init__(self, store_file):

    if not os.path.exists(store_file):

      exit('Error: file not found - ' + store_file)

    self.store = h5py.File(store_file)

  ''' Group recursive iterator (iterate through all groups in all branches and return datasets in dicts) '''

  def h5py_dataset_iterator(self, g, prefix=''):

    for key in g.keys():

      item = g[key]

      path = '{}/{}'.format(prefix, key)

      keys = [i for i in item.keys()]

      if isinstance(item[keys[0]], h5py.Dataset):  # test for dataset

        data = {'path': path}

        for k in keys:

          if not isinstance(item[k], h5py.Group):

            dataset = np.array(item[k].value)

            if type(dataset) is np.ndarray:

              if dataset.size != 0:

                if type(dataset[0]) is np.bytes_:

                  dataset = [a.decode('ascii') for a in dataset]

            data.update({k: dataset})

        yield data

      else:  # test for group (go down)

        for s in self.h5py_dataset_iterator(item, path):

          yield s

  ''' Default class iterator (iterate through all data) '''

  def __iter__(self):

    for data in self.h5py_dataset_iterator(self.store):

      yield data

  ''' Returns a list of all groups in the file '''

  def get_group_list(self):

    return [g for g in self.store.values()]

  ''' Allows interation through the data in a given group '''

  def iter_group(self, g):

    for data in self.h5py_dataset_iterator(g):

      yield data

  ''' Returns the requested dataset '''

  def get_data(self, path, prefix=''):

    item = self.store[path]

    path = '{}/{}'.format(prefix, path)

    keys = [i for i in item.keys()]

    data = {'path': path}

    # print(path)

    for k in keys:

      if not isinstance(item[k], h5py.Group):

        dataset = np.array(item[k].value)

        if type(dataset) is np.ndarray:

          if dataset.size != 0:

            if type(dataset[0]) is np.bytes_:

              dataset = [a.decode('ascii') for a in dataset]

        data.update({k: dataset})

    return data

  ''' Returns the number of groups '''

  def group_size(self):

    return len(self.get_group_list())

  def size(self):

    count = 0

    for g in self.store.values():

      count = count + len(g.items())

    return count

  ''' Close the HDF5 file '''

  def cleanup(self):

    self.store.close()

if __name__ == "__main__":

  base_dir = os.environ["ROITBERG_ANI"]

  # Number of conformations in each file increases exponentially.

  # Start with a smaller dataset before continuing. Use all of them

  # for production

  hdf5files = [

      'ani_gdb_s01.h5', 'ani_gdb_s02.h5', 'ani_gdb_s03.h5', 'ani_gdb_s04.h5',

      'ani_gdb_s05.h5', 'ani_gdb_s06.h5', 'ani_gdb_s07.h5', 'ani_gdb_s08.h5'

  ]

  hdf5files = [os.path.join(base_dir, f) for f in hdf5files]

  for hdf5file in hdf5files:

    print("processing", hdf5file)

    adl = anidataloader(hdf5file)

    for data in adl:

      # Extract the data

      P = data['path']

      R = data['coordinates']

      E = data['energies']

      S = data['species']

      smi = data['smiles']

import numpy as np

import os

import tensorflow as tf

import deepchem as dc

import pyanitools as pya

import app

def convert_species_to_atomic_nums(s):

  PERIODIC_TABLE = {"H": 1, "C": 6, "N": 7, "O": 8}

  res = []

  for k in s:

    res.append(PERIODIC_TABLE[k])

  return np.array(res, dtype=np.float32)

# replace with your own scratch directory

data_dir = "/media/yutong/datablob/datasets"

model_dir = "/media/yutong/datablob/models"

all_dir = os.path.join(data_dir, "all")

test_dir = os.path.join(data_dir, "test")

fold_dir = os.path.join(data_dir, "fold")

train_dir = os.path.join(fold_dir, "train")

valid_dir = os.path.join(fold_dir, "valid")

def load_roiterberg_ANI(mode="atomization"):

  """

  Load the ANI dataset.

  Parameters

  ----------

  mode: str

    Accepted modes are "relative", "atomization", or "absolute". These settings are used

    to adjust the dynamic range of the model, with absolute having the greatest and relative

    having the lowest. Note that for atomization we approximate the single atom energy

    using a different level of theory

  Returns

  -------

  tuples

    Elements returned are 3-tuple (a,b,c) where and b are the train and test datasets, respectively,

    and c is an array of indices denoting the group of each

  """

  if "ROITBERG_ANI" not in os.environ:

    raise ValueError(

        "Please set environment variable ROITBERG_ANI to where the ani_dgb_s0x.h5 files are."

    )

  base_dir = os.environ["ROITBERG_ANI"]

  # Number of conformations in each file increases exponentially.

  # Start with a smaller dataset before continuing. Use all of them

  # for production

  hdf5files = [

      'ani_gdb_s01.h5',

      'ani_gdb_s02.h5',

      # 'ani_gdb_s03.h5',

      # 'ani_gdb_s04.h5',

      # 'ani_gdb_s05.h5',

      # 'ani_gdb_s06.h5',

      # 'ani_gdb_s07.h5',

      # 'ani_gdb_s08.h5'

  ]

  hdf5files = [os.path.join(base_dir, f) for f in hdf5files]

  groups = []

  def shard_generator():

    shard_size = 4096 * 64

    row_idx = 0

    group_idx = 0

    X_cache = []

    y_cache = []

    w_cache = []

    ids_cache = []

    for hdf5file in hdf5files:

      adl = pya.anidataloader(hdf5file)

      for data in adl:

        # Extract the data

        P = data['path']

        R = data['coordinates']

        E = data['energies']

        S = data['species']

        smi = data['smiles']

        if len(S) > 23:

          print("skipping:", smi, "due to atom count.")

          continue

        # Print the data

        print("Processing: ", P)

        print("  Smiles:      ", "".join(smi))

        print("  Symbols:     ", S)

        print("  Coordinates: ", R.shape)

        print("  Energies:    ", E.shape)

        Z_padded = np.zeros((23,), dtype=np.float32)

        nonpadded = convert_species_to_atomic_nums(S)

        Z_padded[:nonpadded.shape[0]] = nonpadded

        if mode == "relative":

          offset = np.amin(E)

        elif mode == "atomization":

          # self-interaction energies taken from

          # https://github.com/isayev/ANI1_dataset README

          atomizationEnergies = {

              0: 0,

              1: -0.500607632585,

              6: -37.8302333826,

              7: -54.5680045287,

              8: -75.0362229210

          }

          offset = 0

          for z in nonpadded:

            offset -= atomizationEnergies[z]

        elif mode == "absolute":

          offset = 0

        else:

          raise Exception("Unsupported mode: ", mode)

        for k in range(len(E)):

          R_padded = np.zeros((23, 3), dtype=np.float32)

          R_padded[:R[k].shape[0], :R[k].shape[1]] = R[k]

          X = np.concatenate([np.expand_dims(Z_padded, 1), R_padded], axis=1)

          y = E[k] - offset

          if len(X_cache) == shard_size:

            yield np.array(X_cache), np.array(y_cache), np.array(

                w_cache), np.array(ids_cache)

            X_cache = []

            y_cache = []

            w_cache = []

            ids_cache = []

          else:

            X_cache.append(X)

            y_cache.append(np.array(y).reshape((1,)))

            w_cache.append(np.array(1).reshape((1,)))

            ids_cache.append(row_idx)

            row_idx += 1

            groups.append(group_idx)

        group_idx += 1

    # flush once more at the end

    if len(X_cache) > 0:

      yield np.array(X_cache), np.array(y_cache), np.array(w_cache), np.array(

          ids_cache)

  tasks = ["ani"]

  dataset = dc.data.DiskDataset.create_dataset(

      shard_generator(), tasks=tasks, data_dir=all_dir)

  print("Number of groups", np.amax(groups))

  splitter = dc.splits.RandomGroupSplitter(groups)

  train_dataset, test_dataset = splitter.train_test_split(

      dataset, train_dir=fold_dir, test_dir=test_dir, frac_train=.8)

  return train_dataset, test_dataset, groups

def broadcast(dataset, metadata):

  new_metadata = []

  for (_, _, _, ids) in dataset.itershards():

    for idx in ids:

      new_metadata.append(metadata[idx])

  return new_metadata

if __name__ == "__main__":

  max_atoms = 23

  batch_size = 64  # CHANGED FROM 16

  layer_structures = [128, 128, 64]

  atom_number_cases = [1, 6, 7, 8]

  metric = [

      dc.metrics.Metric(dc.metrics.mean_absolute_error, mode="regression"),

      dc.metrics.Metric(dc.metrics.pearson_r2_score, mode="regression")

  ]

  print("Fitting new model...")

  train_valid_dataset, test_dataset, all_groups = load_roiterberg_ANI(

      mode="atomization")

  splitter = dc.splits.RandomGroupSplitter(

      broadcast(train_valid_dataset, all_groups))

  print("Performing 1-fold split...")

  train_dataset, valid_dataset = splitter.train_test_split(

      train_valid_dataset, train_dir=train_dir, test_dir=valid_dir)

  transformers = [

      dc.trans.NormalizationTransformer(

          transform_y=True, dataset=train_dataset)

  ]

  print("Total training set shape: ", train_dataset.get_shape())

  for transformer in transformers:

    train_dataset = transformer.transform(train_dataset)

    valid_dataset = transformer.transform(valid_dataset)

    test_dataset = transformer.transform(test_dataset)

  model = dc.models.ANIRegression(

      1,

      max_atoms,

      layer_structures=layer_structures,

      atom_number_cases=atom_number_cases,

      batch_size=batch_size,

      learning_rate=0.001,

      use_queue=True,

      model_dir=model_dir,

      mode="regression")

  #   # For production, set nb_epoch to 100+

  for i in range(10):

    model.fit(train_dataset, nb_epoch=1, checkpoint_interval=100)

  print("Evaluating model")

  train_scores = model.evaluate(train_dataset, metric, transformers)

  valid_scores = model.evaluate(valid_dataset, metric, transformers)

  test_scores = model.evaluate(test_dataset, metric, transformers)

  # print("Train scores")

  # print(train_scores)

  print("Validation scores")

  print(valid_scores)

  print("Test scores")

  print(test_scores)

  coords = np.array([

      [0.3, 0.4, 0.5],

      [0.8, 0.2, 0.3],

      [0.1, 0.3, 0.8],

  ])

  atomic_nums = np.array([1, 8, 1])

  print("Prediction of a single test set structure:")

  print(model.pred_one(coords, atomic_nums))

  print("Gradient of a single test set structure:")

  print(model.grad_one(coords, atomic_nums))

  # print("Minimization of a single test set structure:")

  # print(model.minimize_structure(coords, atomic_nums))

  app.webapp.model = model

  app.webapp.run(host='0.0.0.0', debug=False)