From b161c025606a7a20fdb584ca847bcd67e9060230 Mon Sep 17 00:00:00 2001
From: qzhu2017 <alecfans@gmail.com>
Date: Sun, 28 Jul 2024 18:52:40 -0400
Subject: [PATCH] add sample code for qmc

---
 pyxtal/__init__.py     | 10 ++++++
 pyxtal/lattice.py      | 19 ++++++++++
 pyxtal/optimize/qmc.py | 78 ++++++++++++++++++++++++++++++++++++++++++
 pyxtal/wyckoff_site.py | 32 +++++++++++++++--
 4 files changed, 137 insertions(+), 2 deletions(-)
 create mode 100644 pyxtal/optimize/qmc.py

diff --git a/pyxtal/__init__.py b/pyxtal/__init__.py
index 47b3ccf2..b5beed33 100644
--- a/pyxtal/__init__.py
+++ b/pyxtal/__init__.py
@@ -234,6 +234,16 @@ def get_dof(self):
 
         return self.lattice.dof + dof
 
+    def get_bounds(self, vec=(2.0, 50.0), ang=(30, 150)):
+        """
+        Get the number of dof for the given structures:
+        """
+        bounds = self.lattice.get_bounds(vec[0], vec[1], ang[0], ang[1])
+        sites = self.mol_sites if self.molecular else self.atom_sites
+        for site in sites:
+            bounds.extend(site.get_bounds())
+        return bounds
+
     def get_site_labels(self):
         """
         Get the site_labels as a dictionary
diff --git a/pyxtal/lattice.py b/pyxtal/lattice.py
index facb99a8..993fddab 100644
--- a/pyxtal/lattice.py
+++ b/pyxtal/lattice.py
@@ -163,6 +163,24 @@ def _get_dof(self):
         else:
             self.dof = 1
 
+    def get_bounds(self, min_vec=2.0, max_vec=50.0, min_ang=30, max_ang=150):
+        """
+        get the number of degree of freedom
+        """
+        if self.ltype in ["triclinic"]:
+            self.bounds = [(min_vec, max_vec), (min_vec, max_vec), (min_vec, max_vec),
+                           (min_ang, max_ang), (min_ang, max_ang), (min_ang, max_ang)]
+        elif self.ltype in ["monoclinic"]:
+            self.bounds = [(min_vec, max_vec), (min_vec, max_vec), (min_vec, max_vec),
+                           (min_ang, max_ang)]
+        elif self.ltype in ["orthorhombic"]:
+            self.bounds = [(min_vec, max_vec), (min_vec, max_vec), (min_vec, max_vec)]
+        elif self.ltype in ["tetragonal", "hexagonal", "trigonal"]:
+            self.bounds = [(min_vec, max_vec), (min_vec, max_vec)]
+        else:
+            self.bounds = [(min_vec, max_vec)]
+        return self.bounds
+
     @classmethod
     def get_dofs(self, ltype):
         """
@@ -196,6 +214,7 @@ def get_lengths(self):
         return mat, np.linalg.norm(mat, axis=1)
 
     def scale(self, factor=1.1):
+
         matrix = self.matrix
         return Lattice.from_matrix(matrix * factor, ltype=self.ltype)
 
diff --git a/pyxtal/optimize/qmc.py b/pyxtal/optimize/qmc.py
new file mode 100644
index 00000000..15196c78
--- /dev/null
+++ b/pyxtal/optimize/qmc.py
@@ -0,0 +1,78 @@
+"""
+global optimizer
+ACSALA
+81 11.38  6.48 11.24  96.9 1 0 0.23 0.43 0.03  -44.6   25.0   34.4  -76.6   -5.2  171.5 0 -70594.48
+"""
+# ACSALA has 13 variables (4 cell + 3 xyz + 3 ori + 3 torsions)
+import os
+from scipy.stats import qmc
+from pyxtal.representation import representation
+from pyxtal.optimize.common import optimizer
+from pyxtal.db import database
+from pyxtal.lattice import Lattice
+import pymatgen.analysis.structure_matcher as sm
+
+w_dir = "tmp"
+if not os.path.exists(w_dir):
+    os.makedirs(w_dir)
+
+db = database("pyxtal/database/test.db")
+code = "ACSALA"
+row = db.get_row(code)
+xtal0 = db.get_pyxtal(code)
+if xtal0.has_special_site():
+    xtal0 = xtal0.to_subgroup()
+smiles = row.mol_smi; print(smiles)
+pmg0 = xtal0.to_pymatgen()
+pmg0.remove_species("H")
+
+# Prepare prm
+c_info = row.data["charmm_info"]
+with open(w_dir + "/pyxtal.prm", "w") as prm:
+    prm.write(c_info["prm"])
+with open(w_dir + "/pyxtal.rtf", "w") as rtf:
+    rtf.write(c_info["rtf"])
+
+
+# Prepare the sampling.....
+bounds = xtal0.get_bounds([2.5, 25])
+lb = [b[0] for b in bounds]
+ub = [b[1] for b in bounds]
+sampler = qmc.Sobol(d=xtal0.get_dof(), scramble=False)
+sample = sampler.random_base2(m=10)
+sample = qmc.scale(sample, lb, ub)
+
+l_dof = xtal0.lattice.dof
+est_volume = sum(n*mol.volume for n, mol in zip(xtal0.numMols, xtal0.molecules))
+min_vol, max_vol = 0.25*est_volume, 5.0*est_volume
+
+for i, des in enumerate(sample):
+    cell = [xtal0.group.hall_number] + des[:l_dof].tolist()
+    lat = Lattice.from_1d_representation(cell[1:], xtal0.lattice.ltype)
+
+    if min_vol < lat.volume < max_vol:
+        #print('reset volume0', lat.encode(), lat.volume)
+        coef = (est_volume*0.5/lat.volume)**(1/3)
+        lat = lat.scale(coef)
+        cell = [xtal0.group.hall_number] + lat.encode()
+        #print('reset volume1', lat.encode(), coef, lat.volume)
+
+        x = [cell, [0] + des[l_dof:].tolist() + [False]]
+
+        rep0 = representation(x, [smiles])
+        xtal = rep0.to_pyxtal()
+        if len(xtal.check_short_distances(r=0.75)) == 0:
+            res = optimizer(xtal, c_info, w_dir, skip_ani=True)
+            if res is not None:
+                xtal, eng = res["xtal"], res["energy"]/sum(xtal.numMols)
+                if eng < 1e+3:
+                    rep1 = xtal.get_1D_representation()
+                    strs = f'Opt {i:4d} ' + rep1.to_string(eng) + f' Vol {lat.volume:7.1f} => {xtal.lattice.volume:7.1f}'
+                    pmg1 = xtal.to_pymatgen()
+                    pmg1.remove_species("H")
+                    if sm.StructureMatcher().fit(pmg0, pmg1):
+                        strs += '++++++'
+                    print(strs)
+
+    if i % 10000 == 0:
+        print("Progress", i, lat.volume, xtal0.lattice.volume)
diff --git a/pyxtal/wyckoff_site.py b/pyxtal/wyckoff_site.py
index e9c12e85..5e03fa62 100644
--- a/pyxtal/wyckoff_site.py
+++ b/pyxtal/wyckoff_site.py
@@ -88,6 +88,15 @@ def _get_dof(self):
 
         self.dof = self.wp.get_dof()
 
+    def get_bounds(self):
+        """
+        get the number of dof for the given structures:
+        """
+        self.bounds = []
+        for i in range(self.dof):
+            self.bounds.append([0.0, 1.0])
+        return self.bounds
+
     @classmethod
     def load_dict(cls, dicts):
         """
@@ -363,6 +372,7 @@ def __init__(self, mol, position, orientation, wp, lattice=None, stype=0):
         self.wp = wp
         self.position = position  # fractional coordinate of molecular center
         self.orientation = orientation  # pyxtal.molecule.orientation object
+        self._get_dof()
         if isinstance(lattice, Lattice):
             self.lattice = lattice
         else:
@@ -409,11 +419,29 @@ def __repr__(self):
         return str(self)
 
     def _get_dof(self):
+        """
+        get the number of dof for the given wyckoff site:
+        """
+        dof = np.linalg.matrix_rank(self.wp.ops[0].rotation_matrix)
+        self.dof = dof + 3
+        if self.molecule.torsionlist is not None:
+            self.dof += len(self.molecule.torsionlist)
+
+    def get_bounds(self):
         """
         get the number of dof for the given structures:
         """
-        freedom = np.trace(self.wp.ops[0].rotation_matrix) > 0
-        self.dof = len(freedom[freedom is True])
+        self.bounds = []
+        dof = np.linalg.matrix_rank(self.wp.ops[0].rotation_matrix)
+        for i in range(dof):
+            self.bounds.append([0.0, 1.0])
+        self.bounds.append([-180.0, 180.0])
+        self.bounds.append([-90.0, 90.0])
+        self.bounds.append([-180.0, 180.0])
+        if self.molecule.torsionlist is not None:
+            for i in range(len(self.molecule.torsionlist)):
+                self.bounds.append([0, 180.0])
+        return self.bounds
 
     def save_dict(self):
         dict0 = {