From 63c26c9f926a18124cad1f9cee5b0df3ff57591f Mon Sep 17 00:00:00 2001
From: Matteo Giantomassi <gmatteo@gmail.com>
Date: Sun, 18 Feb 2024 22:19:00 +0100
Subject: [PATCH] Multiprocessing in abiml.py validate

---
 abipy/ml/aseml.py      | 38 ++++++++++++++++++++++++++++++++------
 abipy/scripts/abiml.py | 25 +++++++++++++------------
 2 files changed, 45 insertions(+), 18 deletions(-)

diff --git a/abipy/ml/aseml.py b/abipy/ml/aseml.py
index 0370df9eb..d940218cf 100644
--- a/abipy/ml/aseml.py
+++ b/abipy/ml/aseml.py
@@ -49,6 +49,7 @@
 from abipy.tools.printing import print_dataframe
 from abipy.tools.serialization import HasPickleIO
 from abipy.tools.context_managers import Timer
+from abipy.tools.parallel import get_max_nprocs, pool_nprocs_pmode
 from abipy.abio.enums import StrEnum, EnumMixin
 from abipy.core.mixins import TextFile, NotebookWriter
 from abipy.tools.plotting import (set_axlims, add_fig_kwargs, get_ax_fig_plt, get_axarray_fig_plt, set_grid_legend,
@@ -2845,21 +2846,31 @@ def run(self, nprocs, print_dataframes=True) -> AseResultsComparator:
         labels = ["abinitio"]
         abi_results = self.get_abinitio_results()
         results_list = [abi_results]
+        ntasks = len(results_list)
 
-        ntasks = len(abi_results)
-        nprocs = 1
+        if nprocs <= 0 or nprocs is None:
+            nprocs = get_max_nprocs()
+        print(f"Using {nprocs=}")
+
+        #from abipy.relax_scanner import nprocs_for_ntasks
+        #nprocs = nprocs_for_ntasks(nprocs, ntasks, title="Begin relaxations")
+        #p = pool_nprocs_pmode(ntasks, pmode=pmode)
+        #using_msg = f"Reading {len(directories)} abiml directories {p.using_msg}"
 
         for nn_name in self.nn_names:
-            labels.append(nn_name)
             # Use ML to compute quantities with the same ab-initio trajectory.
+            labels.append(nn_name)
             if nprocs == 1:
                 calc = as_calculator(nn_name)
                 items = [AseResults.from_atoms(res.atoms, calc=calc) for res in abi_results]
             else:
-                raise NotImplementedError("run with multiprocessing!")
-                args_list = [(nn_name, res) for res in abi_results]
+                func = _GetAseResults(nn_name)
                 with Pool(processes=nprocs) as pool:
-                    items = pool.map(_map_run_compare, args_list)
+                    items = pool.map(func, abi_results)
+                #p = pool_nprocs_pmode(len(directories), pmode=pmode)
+                #using_msg = f"Reading {len(directories)} abiml directories {p.using_msg}"
+                #with p.pool_cls(p.nprocs) as pool, Timer(header=using_msg, footer="") as timer:
+                #    return cls(pool.starmap(MdAnalyzer.from_abiml_dir, args))
 
             results_list.append(items)
 
@@ -2877,6 +2888,21 @@ def run(self, nprocs, print_dataframes=True) -> AseResultsComparator:
         return comp
 
 
+class _GetAseResults:
+    """
+    Callable class used to parallelize the computation of AseResults with multiprocessing
+    """
+
+    def __init__(self, nn_name):
+        self.nn_name = nn_name
+        self.calc = None
+
+    def __call__(self, abi_res):
+        if self.calc is None:
+            self.calc = as_calculator(self.nn_name)
+        return AseResults.from_atoms(abi_res.atoms, calc=self.calc)
+
+
 class MolecularDynamics:
     """
     Molecular dynamics class
diff --git a/abipy/scripts/abiml.py b/abipy/scripts/abiml.py
index 188b198a9..d2e8602ec 100755
--- a/abipy/scripts/abiml.py
+++ b/abipy/scripts/abiml.py
@@ -591,7 +591,7 @@ def scan_relax(ctx, filepath, nn_name,
               help="Trajectory range e.g. `5` to select the first 5 iterations, `1:4` to select steps 1,2,3. `1:4:2 for 1,3",
               default=None)
 @click.option("-e", '--exposer', default="mpl", show_default=True, type=click.Choice(["mpl", "panel"]),
-              help='Plotting backend: mpl for matplotlib, panel for web-based')
+              help='Plotting backend: mpl for matplotlib, panel for web-based, None to disable plotting')
 @add_nprocs_opt
 @add_workdir_verbose_opts
 @click.option('--config', default='abiml_validate.yml', type=click.Path(), callback=set_default, is_eager=True, expose_value=False)
@@ -618,17 +618,18 @@ def validate(ctx, filepaths,
     print(ml_comp)
     c = ml_comp.run(nprocs=nprocs)
 
-    with_stress = True
-    from abipy.tools.plotting import Exposer
-    with Exposer.as_exposer(exposer, title=" ".join(os.path.basename(p) for p in filepaths)) as e:
-        e(c.plot_energies(show=False))
-        e(c.plot_forces(delta_mode=True, show=False))
-        e(c.plot_energies_traj(delta_mode=True, show=False))
-        e(c.plot_energies_traj(delta_mode=False, show=False))
-        if with_stress:
-            e(c.plot_stresses(delta_mode=True, show=False))
-        e(c.plot_forces_traj(delta_mode=True, show=False))
-        e(c.plot_stress_traj(delta_mode=True, show=False))
+    if exposer != "None":
+        from abipy.tools.plotting import Exposer
+        with_stress = True
+        with Exposer.as_exposer(exposer, title=" ".join(os.path.basename(p) for p in filepaths)) as e:
+            e(c.plot_energies(show=False))
+            e(c.plot_forces(delta_mode=True, show=False))
+            e(c.plot_energies_traj(delta_mode=True, show=False))
+            e(c.plot_energies_traj(delta_mode=False, show=False))
+            if with_stress:
+                e(c.plot_stresses(delta_mode=True, show=False))
+            e(c.plot_forces_traj(delta_mode=True, show=False))
+            e(c.plot_stress_traj(delta_mode=True, show=False))
 
     return 0