Enhance DCM-Net functionality in app.py by integrating support for multiple DCM layers, updating the rendering of 3D models, and improving the user interface with sliders and buttons. Extend dcm_app.py to create and load models for DCM-3 and DCM-4, and adjust weight loading to accommodate additional models. Update requirements.txt to include new dependencies for molecule visualization.

app.py

@@ -1,5 +1,10 @@
 import streamlit as st
-
+from streamlit_ketcher import st_ketcher
+from stmol import showmol
+import py3Dmol
+import io
+import ase.io as ase_io
+import numpy as np
 try:
     from StringIO import StringIO
 except ImportError:
@@ -7,20 +12,129 @@ except ImportError:
 
 import streamlit.components.v1 as components
 
+
+import streamlit as st
+from streamlit_ketcher import st_ketcher
+
+st.title("DCM-Net")
+smiles = "OC(=O)[C@@H]1CCCN1"
+# Create a molecule editor
+smiles = st_ketcher(smiles)
+
+# Output the SMILES string of the drawn molecule
+st.write(f"SMILES: {smiles}")
+
+
+# color charges by magnitude
+def color_charges(mono, nDCM, atoms):
+    from matplotlib import cm
+    from matplotlib.colors import Normalize
+
+    norm = Normalize(vmin=-1, vmax=1)
+    cmap = cm.get_cmap("bwr")
+    mappable = cm.ScalarMappable(norm=norm, cmap=cmap)
+    elem = atoms.numbers
+    pccolors = mappable.to_rgba(mono.flatten()[: len(atoms.numbers)])
+    from ase.data.colors import jmol_colors
+
+    atomcolors = [jmol_colors[_] for _ in elem]
+    atomcolors_ = []
+    for _ in atomcolors:
+        atomcolors_.append(np.append(_, 0.015))
+    dcmcolors = mappable.to_rgba(mono.flatten()[: len(elem) * nDCM])
+    return dcmcolors
+
+
 from dcm_app import run_dcm
+# Create a 3D view
+def render_3d(atoms, dcm_charges=None, dcmcolors=None):
+    view = py3Dmol.view()
+
+    view.startjs += '''\n
+    let customColorize = function(atom){
+        // attribute elem is from https://3dmol.csb.pitt.edu/doc/AtomSpec.html#elem
+        if (atom.elem === 'X'){
+            return "#FF0000" // red
+        }
+        else if (atom.elem === 'He'){
+            return "#0000FF" // blue
+        }
+        else{
+            return $3Dmol.getColorFromStyle(atom, {colorscheme: "whiteCarbon"});
+        }
+    }    
+    \n'''
+
+    if atoms is not None:
+        # Write structure to xyz file.
+        xyz = io.StringIO()
+        ase_io.write(xyz, atoms, format='xyz')
+        
+        view.addModel(xyz.getvalue(), 'xyz')
+        view.setStyle({'stick': {'radius': 0.1}, 'sphere': {'scale': 0.25}})
+
+    if dcm_charges is not None:
+        xyz_dcm = io.StringIO()
+        ase_io.write(xyz_dcm, dcm_charges, format='xyz')
+        view.addModel(xyz_dcm.getvalue(), 'xyz')
+        view.setStyle({'model': -1}, {'stick': {'radius': 0.01}, 'sphere': {'scale': 0.05}})
+    view.startjs = view.startjs.replace('"customColorize"', 'customColorize')
+    view.zoomTo()
+    return view
+
+# create a slider for n_dcm
+n_dcm = st.slider("Number of DCM layers", min_value=1, max_value=4, value=1)
+
+# create a button to run the DCM-Net
+if st.button("Run DCM-Net"):
+    results = run_dcm(smiles, n_dcm)
+    print(results)
+
+    st.subheader("Model Output")
+    # ase object
+    ase_atoms = results["atoms"]
+    
+    # Display the 3D model
+    
+    col1, col2 = st.columns(2)
+    with col1:
+        showmol(render_3d(ase_atoms), height=300, width=1000)
+    with col2:
+        st.image(results["smiles_image"])
 
-smiles = "C1NCCCC1"
-results = run_dcm(smiles)
-print(results)
-st.image(results["smiles_image"])
-st.write("Click M to see the distributed charges")
+    st.subheader("DCM-1")
+    dcmcolors1 = color_charges(results["mono_dc1"], 1, ase_atoms)
+    # two columns showing combined and just the dcmol
+    col1, col2 = st.columns(2)
+    with col1:
+        showmol(render_3d(ase_atoms, results["dcmol"]), height=300, width=1000)
+    with col2:
+        showmol(render_3d(None, results["dcmol"], dcmcolors1), height=300, width=1000)
 
-output = StringIO()
-(results["atoms"] + results["dcmol"]).write(output, format="html")
-components.html(output.getvalue(), width=1000, height=1000)
+    if n_dcm >= 2:
+        st.subheader("DCM-2")
+        dcmcolors2 = color_charges(results["mono_dc2"], 2, ase_atoms)
+        col1, col2 = st.columns(2)
+        with col1:
+            showmol(render_3d(ase_atoms, results["dcmol2"]), height=300, width=1000)
+        with col2:
+            showmol(render_3d(None, results["dcmol2"], dcmcolors2), height=300, width=1000)
 
-output = StringIO()
-(results["atoms"] + results["dcmol2"]).write(output, format="html")
-components.html(output.getvalue(), width=1000, height=1000)
+    if n_dcm >= 3:
+        st.subheader("DCM-3")
+        dcmcolors3 = color_charges(results["mono_dc3"], 3, ase_atoms)
+        col1, col2 = st.columns(2)
+        with col1:
+            showmol(render_3d(ase_atoms, results["dcmol3"]), height=300, width=1000)
+        with col2:
+            showmol(render_3d(None, results["dcmol3"], dcmcolors3), height=300, width=1000)
 
+    if n_dcm >= 4:
+        st.subheader("DCM-4")
+        dcmcolors4 = color_charges(results["mono_dc4"], 4, ase_atoms)
+        col1, col2 = st.columns(2)
+        with col1:
+            showmol(render_3d(ase_atoms, results["dcmol4"]), height=300, width=1000)
+        with col2:
+            showmol(render_3d(None, results["dcmol4"], dcmcolors4), height=300, width=1000)
 

dcm_app.py

@@ -48,7 +48,23 @@ def create_models():
         cutoff=cutoff,
         n_dcm=2,
     )
-    return dcm1, dcm2
+    dcm3 = MessagePassingModel(
+        features=features,
+        max_degree=max_degree,
+        num_iterations=num_iterations,
+        num_basis_functions=num_basis_functions,
+        cutoff=cutoff,
+        n_dcm=3,
+    )
+    dcm4 = MessagePassingModel(
+        features=features,
+        max_degree=max_degree,
+        num_iterations=num_iterations,
+        num_basis_functions=num_basis_functions,
+        cutoff=cutoff,
+        n_dcm=4,
+    )
+    return dcm1, dcm2, dcm3, dcm4
 
 
 def get_grid_points(coordinates):
@@ -83,7 +99,7 @@ def restore_arrays(obj):
         if any(isinstance(value, dict) for value in restored):
             return restored
         try:
-            return np.array(restored)
+            return jnp.asarray(restored)
         except Exception:
             return restored
     return obj
@@ -96,9 +112,11 @@ def load_json_dict(path):
 
 
 def load_weights():
-    dcm1_weights = load_json_dict("wbs/best_0.0_params.json")
-    dcm2_weights = load_json_dict("wbs/dcm2-best_1000.0_params.json")
-    return dcm1_weights, dcm2_weights
+    dcm1_weights = load_json_dict("wbs/best_0.0_params_dict.json")
+    dcm2_weights = load_json_dict("wbs/dcm2-best_1000.0_params_dict.json")
+    dcm3_weights = load_json_dict("wbs/dcm3-best_1000.0_params_dict.json")
+    dcm4_weights = load_json_dict("wbs/dcm4-best_1000.0_params_dict.json")
+    return dcm1_weights, dcm2_weights, dcm3_weights, dcm4_weights
 
 
 def prepare_inputs(smiles):
@@ -211,30 +229,41 @@ def normalize_batch(batch):
 
     return batch
 
-def run_dcm(smiles="C1NCCCC1"):
-    dcm1, dcm2 = create_models()
-    dcm1_weights, dcm2_weights = load_weights()
+def run_dcm(smiles="C1NCCCC1", n_dcm=1):
+    dcm1, dcm2, dcm3, dcm4 = create_models()
+    dcm1_weights, dcm2_weights, dcm3_weights, dcm4_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 = normalize_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)
+    if n_dcm >= 1:
+        mono_dc1, dipo_dc1 = apply_model(dcm1, dcm1_weights, batch, batch_size)
+        dcm1_results = do_eval(batch, dipo_dc1, mono_dc1, batch_size, n_dcm=1)
 
-    dcm1_results = do_eval(batch, dipo_dc1, mono_dc1, batch_size, n_dcm=1)
-    dcm2_results = do_eval(batch, dipo_dc2, mono_dc2, batch_size, n_dcm=2)
-
-    return {
-        "smiles_image": smiles_image,
+    results = {
         "atoms": dcm1_results["atoms"],
         "dcmol": dcm1_results["dcmol"],
-        "dcmol2": dcm2_results["dcmol"],
+        "smiles_image": smiles_image,
+        "mono_dc1": mono_dc1,
     }
 
+    if n_dcm >= 2:
+        mono_dc2, dipo_dc2 = apply_model(dcm2, dcm2_weights, batch, batch_size)
+        dcm2_results = do_eval(batch, dipo_dc2, mono_dc2, batch_size, n_dcm=2)
+        results["dcmol2"] = dcm2_results["dcmol"]
+        results["mono_dc2"] = mono_dc2
+    if n_dcm >= 3:
+        mono_dc3, dipo_dc3 = apply_model(dcm3, dcm3_weights, batch, batch_size)
+        dcm3_results = do_eval(batch, dipo_dc3, mono_dc3, batch_size, n_dcm=3)
+        results["dcmol3"] = dcm3_results["dcmol"]
+        results["mono_dc3"] = mono_dc3
+    if n_dcm >= 4:
+        mono_dc4, dipo_dc4 = apply_model(dcm4, dcm4_weights, batch, batch_size)
+        dcm4_results = do_eval(batch, dipo_dc4, mono_dc4, batch_size, n_dcm=4)
+        results["dcmol4"] = dcm4_results["dcmol"]
+        results["mono_dc4"] = mono_dc4
+    
 
-# if __name__ == "__main__":
-#     smiles = "C1NCCCC1"
-#     results = run_dcm(smiles)
-#     print(results)
+    return results

requirements.txt

@@ -5,3 +5,7 @@ scipy
 rdkit
 plotnine
 dcmnet @ git+https://github.com/EricBoittier/dcmnet@main
+streamlit_ketcher 
+stmol
+py3Dmol==2.0.0.post2
+ipython_genutils