From c4e98225cb9c9761a9a995c949e147dbeefd6956 Mon Sep 17 00:00:00 2001
From: Muhammed Shuaibi <45150244+mshuaibii@users.noreply.github.com>
Date: Thu, 5 Sep 2024 09:53:49 -0500
Subject: [PATCH] add tags, fix min_ab slab (#829)

* add tags, fix min_ab slab

* add pp setups support
---
 src/fairchem/data/oc/core/interface_config.py |  1 +
 src/fairchem/data/oc/core/slab.py             |  2 +-
 src/fairchem/data/oc/utils/vasp.py            |  3 ++-
 src/fairchem/data/oc/utils/vasp_flags.py      | 25 ++++++-------------
 4 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/src/fairchem/data/oc/core/interface_config.py b/src/fairchem/data/oc/core/interface_config.py
index d34c55d64..91750cf01 100644
--- a/src/fairchem/data/oc/core/interface_config.py
+++ b/src/fairchem/data/oc/core/interface_config.py
@@ -183,6 +183,7 @@ def create_interface_on_sites(
                 "slab_shift": self.slab.shift,
                 "slab_top": self.slab.top,
                 "bulk_idx": self.slab.bulk.bulk_id_from_db,
+                "ase_tags": interface_atoms.get_tags(),
             }
             metadata_interface_list.append(metadata)
 
diff --git a/src/fairchem/data/oc/core/slab.py b/src/fairchem/data/oc/core/slab.py
index 3965477c1..f3f7685c7 100644
--- a/src/fairchem/data/oc/core/slab.py
+++ b/src/fairchem/data/oc/core/slab.py
@@ -93,7 +93,7 @@ def from_bulk_get_random_slab(
             )
             sampled_miller_idx = np.random.randint(len(avail_millers))
             slabs = Slab.from_bulk_get_specific_millers(
-                avail_millers[sampled_miller_idx], bulk
+                avail_millers[sampled_miller_idx], bulk, min_ab
             )
 
             # If multiple shifts exist, sample one randomly
diff --git a/src/fairchem/data/oc/utils/vasp.py b/src/fairchem/data/oc/utils/vasp.py
index c1189db24..c7f3f1d53 100644
--- a/src/fairchem/data/oc/utils/vasp.py
+++ b/src/fairchem/data/oc/utils/vasp.py
@@ -67,7 +67,7 @@ def calculate_surface_k_points(atoms):
     )
 
 
-def write_vasp_input_files(atoms, outdir=".", vasp_flags=None):
+def write_vasp_input_files(atoms, outdir=".", vasp_flags=None, pp_setups="minimal"):
     """
     Effectively goes through the same motions as the `run_vasp` function,
     except it only writes the input files instead of running.
@@ -84,4 +84,5 @@ def write_vasp_input_files(atoms, outdir=".", vasp_flags=None):
 
     atoms, vasp_flags = _clean_up_inputs(atoms, vasp_flags.copy())
     calc = Vasp(directory=outdir, **vasp_flags)
+    calc.input_params["setups"] = pp_setups
     calc.write_input(atoms)
diff --git a/src/fairchem/data/oc/utils/vasp_flags.py b/src/fairchem/data/oc/utils/vasp_flags.py
index fdedd1396..6a66ad28c 100644
--- a/src/fairchem/data/oc/utils/vasp_flags.py
+++ b/src/fairchem/data/oc/utils/vasp_flags.py
@@ -43,20 +43,19 @@
     "ivdw": 11,
     "encut": 400.0,
     "ediff": 1e-6,
-    "ediffg": -0.05,
     "nelm": 100,
     "ismear": 0,
     "sigma": 0.1,
-    "lwave": False,
-    "lcharg": False,
+    "lcharg": True,
+    "lwave": True,
     "isif": 0,
     "ispin": 2,
-    "algo": "Fast",
+    "algo": "All",
     "idipol": 3,
     "ldipol": True,
     "lasph": True,
     "lreal": "Auto",
-    "ncore": 100,
+    "ncore": 100,  # VASP will scale this down to whatever ncores are available.
     "dipol": [0.5, 0.5, 0.5],
 }
 
@@ -82,20 +81,12 @@
     "ncore": 4,
 }
 
-RELAX_FLAGS = {"ibrion": 2, "nsw": 500}
+RELAX_FLAGS = {"ibrion": 2, "nsw": 5}
 
 MD_FLAGS = {
     "ibrion": 0,
-    "nsw": 100,
+    "nsw": 10,
     "smass": 0,
-    "tebeg": 600,
-    "potim": 1,
-}
-
-ML_FLAGS = {
-    "ML_LMLFF": True,
-    "ML_ISTART": 0,
-    "ML_EPS_LOW": 1e-7,
-    "NSW": 1000,
-    "ISIF": 2,
+    "tebeg": 1000,
+    "potim": 2,
 }