Merge pull request #368 from carlocamilloni/post_symmetrization

fix sym for check
multi-ego · Mar 6, 2024 · b39736a · b39736a
2 parents 6347144 + 69a9a0f
commit b39736a
Showing 1 changed file with 12 additions and 4 deletions.
diff --git a/src/multiego/ensemble.py b/src/multiego/ensemble.py
@@ -1281,17 +1281,25 @@ def generate_LJ(meGO_ensemble, train_dataset, check_dataset, parameters):
         # apply symmetries to check contacts
         for sym in symmetries:
             for atypes in itertools.combinations(sym[1:], 2):
-                stmp_df_ai_L = meGO_LJ[meGO_LJ["ai"].str.startswith(f"{atypes[0]}_") & (mglj_resn_ai == sym[0])].copy()
-                stmp_df_aj_L = meGO_LJ[meGO_LJ["aj"].str.startswith(f"{atypes[0]}_") & (mglj_resn_aj == sym[0])].copy()
+                stmp_df_ai_L = meGO_check_contacts[
+                    meGO_check_contacts["ai"].str.startswith(f"{atypes[0]}_") & (mglj_resn_ai == sym[0])
+                ].copy()
+                stmp_df_aj_L = meGO_check_contacts[
+                    meGO_check_contacts["aj"].str.startswith(f"{atypes[0]}_") & (mglj_resn_aj == sym[0])
+                ].copy()
                 stmp_df_ai_L.loc[:, "ai"] = (
                     atypes[1] + "_" + stmp_df_ai_L["ai"].str.split("_").str[1] + "_" + stmp_df_ai_L["ai"].str.split("_").str[2]
                 )
                 stmp_df_aj_L.loc[:, "aj"] = (
                     atypes[1] + "_" + stmp_df_aj_L["aj"].str.split("_").str[1] + "_" + stmp_df_aj_L["aj"].str.split("_").str[2]
                 )
 
-                stmp_df_ai_R = meGO_LJ[meGO_LJ["ai"].str.startswith(f"{atypes[1]}_") & (mglj_resn_ai == sym[0])].copy()
-                stmp_df_aj_R = meGO_LJ[meGO_LJ["aj"].str.startswith(f"{atypes[1]}_") & (mglj_resn_aj == sym[0])].copy()
+                stmp_df_ai_R = meGO_check_contacts[
+                    meGO_check_contacts["ai"].str.startswith(f"{atypes[1]}_") & (mglj_resn_ai == sym[0])
+                ].copy()
+                stmp_df_aj_R = meGO_check_contacts[
+                    meGO_check_contacts["aj"].str.startswith(f"{atypes[1]}_") & (mglj_resn_aj == sym[0])
+                ].copy()
                 stmp_df_ai_R.loc[:, "ai"] = (
                     atypes[0] + "_" + stmp_df_ai_R["ai"].str.split("_").str[1] + "_" + stmp_df_ai_R["ai"].str.split("_").str[2]
                 )