hotfix greater than/less than operations in pdb_manager

a-r-j · Aug 15, 2024 · ca918bf · ca918bf
1 parent 598a6c4
commit ca918bf
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/graphein/ml/datasets/pdb_data.py b/graphein/ml/datasets/pdb_data.py
@@ -1085,7 +1085,10 @@ def oligomeric(
         update: bool = False,
     ) -> pd.DataFrame:
         """Select molecules with a given oligmeric length.
-        I.e. ``df.n_chains ==/</>  oligomer``
+        I.e. ``df.n_chains ==/ =< / >=  oligomer``
+
+        N.b. the `comparison` arguments for `"greater"` and `"less"` are
+        `>=` and `=<` respectively.
 
         :param length: Oligomeric length of molecule, defaults to ``1``.
         :type length: int
@@ -1106,9 +1109,9 @@ def oligomeric(
         if comparison == "equal":
             df = splits_df.loc[splits_df.n_chains == oligomer]
         elif comparison == "less":
-            df = splits_df.loc[splits_df.n_chains < oligomer]
+            df = splits_df.loc[splits_df.n_chains <= oligomer]
         elif comparison == "greater":
-            df = splits_df.loc[splits_df.n_chains > oligomer]
+            df = splits_df.loc[splits_df.n_chains >= oligomer]
         else:
             raise ValueError(
                 "Comparison must be one of 'equal', 'less', or 'greater'."