Refactor app.py to integrate DCM functionality and streamline molecule visualization; update requirements.txt to include plotnine and correct dcmnet repository URL.

app.py

@@ -1,170 +1,26 @@
 import streamlit as st
 
-from ase.visualize import view
-
 try:
     from StringIO import StringIO
 except ImportError:
     from io import StringIO
 
 import streamlit.components.v1 as components
-from scipy.spatial.distance import cdist
-
-import ase
-
-import functools
-import e3x
-from flax import linen as nn
-import jax
-import jax.numpy as jnp
-import matplotlib.pyplot as plt
-import numpy as np
-import optax
-
-
-import pandas as pd
-from dcmnet.modules import MessagePassingModel
-from dcmnet.utils import clip_colors, apply_model
-from dcmnet.data import prepare_batches
-from dcmnet.plotting import plot_model
-
-RANDOM_NUMBER = 0
-filename = "test"
-data_key, train_key = jax.random.split(jax.random.PRNGKey(RANDOM_NUMBER), 2)
-
-
-# Model hyperparameters.
-features = 16
-max_degree = 2
-num_iterations = 2
-num_basis_functions = 8
-cutoff = 4.0
-
-# Create models
-DCM1 = MessagePassingModel(
-    features=features,
-    max_degree=max_degree,
-    num_iterations=num_iterations,
-    num_basis_functions=num_basis_functions,
-    cutoff=cutoff,
-    n_dcm=1,
-)
-
-# Create models
-DCM2 = MessagePassingModel(
-    features=features,
-    max_degree=max_degree,
-    num_iterations=num_iterations,
-    num_basis_functions=num_basis_functions,
-    cutoff=cutoff,
-    n_dcm=2,
-)
-
-
-from rdkit import Chem
-from rdkit.Chem import AllChem
-from rdkit.Chem import Draw
-
-def get_grid_points(coordinates):
-    """
-    create a uniform grid of points around the molecule,
-    starting from minimum and maximum coordinates of the molecule (plus minus some padding)
-    :param coordinates:
-    :return:
-    """
-    bounds = np.array([np.min(coordinates, axis=0),
-                       np.max(coordinates, axis=0)])
-    padding = 3.0
-    bounds = bounds + np.array([-1, 1])[:, None] * padding
-    grid_points = np.meshgrid(*[np.linspace(a, b, 15)
-                                for a, b in zip(bounds[0], bounds[1])])
-
-    grid_points = np.stack(grid_points, axis=0)
-    grid_points = np.reshape(grid_points.T, [-1, 3])
-    #  exclude points that are too close to the molecule
-    grid_points = grid_points[
-        #np.where(np.all(cdist(grid_points, coordinates) >= (2.0 - 1e-1), axis=-1))[0]]
-        np.where(np.all(cdist(grid_points, coordinates) >= (2.5 - 1e-1), axis=-1))[0]]
-
-    return grid_points
-
-
-dcm1_weights = pd.read_pickle("wbs/best_0.0_params.pkl")
-dcm2_weights = pd.read_pickle("wbs/dcm2-best_1000.0_params.pkl")
-
-smiles = 'C1NCCCC1'
-
-smiles_mol = Chem.MolFromSmiles(smiles)
-rdkit_mol = Chem.AddHs(smiles_mol) 
-elements = [a.GetSymbol() for a in rdkit_mol.GetAtoms()]
-# Generate a conformation
-AllChem.EmbedMolecule(rdkit_mol)
-coordinates = rdkit_mol.GetConformer(0).GetPositions()
-surface = get_grid_points(coordinates)
-
-for i, atom in enumerate(smiles_mol.GetAtoms()):
-    # For each atom, set the property "molAtomMapNumber" to a custom number, let's say, the index of the atom in the molecule
-    atom.SetProp("atomNote", str(atom.GetIdx()))
-
-smiles_image = Draw.MolToImage(smiles_mol)
-
-# display molecule
-st.image(smiles_image)
-
-
-vdw_surface = surface 
-max_N_atoms = 60
-max_grid_points = 3143
-max_grid_points - len(vdw_surface)
-try:
-    Z = [np.array([int(_) for _ in elements])]
-except:
-    Z = [np.array([ase.data.atomic_numbers[_.capitalize()] for _ in elements])]
-pad_Z = np.array([np.pad(Z[0], ((0,max_N_atoms - len(Z[0]))))])
-pad_coords = np.array([np.pad(coordinates, ((0, max_N_atoms - len(coordinates)), (0,0)))])
-
-pad_vdw_surface = []
-_ = np.pad(vdw_surface, ((0, max_grid_points - len(vdw_surface)), (0,0)), "constant", constant_values=(0, 10000)) 
-pad_vdw_surface.append(_)
-pad_vdw_surface = np.array(pad_vdw_surface)
-
-
-data_batch = dict(
-    atomic_numbers=jnp.asarray(pad_Z),
-    positions=jnp.asarray(pad_coords),
-    mono=jnp.asarray(pad_Z),
-    ngrid=jnp.array([len(vdw_surface)]),
-    esp=jnp.asarray([np.zeros(max_grid_points)]),
-    vdw_surface=jnp.asarray(pad_vdw_surface),
-)
-
-batch_size = 1
-
-psi4_test_batches = prepare_batches(data_key, data_batch, batch_size)
-
-batchID = 0
-errors_train = []
-batch = psi4_test_batches[batchID]
-
-#mono, dipo = apply_model(DCM1, test_weights, batch, batch_size)
-dcm1results = plot_model(DCM1, dcm1_weights, batch, batch_size, 1, plot=False)
-dcm2results = plot_model(DCM2, dcm2_weights, batch, batch_size, 2, plot=False)
-
-atoms = dcm1results["atoms"]
-dcmol = dcm1results["dcmol"]
-dcmol2 = dcm2results["dcmol"]
 
+from dcm_app import run_dcm
 
+smiles = "C1NCCCC1"
+results = run_dcm(smiles)
 
+st.image(results["smiles_image"])
 st.write("Click M to see the distributed charges")
+
 output = StringIO()
-(atoms+dcmol).write(output, format="html")
-data = output.getvalue()
-components.html(data, width=1000, height=1000)
+(results["atoms"] + results["dcmol"]).write(output, format="html")
+components.html(output.getvalue(), width=1000, height=1000)
 
 output = StringIO()
-(atoms+dcmol2).write(output, format="html")
-data = output.getvalue()
-components.html(data, width=1000, height=1000)
+(results["atoms"] + results["dcmol2"]).write(output, format="html")
+components.html(output.getvalue(), width=1000, height=1000)
 
 

dcm_app.py

@@ -0,0 +1,196 @@
+import ase
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pandas as pd
+from rdkit import Chem
+from rdkit.Chem import AllChem
+from rdkit.Chem import Draw
+from scipy.spatial.distance import cdist
+
+from dcmnet.data import prepare_batches
+from dcmnet.modules import MessagePassingModel
+from dcmnet.loss import (
+    esp_loss_eval,
+    esp_loss_pots,
+    esp_mono_loss_pots,
+    get_predictions,
+)
+from dcmnet.plotting import evaluate_dc, create_plots2
+from dcmnet.multimodel import get_atoms_dcmol
+from dcmnet.multipoles import plot_3d
+from dcmnet.utils import apply_model, clip_colors, reshape_dipole
+RANDOM_NUMBER = 0
+data_key, _ = jax.random.split(jax.random.PRNGKey(RANDOM_NUMBER), 2)
+
+# Model hyperparameters.
+features = 16
+max_degree = 2
+num_iterations = 2
+num_basis_functions = 8
+cutoff = 4.0
+
+
+def create_models():
+    dcm1 = MessagePassingModel(
+        features=features,
+        max_degree=max_degree,
+        num_iterations=num_iterations,
+        num_basis_functions=num_basis_functions,
+        cutoff=cutoff,
+        n_dcm=1,
+    )
+    dcm2 = MessagePassingModel(
+        features=features,
+        max_degree=max_degree,
+        num_iterations=num_iterations,
+        num_basis_functions=num_basis_functions,
+        cutoff=cutoff,
+        n_dcm=2,
+    )
+    return dcm1, dcm2
+
+
+def get_grid_points(coordinates):
+    """
+    create a uniform grid of points around the molecule,
+    starting from minimum and maximum coordinates of the molecule (plus minus some padding)
+    :param coordinates:
+    :return:
+    """
+    bounds = np.array([np.min(coordinates, axis=0), np.max(coordinates, axis=0)])
+    padding = 3.0
+    bounds = bounds + np.array([-1, 1])[:, None] * padding
+    grid_points = np.meshgrid(
+        *[np.linspace(a, b, 15) for a, b in zip(bounds[0], bounds[1])]
+    )
+
+    grid_points = np.stack(grid_points, axis=0)
+    grid_points = np.reshape(grid_points.T, [-1, 3])
+    # exclude points that are too close to the molecule
+    grid_points = grid_points[
+        np.where(np.all(cdist(grid_points, coordinates) >= (2.5 - 1e-1), axis=-1))[0]
+    ]
+
+    return grid_points
+
+
+def load_weights():
+    dcm1_weights = pd.read_pickle("wbs/best_0.0_params.pkl")
+    dcm2_weights = pd.read_pickle("wbs/dcm2-best_1000.0_params.pkl")
+    return dcm1_weights, dcm2_weights
+
+
+def prepare_inputs(smiles):
+    smiles_mol = Chem.MolFromSmiles(smiles)
+    rdkit_mol = Chem.AddHs(smiles_mol)
+    elements = [a.GetSymbol() for a in rdkit_mol.GetAtoms()]
+    AllChem.EmbedMolecule(rdkit_mol)
+    coordinates = rdkit_mol.GetConformer(0).GetPositions()
+    surface = get_grid_points(coordinates)
+
+    for atom in smiles_mol.GetAtoms():
+        atom.SetProp("atomNote", str(atom.GetIdx()))
+
+    smiles_image = Draw.MolToImage(smiles_mol)
+
+    vdw_surface = surface
+    max_n_atoms = 60
+    max_grid_points = 3143
+    try:
+        z_values = [np.array([int(_) for _ in elements])]
+    except Exception:
+        z_values = [np.array([ase.data.atomic_numbers[_.capitalize()] for _ in elements])]
+
+    pad_z = np.array([np.pad(z_values[0], ((0, max_n_atoms - len(z_values[0]))))])
+    pad_coords = np.array(
+        [
+            np.pad(
+                coordinates, ((0, max_n_atoms - len(coordinates)), (0, 0))
+            )
+        ]
+    )
+
+    pad_vdw_surface = []
+    padded_surface = np.pad(
+        vdw_surface,
+        ((0, max_grid_points - len(vdw_surface)), (0, 0)),
+        "constant",
+        constant_values=(0, 10000),
+    )
+    pad_vdw_surface.append(padded_surface)
+    pad_vdw_surface = np.array(pad_vdw_surface)
+    n_atoms = np.sum(pad_z != 0)
+    data_batch = dict(
+        atomic_numbers=jnp.asarray(pad_z),
+        Z=jnp.asarray(pad_z),
+        positions=jnp.asarray(pad_coords),
+        R=jnp.asarray(pad_coords),
+        # N is the number of atoms
+        N=jnp.asarray([n_atoms]),
+        mono=jnp.asarray(pad_z),
+        ngrid=jnp.array([len(vdw_surface)]),
+        n_grid=jnp.array([len(vdw_surface)]),
+        esp=jnp.asarray([np.zeros(max_grid_points)]),
+        vdw_surface=jnp.asarray(pad_vdw_surface),
+        espMask=jnp.asarray([np.ones(max_grid_points)], dtype=jnp.bool_),
+    )
+
+    return data_batch, smiles_image
+
+def do_eval(batch, dipo_dc1, mono_dc1, batch_size):
+    esp_errors, mono_pred, _, _ = evaluate_dc(
+        batch,
+        dipo_dc1,
+        mono_dc1,
+        batch_size,
+        1,
+        plot=False,
+
+    )
+
+    atoms, dcmol, grid, esp, esp_dc_pred, idx_cut = create_plots2(
+        mono_dc1, dipo_dc1, batch, batch_size, 1
+    )
+    outDict = {
+        "mono": mono_dc1,
+        "dipo": dipo_dc1,
+        "esp_errors": esp_errors,
+        "atoms": atoms,
+        "dcmol": dcmol,
+        "grid": grid,
+        "esp": esp,
+        "esp_dc_pred": esp_dc_pred,
+        "esp_mono_pred": mono_pred,
+        "idx_cut": idx_cut,
+    }
+    
+    return outDict
+
+def run_dcm(smiles="C1NCCCC1"):
+    dcm1, dcm2 = create_models()
+    dcm1_weights, dcm2_weights = load_weights()
+    data_batch, smiles_image = prepare_inputs(smiles)
+
+    batch_size = 1
+    psi4_test_batches = prepare_batches(data_key, data_batch, batch_size)
+    batch = psi4_test_batches[0]
+
+    mono_dc1, dipo_dc1 = apply_model(dcm1, dcm1_weights, batch, batch_size)
+    mono_dc2, dipo_dc2 = apply_model(dcm2, dcm2_weights, batch, batch_size)
+
+    dcm1_results = do_eval(batch, dipo_dc1, mono_dc1, batch_size)
+    dcm2_results = do_eval(batch, dipo_dc2, mono_dc2, batch_size)
+
+    return {
+        "smiles_image": smiles_image,
+        "atoms": dcm1_results["atoms"],
+        "dcmol": dcm1_results["dcmol"],
+        "dcmol2": dcm2_results["dcmol"],
+    }
+
+
+if __name__ == "__main__":
+    smiles = "C1NCCCC1"
+    results = run_dcm(smiles)
+    print(results)

requirements.txt

@@ -3,5 +3,5 @@ patchworklib
 ase
 scipy
 rdkit
-#tqdm
-dcmnet @ git+https://github.com/EricBoittier/jaxeq@main
+plotnine
+dcmnet @ git+https://github.com/EricBoittier/dcmnet@main