From b96a3ca025be03679933fac4e6d438442283cb08 Mon Sep 17 00:00:00 2001
From: Xinzijian Liu <liuxin_zijian@163.com>
Date: Tue, 7 May 2024 11:35:52 +0800
Subject: [PATCH] Fix dpdata format of abacus (#218)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit


- **New Features**
- Introduced a new function to enhance data processing by extracting
suffix and calculation details from inputs.
- **Enhancements**
- Updated execution logic to better determine calculation types and
format outputs effectively.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: zjgemi <liuxin_zijian@163.com>
Co-authored-by: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 dpgen2/fp/abacus.py | 29 ++++++++++++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/dpgen2/fp/abacus.py b/dpgen2/fp/abacus.py
index 90159079..1ee4d8bf 100644
--- a/dpgen2/fp/abacus.py
+++ b/dpgen2/fp/abacus.py
@@ -132,6 +132,23 @@ def execute(
         return op.execute(op_in)  # type: ignore in the case of not importing fpop
 
 
+from typing import (
+    Tuple,
+)
+
+
+def get_suffix_calculation(INPUT: List[str]) -> Tuple[str, str]:
+    suffix = "ABACUS"
+    calculation = "scf"
+    for iline in INPUT:
+        sline = iline.split("#")[0].split()
+        if len(sline) >= 2 and sline[0].lower() == "suffix":
+            suffix = sline[1].strip()
+        elif len(sline) >= 2 and sline[0].lower() == "calculation":
+            calculation = sline[1].strip()
+    return suffix, calculation
+
+
 class RunFpOpAbacus(OP):
     @classmethod
     def get_input_sign(cls):
@@ -171,7 +188,17 @@ def execute(
         workdir = op_out["backward_dir"].parent
 
         # convert the output to deepmd/npy format
-        sys = dpdata.LabeledSystem(str(workdir), fmt="abacus/scf")
+        with open("%s/INPUT" % workdir, "r") as f:
+            INPUT = f.readlines()
+        _, calculation = get_suffix_calculation(INPUT)
+        if calculation == "scf":
+            sys = dpdata.LabeledSystem(str(workdir), fmt="abacus/scf")
+        elif calculation == "md":
+            sys = dpdata.LabeledSystem(str(workdir), fmt="abacus/md")
+        elif calculation in ["relax", "cell-relax"]:
+            sys = dpdata.LabeledSystem(str(workdir), fmt="abacus/relax")
+        else:
+            raise ValueError("Type of calculation %s not supported" % calculation)
         out_name = run_config.get("out", fp_default_out_data_name)
         sys.to("deepmd/npy", workdir / out_name)