diff --git a/VOID/dockers/success.py b/VOID/dockers/success.py index e024eae..ee63162 100644 --- a/VOID/dockers/success.py +++ b/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 diff --git a/VOID/fitness/threshold.py b/VOID/fitness/threshold.py index 585cac6..11273d0 100644 --- a/VOID/fitness/threshold.py +++ b/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,12 +46,6 @@ 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, @@ -63,18 +53,6 @@ def add_arguments(parser): 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.