fix CI bugs

Author: zdcao121

main.py

@@ -229,48 +229,27 @@
 else:
     pass
 
-    print("\n========== Print out some test data for the given space group ==========")
+    print("\n========== Calculate the loss of test dataset ==========")
     import numpy as np 
     np.set_printoptions(threshold=np.inf)
 
-    G, L, XYZ, A, W = test_data
-    print (G.shape, L.shape, XYZ.shape, A.shape, W.shape)
-
-    idx = jnp.where(G==args.spacegroup,size=5)
-    G = G[idx]
-    L = L[idx]
-    XYZ = XYZ[idx]
-    A = A[idx]
-    W = W[idx]
-    
-    num_sites = jnp.sum(A!=0, axis=1)
-    print ("num_sites:", num_sites)
-    @jax.vmap
-    def lookup(G, W):
-        return mult_table[G-1, W] # (n_max, )
-    M = lookup(G, W) # (batchsize, n_max)
-    num_atoms = M.sum(axis=-1)
-    print ("num_atoms:", num_atoms)
-
-    print ("G:", G)
-    print ("A:\n", A)
-    for a in A: 
-       print([element_list[i] for i in a])
-    print ("W:\n",W)
-    print ("XYZ:\n",XYZ)
-
-    outputs = jax.vmap(transformer, (None, None, 0, 0, 0, 0, 0, None), (0))(params, key, G, XYZ, A, W, M, False)
-    print ("outputs.shape", outputs.shape)
-
-    h_al = outputs[:, 1::5, :] # (:, n_max, :)
-    a_logit = h_al[:, :, :args.atom_types]
-    l_logit, mu, sigma = jnp.split(h_al[jnp.arange(h_al.shape[0]), num_sites, 
-                                       args.atom_types:args.atom_types+args.Kl+2*6*args.Kl], 
-                                       [args.Kl, args.Kl+6*args.Kl], axis=-1)
-    print ("L:\n",L)
-    print ("exp(l_logit):\n", jnp.exp(l_logit))
-    print ("mu:\n", mu.reshape(-1, args.Kl, 6))
-    print ("sigma:\n", sigma.reshape(-1, args.Kl, 6))
+    test_G, test_L, test_XYZ, test_A, test_W = test_data
+    print (test_G.shape, test_L.shape, test_XYZ.shape, test_A.shape, test_W.shape)
+    test_loss = 0
+    num_samples = len(test_L)
+    num_batches = math.ceil(num_samples / args.batchsize)
+    for batch_idx in range(num_batches):
+        start_idx = batch_idx * args.batchsize
+        end_idx = min(start_idx + args.batchsize, num_samples)
+        G, L, XYZ, A, W = test_G[start_idx:end_idx], \
+                          test_L[start_idx:end_idx], \
+                          test_XYZ[start_idx:end_idx], \
+                          test_A[start_idx:end_idx], \
+                          test_W[start_idx:end_idx]
+        loss, _ = jax.jit(loss_fn, static_argnums=7)(params, key, G, L, XYZ, A, W, False)
+        test_loss += loss
+    test_loss = test_loss / num_batches
+    print ("evaluating loss on test data:" , test_loss)
 
     print("\n========== Start sampling ==========")
     jax.config.update("jax_enable_x64", True) # to get off compilation warning, and to prevent sample nan lattice