add parallelization on get_populated_site

autoadsorbate/autoadsorbate.py

@@ -410,6 +410,7 @@ def get_populated_sites(
         conformers_per_site_cap=None,
         overlap_thr=1.5,
         verbose=False,
+        parallel=None,
     ):
         """
         Populates the specified sites with the given fragment, optimizing the orientation to minimize overlap.
@@ -419,9 +420,10 @@ def get_populated_sites(
         site_index (str or int): The index of the site to be populated. Default is 'all'.
         sample_rotation (bool): Whether to sample different rotations of the fragment. Default is True.
         mode (str): The mode of operation. Can be 'heuristic' or 'all'. Default is 'heuristic'.
-        conformers_per_site_cap (int or None): The maximum number of conformers per site. Default is None.
+        conformers_per_site_cap (int or None): The maximum number of conformers per site. Default None mean the maximum number of conformers.
         overlap_thr (float): The overlap threshold. Default is 1.5.
         verbose (bool): Whether to print detailed information during execution. Default is False.
+        parallel (int): If value >0, parallelize the configuration on the number of CPU specified. Default None
 
         Returns:
         list: A list containing the optimized atoms objects for each site.
@@ -477,13 +479,25 @@ def get_populated_sites(
             print("conformers", len(conformers))
             print("sites", len(sites))
 
+        if parallel != None:
+            from joblib import Parallel, delayed
+
         for site in sites:
             c_trj = []
-            for conformer in conformers:
-                c_trj += conformer_to_site(
-                    self.atoms, site, conformer, mode="optimize", overlap_thr=0
-                )  # the zero is intentional
 
+            if parallel != None:
+
+                c_trj.extend(Parallel(n_jobs=parallel)(delayed(conformer_to_site)(self.atoms, site, conformer, mode="optimize", overlap_thr=0) for conformer in conformers))
+                c_trj = [x[0] for x in c_trj]
+
+            else:
+                for conformer in conformers:
+                    c_trj += conformer_to_site(
+                        self.atoms, site, conformer, mode="optimize", overlap_thr=0
+                    )  # the zero is intentional
+
+            if conformers_per_site_cap == None:
+                conformers_per_site_cap = len(conformers)
             if conformers_per_site_cap != None:
                 c_trj = [atoms for atoms in c_trj if atoms.info["mdf"] > overlap_thr]
 

autoadsorbate/utils.py

@@ -395,23 +395,28 @@ def make_site_info_writable(site):
     return site
 
 
-def docs_plot_conformers(conformer_trajectory, rotation="-90x,0y,0z"):
+def docs_plot_conformers(conformer_trajectory, rotation="-90x,0y,0z", nb_column=5):
     """
     Helper function to plot a series of conformers.
 
     Parameters:
     conformer_trajectory (list): A list of atomic structures representing the conformers.
     rotation (str): The rotation to apply to the plot. Default is '-90x,0y,0z'.
+    nb_column (int): Control the number of column of the plot. The number of row is the number of conformers divided by number of column, rouded to superior.
 
     Returns:
     matplotlib.figure.Figure: The figure object containing the plots.
     """
-    fig, ax = plt.subplots(1, 5, figsize=(10, 2), sharex=True, sharey=True)
+    nb_row = int(np.ceil(len(conformer_trajectory)/nb_column))
+    fig, ax = plt.subplots(nb_row, nb_column, figsize=(2*nb_row, 2*nb_column), sharex=True, sharey=True)
+
+
+    for i in range(len(ax)):
+        for j in range(len(ax[i])):
+            plot_atoms(conformer_trajectory[i], ax=ax[i, j], rotation=rotation)
+            ax[i, j].set_xlim(0, 7), ax[i, j].set_ylim(0, 7)
+            ax[i, j].set_axis_off()
 
-    for i, atoms in enumerate(ax):
-        plot_atoms(conformer_trajectory[i], ax=ax[i], rotation=rotation)
-        ax[i].set_xlim(0, 7), ax[i].set_ylim(0, 7)
-        ax[i].set_axis_off()
     fig.suptitle("Generated structures viewed from +X axis", fontsize=12)
     fig.tight_layout()
     return fig