Merge branch 'master' of https://github.com/qzhu2017/PyXtal

Author: qzhu2017

examples/example_09_QRS_conf.py

@@ -91,7 +91,7 @@ def get_visited_energies(qrs):
     return energies
 
 
-def plot_id_vs_energy(code, energies, match_ids=None, match_energies=None, out_dir="qrs_plots", time_cost_s=None, coverage=None, n_conformers=None):
+def plot_id_vs_energy(code, energies, match_ids=None, match_energies=None, out_dir="qrs_plots", time_cost_s=None, coverage=None, n_conformers=None, energy_unit="kcal/mol"):
     """Save a plot of visited-structure ID vs energy for one QRS run."""
     if not energies:
         print(f"No energies collected for {code}; skipping plot.")
@@ -101,7 +101,7 @@ def plot_id_vs_energy(code, energies, match_ids=None, match_energies=None, out_d
     ids = list(range(1, len(energies) + 1))
 
     fig, ax = plt.subplots(figsize=(7, 4.5))
-    ax.scatter(ids, energies, s=20, alpha=0.8, label="Visited")
+    ax.scatter(ids, energies, s=10, alpha=0.8, label="Visited")
     ax.plot(ids, energies, linewidth=0.8, alpha=0.6)
     if match_ids and match_energies:
         ax.scatter(
@@ -116,12 +116,12 @@ def plot_id_vs_energy(code, energies, match_ids=None, match_energies=None, out_d
             label="Match",
         )
     ax.set_xlabel("Visited Structure ID")
-    ax.set_ylabel("Energy (kcal/mol)")
+    ax.set_ylabel(f"Energy ({energy_unit})")
     title_parts = [f"{code}: "]
     if n_conformers is not None:
         title_parts.append(f"conf: {n_conformers}")
     if time_cost_s is not None:
-        title_parts.append(f"time: {time_cost_s:.1f} s")
+        title_parts.append(f"time: {time_cost_s:.2f} s")
     if coverage is not None:
         title_parts.append(f"coverage: {coverage}")
     if len(title_parts) > 1:
@@ -138,7 +138,7 @@ def plot_id_vs_energy(code, energies, match_ids=None, match_energies=None, out_d
         ax.set_ylim(ymin - margin, ymax + margin)
     else:
         y_max = min(ymin + 30, ymax)
-        ax.set_ylim(ymin - 1, y_max)
+        ax.set_ylim(ymin - 0.25, y_max)
     fig.tight_layout()
 
     fig_path = os.path.join(out_dir, f"{code}.png")
@@ -173,8 +173,7 @@ def plot_id_vs_energy(code, energies, match_ids=None, match_energies=None, out_d
     for code in db.get_all_codes():
         #if code not in ['ACSALA']: continue
         #if code not in ['FUNZOE']: continue
-        if code in ['ACEMID02']: continue
-        #if code not in ['XAFPAY', 'OBEQIX', 'UJIRIO02']: continue
+        if code not in ['XAFQON']: continue
         row = db.get_row(code=code)
         ref_xtal = db.get_pyxtal(code=code)
         if ref_xtal.has_special_site():
@@ -288,11 +287,12 @@ def plot_id_vs_energy(code, energies, match_ids=None, match_energies=None, out_d
             composition = [int(a) for a in ref_xtal.get_zprime()],
             molecules=molecules,
             sites=sites,
-            N_gen=200,
-            N_pop=96,
-            N_cpu=48,
+            N_gen=100,
+            N_pop=48,
+            N_cpu=2,#4,
             cif="all.cif",
-            skip_mlp=True,
+            skip_mlp=False,#True,
+            mlp='MACEOFF',
             verbose=False,
             delta_length=1.0,
             delta_angle=15.0,
@@ -322,6 +322,7 @@ def plot_id_vs_energy(code, energies, match_ids=None, match_energies=None, out_d
             time_cost_s=time_cost_s,
             coverage=coverage,
             n_conformers=n_pregen_total,
+            energy_unit="eV/atom" if not qrs.skip_mlp else "kcal/mol",
         )
 
         with open(csv_path, "a", newline="") as fcsv:

pyxtal/interface/ase_opt.py

@@ -1,10 +1,18 @@
-# --- Optional fix for PyTorch weights_only change (harmless if unused) ---
-from torch.serialization import add_safe_globals
-add_safe_globals([slice])
-# ------------------------------------------------------------------------
+# --- Temporary fix for PyTorch 2.6 weights_only change ---
+import torch as _torch
+import functools as _functools
+_original_torch_load = _torch.load
+@_functools.wraps(_original_torch_load)
+def _patched_torch_load(*args, **kwargs):
+    kwargs.setdefault('weights_only', False)
+    return _original_torch_load(*args, **kwargs)
+_torch.load = _patched_torch_load
+# ----------------------------------------------------------
 
 import signal
 import numpy as np
+import os
+import contextlib
 from ase.constraints import FixSymmetry
 from ase.filters import UnitCellFilter
 from ase.optimize.fire import FIRE
@@ -15,6 +23,13 @@
 _cached_uma = None
 _cached_ani = None
 
+@contextlib.contextmanager
+def _suppress_stdout_stderr():
+    """Redirect stdout and stderr to /dev/null."""
+    with open(os.devnull, 'w') as devnull:
+        with contextlib.redirect_stdout(devnull), contextlib.redirect_stderr(devnull):
+            yield
+
 def get_calculator(calculator):
     """
     Return an ASE calculator instance.
@@ -28,28 +43,32 @@ def get_calculator(calculator):
     if isinstance(calculator, str):
         if calculator == "UMA":
             if _cached_uma is None:
-                from fairchem.core import pretrained_mlip, FAIRChemCalculator
-                predictor = pretrained_mlip.get_predict_unit("uma-s-1p1")
-                _cached_uma = FAIRChemCalculator(predictor,
-                                                 task_name="omc")
+                with _suppress_stdout_stderr():
+                    from fairchem.core import pretrained_mlip, FAIRChemCalculator
+                    predictor = pretrained_mlip.get_predict_unit("uma-s-1p1")
+                    _cached_uma = FAIRChemCalculator(predictor,
+                                                     task_name="omc")
             calc = _cached_uma
 
         elif calculator == "ANI":
             if _cached_ani is None:
-                import torchani
-                _cached_ani = torchani.models.ANI2x().ase()
+                with _suppress_stdout_stderr():
+                    import torchani
+                    _cached_ani = torchani.models.ANI2x().ase()
             calc = _cached_ani
 
         elif calculator == "MACE":
             if _cached_mace is None:
-                from mace.calculators import mace_mp
-                _cached_mace = mace_mp(model="small", dispersion=True)
+                with _suppress_stdout_stderr():
+                    from mace.calculators import mace_mp
+                    _cached_mace = mace_mp(model="small", dispersion=True)
             calc = _cached_mace
 
         elif calculator == "MACEOFF":
             if _cached_mace is None:
-                from mace.calculators import mace_off
-                _cached_mace = mace_off(model="medium")#, device="cpu")
+                with _suppress_stdout_stderr():
+                    from mace.calculators import mace_off
+                    _cached_mace = mace_off(model="medium")#, device="cpu")
             calc = _cached_mace
 
         else:

pyxtal/optimize/base.py

@@ -991,38 +991,33 @@ def local_optimization_mproc(self, xtals, ncpu, ids=None, qrs=False, pool=None):
         if ids is None:
             ids = range(len(xtals))
 
-        N_cycle = int(np.ceil(len(xtals) / ncpu))
-        # Generator to create arg_lists for multiprocessing tasks
+        # Interleaved assignment: worker i handles indices i, i+ncpu, i+2*ncpu, …
+        # so all workers stay busy and results can be sorted to 0,1,2,3… order.
         def generate_args_lists():
             for i in range(ncpu):
-                id1 = i * N_cycle
-                id2 = min([id1 + N_cycle, len(xtals)])
-                _ids = ids[id1: id2]
+                _indices = list(range(i, len(xtals), ncpu))
+                _ids = [ids[j] for j in _indices]
                 job_tags = [self.tag + "-g" + str(gen)
                             + "-p" + str(id) for id in _ids]
-                _xtals = [xtals[id][0] for id in range(id1, id2)]
+                _xtals = [xtals[j][0] for j in _indices]
                 mutates = []
                 labels = []
-                for i in range(id1, id2):
-                    orig_tag = xtals[i][1]
+                for j in _indices:
+                    orig_tag = xtals[j][1]
                     if qrs:
                         mutates.append(False)
                         labels.append(orig_tag if orig_tag != "Random" else None)
                     else:
-                        if orig_tag == "Mutation":
-                            mutates.append(True)
-                        else:
-                            mutates.append(False)
+                        mutates.append(orig_tag == "Mutation")
                         labels.append(None)
                 my_args = [_xtals, _ids, mutates, job_tags, labels, *args, self.rank, self.timeout]
-                yield tuple(my_args)  # Yield args instead of appending to a list
+                yield tuple(my_args)
 
         gen_results = []
         for result in pool.imap_unordered(process_task, generate_args_lists()):
             if result is not None:
                 for _res in result:
                     gen_results.append(_res)
-
         return gen_results
 
     def gen_summary(self, t0, gen_results, xtals):

pyxtal/optimize/common.py

@@ -424,8 +424,8 @@ def optimizer(
             and 0.25 < struc.get_density() < 3.0
         ):
             s = struc.to_ase()
-            step = 50 if mlp in ['MACE', 'ANI'] else 10
-            s = ASE_relax(s, mlp, step=step, fmax=0.1, logfile="ase.log")
+            step = 50 if mlp in ['MACE', 'ANI'] else 25
+            s = ASE_relax(s, mlp, opt_lat=opt_lat, step=step, fmax=0.1, logfile="ase.log")
             if s is None: return None
             eng = s.get_potential_energy()
             stress = max(abs(s.get_stress())) / units.GPa
@@ -662,7 +662,7 @@ def optimizer_single(
             strs += f" {match:.3f}"
 
         xtal.energy = eng
-        print(f"{id:3d} " + strs)#; import sys; sys.exit()
+        print(f"{id:3d} " + strs)
         return xtal, match, stable
     else:
         return None, match, stable
@@ -677,9 +677,9 @@ def refine_struc(xtal, smiles, calculator, mlp):
         calculator: ANI_relax or MACE_relax
     """
     s = xtal.to_ase()
-    s = calculator(s, mlp, step=50, fmax=0.1, logfile="ase.log")
-    s = calculator(s, mlp, step=250, opt_cell=True, logfile="ase.log")
-    s = calculator(s, mlp, step=50, fmax=0.1, logfile="ase.log")
+    s = calculator(s, mlp, opt_lat=False, step=50, fmax=0.1, logfile="ase.log")
+    s = calculator(s, mlp, opt_lat=True, step=250, logfile="ase.log")
+    s = calculator(s, mlp, opt_lat=False, step=50, fmax=0.1, logfile="ase.log")
     eng1 = s.get_potential_energy()  # /sum(xtal.numMols)
 
     xtal = pyxtal(molecular=True)