From f1c1392f276a1a8639befa321185363f806f3f59 Mon Sep 17 00:00:00 2001
From: Pau Ferri-Vicedo <paufv@mit.edu>
Date: Wed, 29 May 2024 17:35:13 -0400
Subject: [PATCH] threshold for cation-anion distance added as an argument
 threshold_catan in def __init__ for ThresholdFitness class

---
 VOID/fitness/threshold.py | 29 ++++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/VOID/fitness/threshold.py b/VOID/fitness/threshold.py
index 9d33a00..f68b5eb 100644
--- a/VOID/fitness/threshold.py
+++ b/VOID/fitness/threshold.py
@@ -6,13 +6,21 @@
 
 
 THRESHOLD = 1.5
+THRESHOLD_CATAN = 2.0
 DEFAULT_STRUCTURE = "complex"
 STRUCTURE_CHOICES = ["complex", "guest", "host"]
 DEFAULT_STEP = False
 
 
 class ThresholdFitness(Fitness):
-    def __init__(self, threshold=THRESHOLD, structure="complex", step=False, **kwargs):
+    def __init__(
+        self,
+        threshold=THRESHOLD,
+        threshold_catan=THRESHOLD_CATAN,
+        structure="complex",
+        step=False,
+        **kwargs,
+    ):
         """Fitness is positive if the minimum distance is above
             the given threshold.
 
@@ -23,6 +31,7 @@ def __init__(self, threshold=THRESHOLD, structure="complex", step=False, **kwarg
         """
         super().__init__()
         self.threshold = threshold
+        self.threshold_catan = threshold_catan
         self.step = step
 
         if structure not in STRUCTURE_CHOICES:
@@ -39,6 +48,12 @@ def add_arguments(parser):
             help="threshold for distance calculations (default: %(default)s)",
             default=THRESHOLD,
         )
+        parser.add_argument(
+            "--threshold_catan",
+            type=float,
+            help="threshold for cation-anion distance calculations (default: %(default)s)",
+            default=THRESHOLD_CATAN,
+        )
         parser.add_argument(
             "--structure",
             type=str,
@@ -288,14 +303,14 @@ def get_catan_distances(
             distances_catan.append(distances_cation_anion)
 
         if (
-            any(dist < 2.0 for sublist in distances_catan for dist in sublist)
+            any(
+                dist < self.threshold_catan
+                for sublist in distances_catan
+                for dist in sublist
+            )
             and self.normalize(self.get_distances(complex).min() - self.threshold) > 0
         ):
-            print("Optimal distance found! Aborting the run")
-            print(
-                "Distances between cation and acid oxygens are:",
-                distances_catan,
-            )
+            print("Optimal cation-anion distance found! Aborting the run")
             return True, distances_catan
 
         else: