rsspolymlp 0.3.2__tar.gz → 0.3.4__tar.gz

{rsspolymlp-0.3.2/src/rsspolymlp.egg-info → rsspolymlp-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rsspolymlp
-Version: 0.3.2
+Version: 0.3.4
 Summary: A framework for random structure search using polynomial MLPs
 Author-email: Hayato Wakai <wakai@cms.mtl.kyoto-u.ac.jp>
 Maintainer-email: Hayato Wakai <wakai@cms.mtl.kyoto-u.ac.jp>

{rsspolymlp-0.3.2 → rsspolymlp-0.3.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "rsspolymlp"
-version = "0.3.2"
+version = "0.3.4"
 description = "A framework for random structure search using polynomial MLPs"
 license = "MIT"
 authors = [

{rsspolymlp-0.3.2 → rsspolymlp-0.3.4}/src/rsspolymlp/analysis/phase_analysis.py

@@ -1,3 +1,4 @@
+import ast
 import json
 import os
 from collections import defaultdict
@@ -18,6 +19,7 @@ class ConvexHullAnalyzer:
         self.elements = np.array(elements)
         # Initialize per-composition storage
         self._composition_data = {}
+        self.endmember_energies = None
         # Convex hull object
         self.hull = None
         # Results for global minima on the hull
@@ -33,9 +35,9 @@ class ConvexHullAnalyzer:
     def composition_data(self, data):
         """
         composition_data maps each composition ratio (tuple) to a dict with keys:
-            'energy': np.ndarray of energies,
-            'input_path': np.ndarray of file paths,
-            'struct_tag': np.ndarray of structure tags.
+            - 'energy': np.ndarray of energies,
+            - 'input_path': np.ndarray of file paths,
+            - 'struct_tag': np.ndarray of structure tags.
         """
         if not isinstance(data, dict):
             raise TypeError("composition_data must be a dictionary.")
@@ -60,53 +62,53 @@ class ConvexHullAnalyzer:
         self._composition_data = data
 
     def run_analysis(self):
-        """Compute formation energies, build convex hull, and evaluate energies above hull."""
+        """
+        Construct the convex hull, compute formation energies, and evaluate energies above the hull.
+        The convex-hull results are stored in the following attributes:
+            - 'self.hull_energies': formation energies of structures on the hull
+            - 'self.hull_compositions': compositions of structures on the hull
+            - 'self.hull_paths': paths corresponding to structures on the hull
+        self._composition_data maps each composition ratio to a dictionary containing:
+            - 'formation_e': np.ndarray of formation energies
+            - 'fe_above_ch': np.ndarray of energies above the convex hull
+        """
         if len(self.elements) == 0:
             raise ValueError("elements is none. Please set elements.")
         if not self._composition_data:
             raise ValueError("composition_data is empty. Please set results first.")
-        self.compute_formation_energies()
+
+        self.set_endmember_energies()
         self.compute_convex_hull()
+        self.compute_formation_energies()
         self.compute_fe_above_ch()
 
-    def compute_formation_energies(self):
-        """Compute formation energy for each structure relative to pure elements."""
-        # Sort compositions for reproducible output
-        comps = np.array(list(self._composition_data.keys()))
-        order = np.lexsort(comps.T)
-        sorted_keys = [tuple(comps[i].tolist()) for i in order]
-        self._composition_data = {k: self._composition_data[k] for k in sorted_keys}
-
+    def set_endmember_energies(self):
         # Reference energies from pure elements (where any ratio == 1)
+        os.makedirs("phase_analysis/data", exist_ok=True)
+        comps = np.array(list(self._composition_data.keys()))
         pure = comps[np.any(comps == 1, axis=1)]
         pure_sorted = sorted(pure, key=lambda x: np.argmax(x))
-        ref_energies = [
+        self.endmember_energies = [
             np.min(self._composition_data[tuple(r)]["energy"]) for r in pure_sorted
         ]
-        ref_energies = np.array(ref_energies)
-
-        # Calculate and store formation energies
-        for ratio, data in self._composition_data.items():
-            formation = data["energy"] - np.dot(ref_energies, np.array(ratio))
-            data["formation_e"] = formation
-
-        # Save JSON and elements
-        os.makedirs("phase_analysis/data", exist_ok=True)
-        serial = {
-            str(r): {k: v.tolist() for k, v in d.items()}
-            for r, d in self._composition_data.items()
-        }
-        with open("phase_analysis/data/composition_data.json", "w") as f:
-            json.dump(serial, f)
-        np.save("phase_analysis/data/elements.npy", self.elements)
+        self.endmember_energies = np.array(self.endmember_energies)
 
     def compute_convex_hull(self):
         """Build the convex hull and extract vertex compositions and energies."""
+        os.makedirs("phase_analysis/data", exist_ok=True)
+        # Sort compositions for reproducible output
+        comps = np.array(list(self._composition_data.keys()))
+        order = np.lexsort(comps.T)
+        sorted_keys = [tuple(comps[i].tolist()) for i in order]
+        self._composition_data = {k: self._composition_data[k] for k in sorted_keys}
+
         comps, min_e, paths = [], [], []
         for ratio, data in self._composition_data.items():
             comps.append(ratio)
-            idx = np.argmin(data["formation_e"])
-            min_e.append(data["formation_e"][idx])
+            idx = np.argmin(data["energy"])
+            min_e.append(
+                data["energy"][idx] - np.dot(self.endmember_energies, np.array(ratio))
+            )
             paths.append(data["input_path"][idx])
 
         comp_arr = np.array(comps)
@@ -140,10 +142,28 @@ class ConvexHullAnalyzer:
                 f.write(f"    formation_e: {energy}\n")
 
         # Save NumPy arrays
-        os.makedirs("phase_analysis/data", exist_ok=True)
         np.save("phase_analysis/data/hull_e.npy", self.hull_energies)
         np.save("phase_analysis/data/hull_comp.npy", self.hull_compositions)
 
+    def compute_formation_energies(self, json_output=True):
+        """Compute formation energy for each structure relative to pure elements."""
+        # Calculate and store formation energies
+        for ratio, data in self._composition_data.items():
+            formation = data["energy"] - np.dot(
+                self.endmember_energies, np.array(ratio)
+            )
+            data["formation_e"] = formation
+
+        # Save JSON and elements
+        if json_output:
+            serial = {
+                str(r): {k: v.tolist() for k, v in d.items()}
+                for r, d in self._composition_data.items()
+            }
+            with open("phase_analysis/data/composition_data.json", "w") as f:
+                json.dump(serial, f)
+            np.save("phase_analysis/data/elements.npy", self.elements)
+
     def compute_fe_above_ch(self):
         """Compute each structure's energy above the convex hull."""
         for ratio, data in self._composition_data.items():
@@ -309,6 +329,9 @@ class ConvexHullAnalyzer:
             for element in structure.elements:
                 energy_dft -= atomic_energy(element)
             energy_dft /= len(structure.elements)
+            if energy_dft < -15:
+                print(vasprun_path, "exhibits very low energy:", energy_dft, "eV/atom")
+                continue
 
             comp_res = compute_composition(structure.elements, self.elements)
             comp_ratio = tuple(
@@ -338,3 +361,35 @@ class ConvexHullAnalyzer:
             }
 
         self._composition_data = dft_dict_array
+
+
+def load_convexhull_data(path="./phase_analysis", threshold=None):
+    res = {}
+
+    res["elements"] = np.load(f"{path}/data/elements.npy")
+    res["hull_comp"] = np.load(f"{path}/data/hull_comp.npy")
+    res["hull_e"] = np.load(f"{path}/data/hull_e.npy")
+    with open(f"{path}/data/composition_data.json", "r") as f:
+        data = json.load(f)
+    res["composition_data"] = convert_json_to_ndarray(data)
+
+    res["not_near_ch"] = None
+    res["near_ch"] = None
+    if threshold is not None:
+        thre = float(threshold)
+        with open(f"{path}/threshold_{thre}meV/not_near_ch.json", "r") as f:
+            data1 = json.load(f)
+        with open(f"{path}/threshold_{thre}meV/near_ch.json", "r") as f:
+            data2 = json.load(f)
+        res["not_near_ch"] = convert_json_to_ndarray(data1)
+        res["near_ch"] = convert_json_to_ndarray(data2)
+
+    return res
+
+
+def convert_json_to_ndarray(data):
+    converted = {
+        ast.literal_eval(k): {key: np.array(val) for key, val in v.items()}
+        for k, v in data.items()
+    }
+    return converted

{rsspolymlp-0.3.2 → rsspolymlp-0.3.4}/src/rsspolymlp/analysis/rss_summarize.py

@@ -80,9 +80,7 @@ class RSSResultSummarizer:
             parent_path_ref[composition_tag] = json_path
 
         if not self.parse_vasp:
-            paths_same_comp, results_same_comp = self._parse_json_result(
-                parse_rss_result=True
-            )
+            paths_same_comp, results_same_comp = self._parse_json_result()
             axis_tol = 0.03
             pos_tol = 0.03
         else:
@@ -187,10 +185,12 @@ class RSSResultSummarizer:
         paths_same_comp, results_same_comp = self._parse_json_result()
 
         for composition_tag, res_paths in paths_same_comp.items():
+            print(f"Composition {composition_tag}: summarizing...")
             self.analyzer = UniqueStructureAnalyzer()
 
             time_start = time()
             for res_path in res_paths:
+                print(f" - Processing result file (parent): {res_path}")
                 self.remove_duplicates_in_comp(
                     results_same_comp[composition_tag][res_path],
                     standardize_axis=True,
@@ -364,7 +364,7 @@ class RSSResultSummarizer:
 
         return struct_counts
 
-    def _parse_json_result(self, parse_rss_result=False):
+    def _parse_json_result(self):
 
         def resolve_path(base: Path, p):
             cwd = Path.cwd()
@@ -382,15 +382,12 @@ class RSSResultSummarizer:
             with open(path_name) as f:
                 loaded_dict = json.load(f)
 
-            if parse_rss_result:
-                base = Path(path_name).parents[1]
-                for r in loaded_dict["rss_results"]:
-                    r["struct_path"] = resolve_path(base, r["struct_path"])
-                    r["dupstr_paths"] = {
-                        resolve_path(base, p) for p in r["dupstr_paths"]
-                    }
-                    r["structure"] = polymlp_struct_from_dict(r["structure"])
-                    r["struct_no"] = None
+            base = Path(path_name).parents[1]
+            for r in loaded_dict["rss_results"]:
+                r["struct_path"] = resolve_path(base, r["struct_path"])
+                r["dupstr_paths"] = {resolve_path(base, p) for p in r["dupstr_paths"]}
+                r["structure"] = polymlp_struct_from_dict(r["structure"])
+                r["struct_no"] = None
 
             target_elements = loaded_dict["elements"]
             comp_ratio = tuple(loaded_dict["comp_ratio"])

{rsspolymlp-0.3.2 → rsspolymlp-0.3.4}/src/rsspolymlp/analysis/struct_matcher/struct_match.py

@@ -63,6 +63,7 @@ def struct_match(
     struct_match = False
     axis_d_min = None
     pos_d_min = None
+    min_axis_l2_norm = 1e8
     for st_1 in st_1_set:
         for st_2 in st_2_set:
             if struct_match or st_1.element_count != st_2.element_count:
@@ -75,6 +76,7 @@ def struct_match(
             )  # (N_symp1, N_symp2, 6)
             axis_d_flat = axis_d.reshape(-1, axis_d.shape[2])  # (N_symp1*N_symp2, 6)
             l2_norm = np.linalg.norm(axis_d_flat, axis=1)
+            min_axis_l2_norm = min(min_axis_l2_norm, np.min(l2_norm))
 
             match_axis = l2_norm < axis_tol
             if not np.any(match_axis):
@@ -117,32 +119,31 @@ def struct_match(
                 ["a", "b", "c"], np.round(positions.reshape(3, -1), 3).tolist()
             ):
                 formatted = ",".join(f"{val:6.3f}" for val in p)
-                print(f"       - [{formatted}]")
+                print(f"      - [{formatted}]")
 
+        print("--- results of rsspolymlp.analysis.struct_matcher ---")
         print("tolerance:")
         print("  axis_tol:", axis_tol)
         print("  pos_tol:", pos_tol)
-        print("")
         print("structures:")
         for i, st_set in enumerate([st_1_set, st_2_set]):
-            print(f" - struct_No: {i+1}")
+            print(f"  - struct_No: {i+1}")
             for st in st_set:
-                print("   spg_number:", st.spg_number)
-                print("   representations:")
+                print("    spg_number:", st.spg_number)
+                print("    representations:")
                 for h, pos in enumerate(st.positions):
                     log_axis_positions(st.symprec_set[h], st.axis[h], pos)
+        print("min_axis_l2_norm:", np.round(min_axis_l2_norm, 3))
         if axis_d_min is not None:
-            print("")
             print("difference_log:")
-            print(" - axis_l2_norm:", np.round(axis_d_min[0], 3))
-            print("   pos_max_abs:", np.round(pos_d_min[0], 3))
-            print("   structure_1:")
+            print("  - axis_l2_norm:", np.round(axis_d_min[0], 3))
+            print("    pos_max_abs:", np.round(pos_d_min[0], 3))
+            print("    structure_1:")
             log_axis_positions(pos_d_min[1][0], axis_d_min[1][0], pos_d_min[1][1])
-            print("   structure_2:")
+            print("    structure_2:")
             log_axis_positions(pos_d_min[1][2], axis_d_min[1][1], pos_d_min[1][3])
-            print("   diffs:")
+            print("    diffs:")
             log_axis_positions([], axis_d_min[1][2], pos_d_min[1][4])
-        print("")
         print("Match:", struct_match)
 
     return struct_match

{rsspolymlp-0.3.2 → rsspolymlp-0.3.4}/src/rsspolymlp/api/cli_rsspolymlp.py

@@ -165,12 +165,6 @@ def run():
         default=-1,
         help="Number of optimized structures to analyze (-1 means all)",
     )
-    parser.add_argument(
-        "--cutoff",
-        type=float,
-        default=None,
-        help="Cutoff radius used in the MLP model (optional)",
-    )
 
     # Target paths for rsspolymlp postprocessing
     parser.add_argument(
@@ -275,14 +269,11 @@ def run():
             c_maxiter=args.c_maxiter,
             n_opt_str=args.n_opt_str,
             not_stop_rss=args.not_stop_rss,
-            num_process=args.num_process,
-            backend=args.backend,
         )
 
     if args.uniq_struct:
         rss_uniq_struct(
             num_str=args.num_str,
-            cutoff=args.cutoff,
             num_process=args.num_process,
             backend=args.backend,
         )

{rsspolymlp-0.3.2 → rsspolymlp-0.3.4}/src/rsspolymlp/api/cli_rsspolymlp_devkit.py

@@ -208,8 +208,18 @@ def run():
         default=4,
         help="Number of processes to use with joblib.",
     )
+    parser.add_argument(
+        "--not_check",
+        action="store_true",
+    )
 
     # --divide_data mode
+    parser.add_argument(
+        "--target_pressure",
+        type=float,
+        default=0.0,
+        help="",
+    )
     parser.add_argument(
         "--th_e_high",
         type=float,
@@ -305,11 +315,14 @@ def run():
             vasp_paths=args.paths,
             output_dir=args.output_dir,
             num_process=args.num_process,
+            not_check=args.not_check,
         )
 
     if args.divide_data:
         divide_dft_dataset(
+            elements=args.elements,
             target_dirs=args.paths,
+            target_pressure=args.target_pressure,
             threshold_e_high=args.th_e_high,
             threshold_e_low=args.th_e_low,
             threshold_f_small=args.th_f_small,

rsspolymlp-0.3.4/src/rsspolymlp/api/cli_rsspolymlp_utils.py

@@ -0,0 +1,148 @@
+import argparse
+
+from rsspolymlp.api.rsspolymlp_utils import geometry_opt, struct_compare, struct_matcher
+
+
+def run():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--struct_matcher",
+        action="store_true",
+        help="Mode: struct_matcher",
+    )
+    parser.add_argument(
+        "--struct_compare",
+        action="store_true",
+        help="Mode: struct_compare",
+    )
+    parser.add_argument(
+        "--geometry_opt",
+        action="store_true",
+        help="Mode: geometry_optimizations",
+    )
+
+    # --struct_matcher mode
+    parser.add_argument(
+        "--poscars",
+        type=str,
+        nargs="+",
+        default=None,
+        help="Paths of target POSCAR files.",
+    )
+    parser.add_argument(
+        "--num_process",
+        type=int,
+        default=-1,
+        help="Number of processes to use with joblib. Use -1 to use all available CPU cores.",
+    )
+    parser.add_argument(
+        "--backend",
+        type=str,
+        choices=["loky", "threading", "multiprocessing"],
+        default="loky",
+        help="Backend for joblib parallelization",
+    )
+    parser.add_argument(
+        "--primitive_symprecs",
+        nargs="*",
+        type=float,
+        default=[1e-5, 1e-4, 1e-3, 1e-2],
+        help="List of symmetry tolerances used to identify distinct primitive cells.",
+    )
+    parser.add_argument(
+        "--output_file",
+        type=str,
+        default="unique_struct.yaml",
+        help="Output file name (default: unique_struct.yaml).",
+    )
+
+    # --struct_compare mode
+    parser.add_argument(
+        "--refine_symprec",
+        type=float,
+        default=1e-3,
+        help="Symmetry tolerance used to identify refine structure.",
+    )
+    parser.add_argument(
+        "--reduced_symprecs",
+        nargs="*",
+        type=float,
+        default=None,
+        help="List of symmetry tolerances used to identify reduced structure representation.",
+    )
+    parser.add_argument(
+        "--axis_tol",
+        type=float,
+        default=1e-2,
+        help="",
+    )
+    parser.add_argument(
+        "--pos_tol",
+        type=float,
+        default=1e-2,
+        help="",
+    )
+    parser.add_argument("--standardize_axis", action="store_true", help="")
+    parser.add_argument("--original_axis", action="store_true", help="")
+    parser.add_argument("--frac_coords", action="store_true", help="")
+    parser.add_argument("--not_verbose", action="store_true", help="")
+
+    # --geometry_opt mode
+    parser.add_argument(
+        "--pot",
+        nargs="*",
+        type=str,
+        default=["polymlp.yaml"],
+        help="Potential file for polynomial MLP",
+    )
+    parser.add_argument(
+        "--pressure", type=float, default=0.0, help="Pressure term (in GPa)"
+    )
+    parser.add_argument(
+        "--symmetry",
+        action="store_true",
+        help="If enabled, the optimization is comducted with using symmetry constraints.",
+    )
+    parser.add_argument(
+        "--solver_method", type=str, default="CG", help="Type of solver"
+    )
+    parser.add_argument(
+        "--c_maxiter",
+        type=int,
+        default=100,
+        help="Maximum number of iterations when c1 and c2 values are changed",
+    )
+
+    args = parser.parse_args()
+
+    if args.struct_matcher:
+        struct_matcher(
+            poscar_paths=args.poscars,
+            num_process=args.num_process,
+            backend=args.backend,
+            primitive_symprecs=args.primitive_symprecs,
+            output_file=args.output_file,
+        )
+
+    if args.struct_compare:
+        struct_compare(
+            poscar_paths=args.poscars,
+            refine_symprec=args.refine_symprec,
+            reduced_symprecs=args.reduced_symprecs,
+            axis_tol=args.axis_tol,
+            pos_tol=args.pos_tol,
+            standardize_axis=args.standardize_axis,
+            original_axis=args.original_axis,
+            frac_coords=args.frac_coords,
+            verbose=not args.not_verbose,
+        )
+
+    if args.geometry_opt:
+        geometry_opt(
+            poscar_paths=args.poscars,
+            pot=args.pot,
+            pressure=args.pressure,
+            with_symmetry=args.symmetry,
+            solver_method=args.solver_method,
+            c_maxiter=args.c_maxiter,
+        )

{rsspolymlp-0.3.2 → rsspolymlp-0.3.4}/src/rsspolymlp/api/rsspolymlp.py

@@ -1,11 +1,7 @@
 import glob
-import multiprocessing
 import os
 import time
 
-from joblib import Parallel, delayed
-from joblib.externals.loky import get_reusable_executor
-
 from rsspolymlp.analysis.ghost_minima import (
     detect_actual_ghost_minima,
     ghost_minima_candidates,
@@ -43,8 +39,6 @@ def rss_run_parallel(
     c_maxiter=100,
     n_opt_str=1000,
     not_stop_rss=False,
-    num_process=-1,
-    backend="loky",
 ):
     """
     Performing Random Structure Search (RSS) on multiple tasks in parallel
@@ -73,11 +67,6 @@ def rss_run_parallel(
     if len(poscar_path_all) == 0:
         return
 
-    if num_process == 1:
-        not_stop_rss = True
-    elif num_process == -1:
-        num_process = multiprocessing.cpu_count()
-
     rssobj = OptimizationMLP(
         pot=pot,
         pressure=pressure,
@@ -89,36 +78,22 @@ def rss_run_parallel(
     )
 
     time_start = time.time()
-    if num_process == 1:
-        for poscar in poscar_path_all:
-            rssobj.run_optimization(poscar)
-    else:
-        # Perform parallel optimization with joblib
-        time_start = time.time()
-        Parallel(n_jobs=num_process, backend=backend, verbose=100)(
-            delayed(rssobj.run_optimization)(poscar) for poscar in poscar_path_all
-        )
-        executor = get_reusable_executor(max_workers=num_process)
-        executor.shutdown(wait=True)
+    for poscar in poscar_path_all:
+        rssobj.run_optimization(poscar)
     elapsed = time.time() - time_start
 
-    with open("rss_result/parallel_time.log", "a") as f:
-        print("Number of CPU cores:", num_process, file=f)
-        print("Number of the structures:", len(glob.glob("log/*")), file=f)
+    with open("rss_result/elapsed_time.log", "a") as f:
+        print("Number of structures:", len(glob.glob("log/*")), file=f)
         print("Computational time:", elapsed, file=f)
         print("", file=f)
 
 
 def rss_uniq_struct(
     num_str=-1,
-    cutoff=None,
     num_process=-1,
     backend="loky",
 ):
     analyzer = RSSResultAnalyzer()
-    if cutoff is not None:
-        analyzer.cutoff = cutoff
-
     analyzer.run_rss_uniq_struct(
         num_str=num_str,
         num_process=num_process,
@@ -182,8 +157,6 @@ def rss_polymlp(
             c_maxiter=c_maxiter,
             n_opt_str=n_opt_str,
             not_stop_rss=not_stop_rss,
-            num_process=num_process,
-            backend=backend,
         )
 
         with open("rss_result/success.dat") as f:
@@ -192,7 +165,10 @@ def rss_polymlp(
             print("Target number of optimized structures reached. Exiting.")
             break
 
-    rss_uniq_struct()
+    rss_uniq_struct(
+        num_process=num_process,
+        backend=backend,
+    )
 
 
 def rss_summarize(

{rsspolymlp-0.3.2 → rsspolymlp-0.3.4}/src/rsspolymlp/api/rsspolymlp_devkit.py

@@ -140,9 +140,21 @@ def mlp_dataset(
 
 
 def compress_vasprun(
-    vasp_paths: list[str], output_dir: str = "compress_dft_data", num_process: int = 4
+    vasp_paths: list[str],
+    output_dir: str = "compress_dft_data",
+    num_process: int = 4,
+    not_check: bool = False,
 ):
-    valid_paths, vasprun_status = check_convergence(vasp_paths=vasp_paths)
+    if not not_check:
+        valid_paths, vasprun_status = check_convergence(vasp_paths=vasp_paths)
+    else:
+        valid_paths = vasp_paths
+        vasprun_status = {
+            "fail": 0,
+            "fail_iteration": 0,
+            "parse": 0,
+            "success": 0,
+        }
 
     judge_list = Parallel(n_jobs=num_process)(
         delayed(compress)(vasp_path + "/vasprun.xml", output_dir)
@@ -168,7 +180,9 @@ def compress_vasprun(
 
 
 def divide_dft_dataset(
+    elements: list[str],
     target_dirs: str,
+    target_pressure: float = 0.0,
     threshold_e_high: float = 10.0,  # in eV/atom
     threshold_e_low: Optional[float] = None,
     threshold_f_small: float = 3.0,  # in eV/ang
@@ -183,7 +197,9 @@ def divide_dft_dataset(
         vasprun_paths.extend(sorted(glob.glob(target_dir + "/*")))
 
     vasprun_dict = divide_dataset(
+        elements=elements,
         vasprun_paths=vasprun_paths,
+        target_pressure=target_pressure,
         threshold_e_high=threshold_e_high,
         threshold_e_low=threshold_e_low,
         threshold_f_small=threshold_f_small,
@@ -198,7 +214,9 @@ def divide_dft_dataset(
 
     with open(f"{output_dir}/dataset.yaml", "w") as f:
         print("arguments:", file=f)
+        print("  elements:", elements, file=f)
         print("  path:", target_dir, file=f)
+        print("  target_pressure:", target_pressure, file=f)
         print("  threshold_e_high:", threshold_e_high, file=f)
         print("  threshold_e_low:", threshold_e_low, file=f)
         print("  threshold_f_small:", threshold_f_small, file=f)