From f216c88cd68def4369b52da8820d54dcb7fde534 Mon Sep 17 00:00:00 2001
From: Pau Ferri-Vicedo <paufv@mit.edu>
Date: Tue, 28 May 2024 16:39:51 -0400
Subject: [PATCH] Initial functions for the MinDistanceCationAnionFitness class
 in the MC docking

---
 VOID/fitness/__init__.py  | 14 ++++++-
 VOID/fitness/threshold.py | 80 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 92 insertions(+), 2 deletions(-)

diff --git a/VOID/fitness/__init__.py b/VOID/fitness/__init__.py
index 29d18cf..b9f2645 100644
--- a/VOID/fitness/__init__.py
+++ b/VOID/fitness/__init__.py
@@ -1,10 +1,20 @@
 from .base import Fitness
-from .threshold import MinDistanceFitness, MeanDistanceFitness, SumInvDistanceFitness
-from .target import MinDistanceGaussianTarget, MeanDistanceGaussianTarget, MaxDistanceGaussianTarget
+from .threshold import (
+    MinDistanceFitness,
+    MeanDistanceFitness,
+    SumInvDistanceFitness,
+    MinDistanceCationAnionFitness,
+)
+from .target import (
+    MinDistanceGaussianTarget,
+    MeanDistanceGaussianTarget,
+    MaxDistanceGaussianTarget,
+)
 from .union import MultipleFitness
 
 __all__ = [
     MinDistanceFitness,
+    MinDistanceCationAnionFitness,
     MeanDistanceFitness,
     SumInvDistanceFitness,
     MinDistanceGaussianTarget,
diff --git a/VOID/fitness/threshold.py b/VOID/fitness/threshold.py
index 5d85f23..cecc4ba 100644
--- a/VOID/fitness/threshold.py
+++ b/VOID/fitness/threshold.py
@@ -2,6 +2,8 @@
 
 from .base import Fitness
 
+from pymatgen.core.sites import Site
+
 
 THRESHOLD = 1.5
 DEFAULT_STRUCTURE = "complex"
@@ -60,6 +62,62 @@ def get_distances(self, complex):
         else:
             raise ValueError("structure type not supported")
 
+    def get_zeolite_oxygens(self, pose):
+        """Collect all the O atoms in the structure."""
+        return [index for index, site in enumerate(pose) if site.species_string == "O"]
+
+    def find_cation_index(self, pose, distance_matrices):
+        """Identify the cation position in the guest."""
+        for index, site in enumerate(pose):
+            element = site.species_string
+            if element == "C":
+                bonds = sum(1 for dist in distance_matrices[index] if 0 < dist < 1.6)
+                if bonds == 3:
+                    return index
+        return None
+
+    def find_acid_sites(self, pose, distance_matrices, zeolite_oxygens):
+        """Identify the acid sites in the zeolite."""
+        acid_oxygens = []
+        acid_al_indexes = []
+
+        for index, site in enumerate(pose):
+            if site.species_string == "Al":
+                candidate_oxygens = [
+                    dist_index
+                    for dist_index, dist in enumerate(distance_matrices[index])
+                    if 0 < dist < 1.8 and dist_index in zeolite_oxygens
+                ]
+                if len(candidate_oxygens) == 4 and all(
+                    all(
+                        not (bond_dist < 1.15 and bond_dist != 0.0)
+                        for bond_dist in distance_matrices[ox_index]
+                    )
+                    for ox_index in candidate_oxygens
+                ):  # 1.15 accounts for O-H bond
+                    acid_oxygens.append(candidate_oxygens)
+                    acid_al_indexes.append(index)
+
+        return acid_oxygens, acid_al_indexes
+
+    def get_catan_distances(self, acid_oxygens, cation_index, distance_matrices):
+        """Check the cation-anion distances for the different acid sites in the zeolite."""
+        distances_catan = []
+        for acid_al in acid_oxygens:
+            distances_cation_anion = [
+                distance_matrices[cation_index][ox_index] for ox_index in acid_al
+            ]
+            print(
+                "Distances between cation and acid oxygens are:", distances_cation_anion
+            )
+            distances_catan.append(distances_cation_anion)
+            # if any(dist < 2.0 for dist in distances_cation_anion):
+            #    print("Optimal distance found! Aborting the run")
+            #    return True
+        # return False
+
+        return distances_catan
+
     def normalize(self, value):
         if self.step:
             return 0 if value > 0 else -np.inf
@@ -74,6 +132,28 @@ def __call__(self, complex):
         return self.normalize(self.get_distances(complex).min() - self.threshold)
 
 
+class MinDistanceCationAnionFitness(ThresholdFitness):
+    PARSER_NAME = "min_catan_distance"
+    HELP = "Complexes have positive score if the minimum distance between host anion and guest cation is above the given threshold"
+
+    def __call__(self, complex):
+        # print(complex.pose)
+        pose = complex.pose
+        distance_matrices = complex.pose.distance_matrix
+        zeolite_oxygens = self.get_zeolite_oxygens(pose)
+        cation_index = self.find_cation_index(
+            pose,
+            distance_matrices,
+        )
+        acid_sites, acid_al_indexes = self.find_acid_sites(
+            pose, distance_matrices, zeolite_oxygens
+        )
+        return self.normalize(
+            min(self.get_catan_distances(acid_sites, cation_index, distance_matrices))
+            - self.threshold
+        )
+
+
 class MeanDistanceFitness(ThresholdFitness):
     PARSER_NAME = "mean_distance"
     HELP = "Complexes have positive score if the mean distance between host and guest is above the given threshold"