.

app.py

@@ -50,6 +50,15 @@ DCM1 = MessagePassingModel(
     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
@@ -80,7 +89,8 @@ def get_grid_points(coordinates):
     return grid_points
 
 
-test_weights = pd.read_pickle("wbs/best_0.0_params.pkl")
+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'
 
@@ -137,19 +147,24 @@ errors_train = []
 batch = psi4_test_batches[batchID]
 
 #mono, dipo = apply_model(DCM1, test_weights, batch, batch_size)
-dcm1results = plot_model(DCM1, test_weights, batch, batch_size, 1, plot=False)
+dcm1results = plot_model(DCM1, dcm1_weights, batch, batch_size, 1, plot=False)
+dcm2results = plot_model(DCM2, dcm2_weights, batch, batch_size, 1, plot=False)
 
-dipo = dcm1results["dipo"]
-mono = dcm1results["mono"]
 atoms = dcm1results["atoms"]
 dcmol = dcm1results["dcmol"]
+dcmol2 = dcm2results["dcmol"]
 
 
 
-
+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)
 
+output = StringIO()
+(atoms+dcmol2).write(output, format="html")
+data = output.getvalue()
 components.html(data, width=1000, height=1000)
 
+