Class inheritance and argument parsing distributed between ThresholdF…

…itness and MinCationAnionFitness so --threshold_catan, --cation_index and --acid_sites will be ignored if any other kind of Fitness is called

VOID/dockers/success.py

@@ -1,5 +1,8 @@
 import numpy as np
 
+from pymatgen.core.sites import Site
+from pymatgen.core import Lattice, Structure
+
 from VOID.structure import Complex
 from VOID.dockers.serial import SerialDocker
 from VOID.dockers.mcdocker import MonteCarloDocker
@@ -45,3 +48,57 @@ def dock(self, attempts):
                 return [cpx]
 
         return []
+
+    def rescale(complex):
+        """Rescale the complex to the 0-1 range so results can be visualized in direct and xyz format.
+
+        Args:
+            complex (Complex): The host-guest complex object containing the host and guest molecules.
+
+        Returns:
+            Complex: The rescaled host-guest complex object.
+        """
+        lattice = complex.pose.lattice
+        frac_coords = []
+        species_list = []
+        site_labels = []
+
+        for site in complex.pose:
+            site_labels.append(site.label)
+            species_list.append(site.species)
+            coords = (
+                site.frac_coords
+                if site.label == "host"
+                else np.mod(site.frac_coords, 1.0)
+            )
+            frac_coords.append(coords)
+
+        site_properties = {"label": site_labels}
+
+        updated_structure = Structure(
+            lattice, species_list, frac_coords, site_properties=site_properties
+        )
+
+        num_host_atoms = len(complex.host)
+
+        species = updated_structure.species
+        cart_coords = updated_structure.cart_coords
+
+        # Update the host and guest with the rescaled 0-1 species and coordinates
+        complex.host = Structure(
+            lattice=complex.host.lattice,
+            species=species[:num_host_atoms],
+            coords=cart_coords[:num_host_atoms],
+            coords_are_cartesian=True,
+            site_properties=complex.host.site_properties,
+        )
+
+        complex.guest = Structure(
+            lattice=complex.host.lattice,
+            species=species[num_host_atoms:],
+            coords=cart_coords[num_host_atoms:],
+            coords_are_cartesian=True,
+            site_properties=complex.guest.site_properties,
+        )
+
+        return complex

VOID/fitness/threshold.py

@@ -16,9 +16,6 @@ class ThresholdFitness(Fitness):
     def __init__(
         self,
         threshold=THRESHOLD,
-        threshold_catan=THRESHOLD_CATAN,
-        cation_index=CATION_INDEX,
-        acid_sites=ACID_SITES,
         structure="complex",
         step=False,
         **kwargs,
@@ -33,7 +30,6 @@ def __init__(
         """
         super().__init__()
         self.threshold = threshold
-        self.threshold_catan = threshold_catan
         self.step = step
 
         if structure not in STRUCTURE_CHOICES:
@@ -50,31 +46,13 @@ 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,
             choices=STRUCTURE_CHOICES,
             help="threshold for distance calculations (default: %(default)s)",
             default=DEFAULT_STRUCTURE,
         )
-        parser.add_argument(
-            "--cation_index",
-            type=int,
-            help="index for the atom holding the positive charge in the molecule (default: %(default)s)",
-            default=CATION_INDEX,
-        )
-        parser.add_argument(
-            "--acid_sites",
-            type=list,
-            help="list of indexes for the O atoms that hold a negative charge (default: %(default)s)",
-            default=ACID_SITES,
-        )
 
     def get_distances(self, complex):
         if self.structure == "complex":
@@ -130,21 +108,39 @@ class MinDistanceCationAnionFitness(ThresholdFitness):
 
     def __init__(
         self,
-        threshold,
-        threshold_catan,
-        cation_index,
-        acid_sites,
-        structure,
-        step,
-        **kwargs,
+        threshold=THRESHOLD,
+        threshold_catan=THRESHOLD_CATAN,
+        structure=DEFAULT_STRUCTURE,
+        cation_index=None,
+        acid_sites=None,
     ):
-        super().__init__(**kwargs)
-        self.threshold = threshold
+        super().__init__(threshold)
         self.threshold_catan = threshold_catan
         self.cation_index = cation_index
         self.acid_sites = acid_sites
-        self.structure = structure
-        self.step = step
+
+    @staticmethod
+    def add_arguments(parser):
+        ThresholdFitness.add_arguments(parser)
+
+        parser.add_argument(
+            "--threshold_catan",
+            type=float,
+            help="threshold for cation-anion distance calculations (default: %(default)s)",
+            default=THRESHOLD_CATAN,
+        )
+        parser.add_argument(
+            "--cation_index",
+            type=int,
+            help="index for the atom holding the positive charge in the molecule (default: %(default)s)",
+            default=CATION_INDEX,
+        )
+        parser.add_argument(
+            "--acid_sites",
+            type=list,
+            help="list of indexes for the O atoms that hold a negative charge (default: %(default)s)",
+            default=ACID_SITES,
+        )
 
     def __call__(self, complex):
         """Docks a guest cation into a host with anionic spots while ensuring a minimal distance between them.