Merge pull request #52 from ahxbcn/develop

Update run_abacus_calculation

Author: pxlxingliang

README.md

@@ -20,7 +20,7 @@ conda activate abacus-agent
 ```
 2. Install necessary dependencies:
 ```bash
-pip install mcp google-adk litellm bohr-agent-sdk pymatgen abacustest==0.4.24
+pip install mcp google-adk litellm bohr-agent-sdk pymatgen abacustest
 ```
 3. Install ABACUS-agent-tools:
 ```bash
@@ -60,9 +60,13 @@ Note: When running `abacusagent`, it will automatically check if the file exists
     "_comments": {
         "ABACUS_WORK_PATH": "The working directory for AbacusAgent, where all temporary files will be stored.",
         "ABACUS_SUBMIT_TYPE": "The type of submission for ABACUS, can be local or bohrium.",
+        "ABACUSAGENT_TRANSPORT": "The transport protocol for AbacusAgent, can be 'sse' or 'streamable-http'.",
         "ABACUSAGENT_HOST": "The host address for the AbacusAgent server.",
         "ABACUSAGENT_PORT": "The port number for the AbacusAgent server.",
         "ABACUSAGENT_MODEL": "The model to use for AbacusAgent, can be 'fastmcp', 'test', or 'dp'.",
+        "LLM_MODEL": "The model name for the LLM to use. Like: openai/qwen-turbo, deepseek/deepseek-chat",
+        "LLM_API_KEY": "The API key for the LLM service.",
+        "LLM_BASE_URL": "The base URL for the LLM service, if applicable.",
         "BOHRIUM_USERNAME": "The username for Bohrium.",
         "BOHRIUM_PASSWORD": "The password for Bohrium.",
         "BOHRIUM_PROJECT_ID": "The project ID for Bohrium.",
@@ -71,7 +75,10 @@ Note: When running `abacusagent`, it will automatically check if the file exists
         "BOHRIUM_ABACUS_COMMAND": "The command to run Abacus on Bohrium",
         "ABACUS_COMMAND": "The command to execute Abacus on local machine.",
         "ABACUS_PP_PATH": "The path to the pseudopotential library for Abacus.",
-        "ABACUS_ORB_PATH": "The path to the orbital library for Abacus.",
+        "ABACUS_ORB_PATH": "The path to the orbital library for ABACUS_PP_PATH",
+        "ABACUS_SOC_PP_PATH": "The path to the SOC pseudopotential library for Abacus.",
+        "ABACUS_SOC_ORB_PATH": "The path to the orbital library for ABACUS_SOC_PP_PATH.",
+        "PYATB_COMMAND": "The command to execute PYATB on local machine.",
         "_comments": "This dictionary contains the default environment variables for AbacusAgent."
     }
 }
@@ -86,13 +93,17 @@ INFO:     Application startup complete.
 INFO:     Uvicorn running on http://localhost:50001 (Press CTRL+C to quit)
 ```
 #### Preparing Google ADK Agent
-Organize your agent files in the following structure:
+You can use the following command to create files needed by a new agent:
 ```
-name_of_your_agent/
+abacusagent --create
+```
+Then a directory containing all necessary files will be generated:
+```
+abacus-agent/
 ├── __init__.py
 └── agent.py
 ```
-See example files in `examples\my_agent` to prepare `__init__.py` and `agent.py`.
+Then you can edit the `agent.py` file to customize the agent.
 #### Starting Google ADK
 ```bash
 >>> adk web
@@ -146,3 +157,6 @@ Functions of ABACUS Agent tools are in active development. Currently, the follow
 - Charge density difference
 - Using Birch-Murganhan equation to fit equation of state for cubic crystals
 
+Besides, a wrapper function which accepts a structure file and some key parameters to do ABACUS calculation is also provided.
+
+You can use `abacusagent --screen-modules` to hide tool function in some modules. The `--matmaster` option will only allow the wrapper function to be loaded.

src/abacusagent/modules/submodules/dos.py

@@ -334,8 +334,8 @@ def plot_pdos_species(shifted_energy, orbitals, output_dir, nspin, dpi):
             plt.plot(shifted_energy, -pdos_data[1::2], label=f'{species_name} ' + r'$\downarrow$', linestyle='--', linewidth=1.0)
 
     plt.axvline(x=0, color='black', linestyle=':', linewidth=1.0)
-    plt.xlabel('Energy/eV', fontsize=10)
-    plt.ylabel(r'States/ev${^{-1}}$', fontsize=10)
+    plt.xlabel('Energy (eV)', fontsize=10)
+    plt.ylabel(r"States ($eV^{-1}$)", fontsize=10)
     plt.xlim(max(min(shifted_energy), -20), min(20, max(shifted_energy)))
     if nspin == 1:
         plt.ylim(bottom=0)
@@ -382,7 +382,7 @@ def plot_pdos_species_shell(shifted_energy, orbitals, output_dir, nspin, dpi):
 
         ax.axvline(x=0, color='black', linestyle=':', linewidth=1.0)
         ax.set_title(f'PDOS for {species}', fontsize=12, pad=10)
-        ax.set_ylabel(r'States/ev${^{-1}}$', fontsize=10)
+        ax.set_ylabel(r"States ($eV^{-1}$)", fontsize=10)
         ax.set_xlim(max(min(shifted_energy), -20), min(20, max(shifted_energy)))
         #if nspin == 1:
         #    ax.set_ylim(bottom=0)
@@ -391,7 +391,7 @@ def plot_pdos_species_shell(shifted_energy, orbitals, output_dir, nspin, dpi):
 
         #ax.set_ylim(bottom=0)
 
-    axes[-1].set_xlabel('Energy/eV', fontsize=10)
+    axes[-1].set_xlabel('Energy (eV)', fontsize=10)
 
     plt.tight_layout()
     pdos_pic_file = os.path.join(output_dir, 'PDOS.png')
@@ -460,13 +460,13 @@ def plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label
             ax.set_xlim(max(min(shifted_energy), -20), min(20, max(shifted_energy)))
             if nspin == 1:
                 ax.set_ylim(bottom=0)
-            ax.set_ylabel(r'States/ev${^{-1}}$', fontsize=10)
+            ax.set_ylabel(r"States ($eV^{-1}$)", fontsize=10)
             ax.legend(fontsize=8, ncol=nspin)
             ax.grid(alpha=0.3)
 
             subplot_count += 1
 
-    axes[-1].set_xlabel('Energy/eV', fontsize=10)
+    axes[-1].set_xlabel('Energy (eV)', fontsize=10)
 
     plt.tight_layout()
     pdos_pic_file = os.path.join(output_dir, 'PDOS.png')
@@ -500,8 +500,8 @@ def plot_dos(file_path: List[Path],
         plt.plot(energy, dos, linestyle='-', label='spin up')
         plt.plot(energy, -dos_dn, linestyle='--', label='spin down')
     plt.axvline(x=0, color='k', linestyle='--', alpha=0.5)
-    plt.xlabel('Energy/eV')
-    plt.ylabel(r'States/ev${^{-1}}$')
+    plt.xlabel('Energy (eV)')
+    plt.ylabel(r'States ($eV^{-1}$)')
     plt.title('Density of States')
     plt.grid(True, alpha=0.3)
     plt.xlim(x_min, x_max)

src/abacusagent/modules/submodules/eos.py

@@ -1,15 +1,15 @@
 import os
 from pathlib import Path
 from typing import Literal, Optional, TypedDict, Dict, Any, List
+import copy
 import numpy as np
 import matplotlib.pyplot as plt
 from abacustest.lib_prepare.abacus import AbacusStru, ReadInput, WriteInput
 from abacustest.lib_model.comm_eos import eos_fit
 from abacustest.lib_model.comm import check_abacus_inputs
 
 from abacusagent.init_mcp import mcp
-from abacusagent.modules.abacus import abacus_modify_input, abacus_modify_stru, abacus_collect_data
-from abacusagent.modules.util.comm import run_abacus, link_abacusjob, generate_work_path
+from abacusagent.modules.util.comm import run_abacus, link_abacusjob, generate_work_path, collect_metrics
 
 def is_cubic(cell: List[List[float]]) -> bool:
     """
@@ -76,7 +76,8 @@ def abacus_eos(
             raise ValueError("The structure is not cubic. Implemented EOS calculation requires a cubic structure.")
 
         # Generated lattice parameters for EOS calculation
-        original_cell_param = input_stru.get_cell()[0][0]  # Assuming cubic structure, take one cell parameter
+        original_cell = input_stru.get_cell()
+        original_cell_param = np.linalg.norm(original_cell[0])
         if stru_scale_type == 'percentage':
             scales = [1 + i * scale_stepsize for i in range(-stru_scale_number, stru_scale_number + 1)]
         elif stru_scale_type == 'angstrom':
@@ -86,9 +87,11 @@ def abacus_eos(
 
         scaled_lat_params = [original_cell_param * scale for scale in scales]
 
-        output = abacus_modify_input(abacus_inputs_dir, extra_input={'calculation': 'scf'})
+        input_params['calculation'] = 'scf'
+        WriteInput(input_params, os.path.join(abacus_inputs_dir, "INPUT"))
 
         scale_cell_job_dirs = []
+        stru = copy.deepcopy(input_stru)
         for i in range(len(scales)):
             dir_name = Path(os.path.join(work_path, f"scale_cell_{i}")).absolute()
             os.makedirs(dir_name, exist_ok=True)
@@ -103,13 +106,14 @@ def abacus_eos(
             )
 
             new_cell = (np.array(input_stru.get_cell()) * scales[i]).tolist()
-            output = abacus_modify_stru(dir_name, cell=new_cell, coord_change_type='scale')
+            stru.set_cell(new_cell, bohr=False, change_coord=True)
+            stru.write(os.path.join(dir_name, input_stru_file))
 
         run_abacus(scale_cell_job_dirs)
 
         energies = []
         for i, job_dir in enumerate(scale_cell_job_dirs):
-            metrics = abacus_collect_data(job_dir)['collected_metrics']
+            metrics = collect_metrics(job_dir)
             if metrics['normal_end'] is not True or metrics['converge'] is not True:
                 raise RuntimeError(f"Job {i} did not end normally or did not converge. Please check the job directory: {job_dir}")
             energies.append(metrics['energy'])
@@ -139,9 +143,8 @@ def abacus_eos(
             )
         optimal_lat_param = V0 ** (1.0 / 3)
         optimal_cell = (np.array(input_stru.get_cell()) * optimal_lat_param / original_cell_param).tolist()
-        output = abacus_modify_stru(new_abacus_inputs_dir,
-                                    cell = optimal_cell,
-                                    coord_change_type = 'scale')
+        stru.set_cell(optimal_cell, bohr=False, change_coord=True)
+        stru.write(os.path.join(new_abacus_inputs_dir, input_stru_file))
 
         return {
             "eos_work_path": work_path.absolute(),

src/abacusagent/modules/submodules/md.py

@@ -8,8 +8,7 @@
 from abacustest.lib_prepare.abacus import ReadInput, WriteInput, AbacusStru
 
 from abacusagent.init_mcp import mcp
-from abacusagent.modules.util.comm import generate_work_path, link_abacusjob, run_abacus
-from abacusagent.modules.abacus import abacus_collect_data
+from abacusagent.modules.util.comm import generate_work_path, link_abacusjob, run_abacus, collect_metrics
 
 
 def get_last_md_stru(md_stru_outputdir: Path) -> Path:
@@ -181,7 +180,7 @@ def abacus_run_md(
 
         run_abacus(work_path)
 
-        metrics = abacus_collect_data(work_path)['collected_metrics']
+        metrics = collect_metrics(work_path)
         suffix = input_params.get('suffix', 'ABACUS')
         md_traj_file, traj_frame_nums = convert_md_dump_to_ase_traj(Path(os.path.join(work_path, f'OUT.{suffix}/MD_dump')).absolute())
         return {'md_work_path': work_path,