performance improvements

src/multiego/ensemble.py

@@ -333,29 +333,25 @@ def init_meGO_matrices(ensemble, args, custom_dict):
         lj_data = topology.get_lj_params(topol)
         # these are the combined cases (c6_ij, c12_ij)
         lj_pairs = topology.get_lj_pairs(topol)
+        # Create reversed pairs
+        reversed_lj_pairs = lj_pairs.rename(columns={"ai": "aj", "aj": "ai"})
+        # Combine original and reversed
+        symmetric_lj_pairs = pd.concat([lj_pairs, reversed_lj_pairs])
+        # Remove duplicates to avoid duplication of symmetric pairs
+        # This step ensures that if a pair already exists in both directions, it's not duplicated
+        symmetric_lj_pairs = symmetric_lj_pairs.drop_duplicates(subset=["ai", "aj"]).reset_index(drop=True)
+
         # these are the combined cases in the [pairs] section (c6_ij, c12_ij)
         lj14_pairs = topology.get_lj14_pairs(topol)
+        # Create reversed pairs
+        reversed_lj14_pairs = lj14_pairs.rename(columns={"ai": "aj", "aj": "ai"})
+        # Combine original and reversed
+        symmetric_lj14_pairs = pd.concat([lj14_pairs, reversed_lj14_pairs])
+        # Remove duplicates to avoid duplication of symmetric pairs
+        # This step ensures that if a pair already exists in both directions, it's not duplicated
+        symmetric_lj14_pairs = symmetric_lj14_pairs.drop_duplicates(subset=["ai", "aj"]).reset_index(drop=True)
+
         lj_data_dict = {str(key): val for key, val in zip(lj_data["ai"], lj_data[["c6", "c12"]].values)}
-        #lj_pairs_dict = {
-        #    (ai, aj): (epsilon, sigma) for ai, aj, epsilon, sigma in lj_pairs[["ai", "aj", "epsilon", "sigma"]].to_numpy()
-        #}
-
-        # Create the lj_pairs_dict and make it symmetric
-        lj_pairs_dict = {}
-        for ai, aj, epsilon, sigma in lj_pairs[["ai", "aj", "epsilon", "sigma"]].to_numpy():
-        # Add both (ai, aj) and (aj, ai) with the same epsilon, sigma
-            lj_pairs_dict[(ai, aj)] = (epsilon, sigma)
-            lj_pairs_dict[(aj, ai)] = (epsilon, sigma)  # Ensure symmetry
-        lj14_pairs_dict = {}
-        for ai, aj, epsilon, sigma in lj14_pairs[["ai", "aj", "epsilon", "sigma"]].to_numpy():
-        # Add both (ai, aj) and (aj, ai) with the same epsilon, sigma
-            lj14_pairs_dict[(ai, aj)] = (epsilon, sigma)
-            lj14_pairs_dict[(aj, ai)] = (epsilon, sigma)  # Ensure symmetry
-
-
-        #lj14_pairs_dict = {
-        #    (ai, aj): (epsilon, sigma) for ai, aj, epsilon, sigma in lj14_pairs[["ai", "aj", "epsilon", "sigma"]].to_numpy()
-        #}
 
         ensemble["topology_dataframe"]["c6"] = lj_data["c6"].to_numpy()
         ensemble["topology_dataframe"]["c12"] = lj_data["c12"].to_numpy()
@@ -398,15 +394,19 @@ def init_meGO_matrices(ensemble, args, custom_dict):
                 0,
             )
 
-            # apply the LJ pairs
-            # apply the LJ14 pairs (only if same_chain = True)
-            for i, row in reference_contact_matrices[name].iterrows():
-                if (row["rc_ai"], row["rc_aj"]) in lj_pairs_dict.keys():
-                    reference_contact_matrices[name].loc[i, "epsilon_prior"] = lj_pairs_dict[(row["rc_ai"], row["rc_aj"])][0]
-                    reference_contact_matrices[name].loc[i, "sigma_prior"] = lj_pairs_dict[(row["rc_ai"], row["rc_aj"])][1]
-                if row["rc_same_chain"] and (row["rc_ai"], row["rc_aj"]) in lj14_pairs_dict.keys():
-                    reference_contact_matrices[name].loc[i, "epsilon_prior"] = lj14_pairs_dict[(row["rc_ai"], row["rc_aj"])][0]
-                    reference_contact_matrices[name].loc[i, "sigma_prior"] = lj14_pairs_dict[(row["rc_ai"], row["rc_aj"])][1]
+            # Create a mapping from lj_pairs for sigma and epsilon
+            lj_sigma_map = symmetric_lj_pairs.set_index(["ai", "aj"])["sigma"]
+            lj_epsilon_map = symmetric_lj_pairs.set_index(["ai", "aj"])["epsilon"]
+            lj14_sigma_map = symmetric_lj14_pairs.set_index(["ai", "aj"])["sigma"]
+            lj14_epsilon_map = symmetric_lj14_pairs.set_index(["ai", "aj"])["epsilon"]
+
+            # Update sigma values where they exist in lj_pairs
+            reference_contact_matrices[name].loc[lj_sigma_map.index, "sigma_prior"] = lj_sigma_map.astype("float64")
+            reference_contact_matrices[name].loc[lj14_sigma_map.index, "sigma_prior"] = lj14_sigma_map.astype("float64")
+
+            # Update epsilon values where they exist in lj_pairs
+            reference_contact_matrices[name].loc[lj_epsilon_map.index, "epsilon_prior"] = lj_epsilon_map.astype("float64")
+            reference_contact_matrices[name].loc[lj14_epsilon_map.index, "epsilon_prior"] = lj14_epsilon_map.astype("float64")
 
             reference_contact_matrices[name].drop(columns=["c6_i", "c6_j", "c12_i", "c12_j", "c6", "c12"], inplace=True)