Enhance DCM evaluation by adding normalization for batch inputs and updating the evaluation function to accept a variable number of DCMs; include print statement for results in app.py.

app.py

@@ -11,7 +11,7 @@ from dcm_app import run_dcm
 
 smiles = "C1NCCCC1"
 results = run_dcm(smiles)
-
+print(results)
 st.image(results["smiles_image"])
 st.write("Click M to see the distributed charges")
 

dcm_app.py

@@ -138,19 +138,28 @@ def prepare_inputs(smiles):
 
     return data_batch, smiles_image
 
-def do_eval(batch, dipo_dc1, mono_dc1, batch_size):
+def do_eval(batch, dipo_dc1, mono_dc1, batch_size, n_dcm):
     esp_errors, mono_pred, _, _ = evaluate_dc(
         batch,
         dipo_dc1,
         mono_dc1,
         batch_size,
-        1,
+        n_dcm,
         plot=False,
 
     )
 
-    atoms, dcmol, grid, esp, esp_dc_pred, idx_cut = create_plots2(
-        mono_dc1, dipo_dc1, batch, batch_size, 1
+    n_atoms = int(batch.get("N", jnp.array([jnp.count_nonzero(batch["Z"])]))[0])
+    n_dcm = mono_dc1.shape[-1]
+    atoms = ase.Atoms(
+        numbers=np.array(batch["Z"][:n_atoms]),
+        positions=np.array(batch["R"][:n_atoms]),
+    )
+    dcm_positions = np.array(dipo_dc1).reshape(-1, 3)[: n_atoms * n_dcm]
+    dcm_charges = np.array(mono_dc1).reshape(-1)[: n_atoms * n_dcm]
+    dcmol = ase.Atoms(
+        ["X" if _ > 0 else "He" for _ in dcm_charges],
+        dcm_positions,
     )
     outDict = {
         "mono": mono_dc1,
@@ -158,15 +167,30 @@ def do_eval(batch, dipo_dc1, mono_dc1, batch_size):
         "esp_errors": esp_errors,
         "atoms": atoms,
         "dcmol": dcmol,
-        "grid": grid,
-        "esp": esp,
-        "esp_dc_pred": esp_dc_pred,
+        "grid": None,
+        "esp": None,
+        "esp_dc_pred": None,
         "esp_mono_pred": mono_pred,
-        "idx_cut": idx_cut,
+        "idx_cut": None,
     }
     
     return outDict
 
+def normalize_batch(batch):
+    vdw_surface = batch.get("vdw_surface")
+    if vdw_surface is not None and vdw_surface.ndim == 4 and vdw_surface.shape[1] == 1:
+        batch["vdw_surface"] = vdw_surface.squeeze(axis=1)
+
+    esp = batch.get("esp")
+    if esp is not None and esp.ndim == 3 and esp.shape[1] == 1:
+        batch["esp"] = esp.squeeze(axis=1)
+
+    esp_mask = batch.get("espMask")
+    if esp_mask is not None and esp_mask.ndim == 3 and esp_mask.shape[1] == 1:
+        batch["espMask"] = esp_mask.squeeze(axis=1)
+
+    return batch
+
 def run_dcm(smiles="C1NCCCC1"):
     dcm1, dcm2 = create_models()
     dcm1_weights, dcm2_weights = load_weights()
@@ -174,13 +198,13 @@ def run_dcm(smiles="C1NCCCC1"):
 
     batch_size = 1
     psi4_test_batches = prepare_batches(data_key, data_batch, batch_size)
-    batch = psi4_test_batches[0]
+    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)
 
-    dcm1_results = do_eval(batch, dipo_dc1, mono_dc1, batch_size)
-    dcm2_results = do_eval(batch, dipo_dc2, mono_dc2, batch_size)
+    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,
@@ -190,7 +214,7 @@ def run_dcm(smiles="C1NCCCC1"):
     }
 
 
-if __name__ == "__main__":
-    smiles = "C1NCCCC1"
-    results = run_dcm(smiles)
-    print(results)
+# if __name__ == "__main__":
+#     smiles = "C1NCCCC1"
+#     results = run_dcm(smiles)
+#     print(results)