From a390ba64912df3b4b0abcab4d89e123fde72e350 Mon Sep 17 00:00:00 2001
From: Anup Kumar <anukumar94663@gmail.com>
Date: Thu, 4 Jul 2024 05:36:33 -0700
Subject: [PATCH] added workflow code for neb ts hess nn

---
 .../neb_ts_with_hessian/using_newtonnet.py    | 100 ++++++++++++++++++
 1 file changed, 100 insertions(+)
 create mode 100644 src/ts-workflow-examples/neb_ts_with_hessian/using_newtonnet.py

diff --git a/src/ts-workflow-examples/neb_ts_with_hessian/using_newtonnet.py b/src/ts-workflow-examples/neb_ts_with_hessian/using_newtonnet.py
new file mode 100644
index 0000000..4d3b408
--- /dev/null
+++ b/src/ts-workflow-examples/neb_ts_with_hessian/using_newtonnet.py
@@ -0,0 +1,100 @@
+import os
+import logging
+import toml
+from ase.io import read
+from quacc.recipes.newtonnet.ts import ts_job, irc_job, neb_job
+import jobflow as jf
+
+# Load configuration from TOML file
+config = toml.load('config.toml')
+
+# Constants from TOML file
+INPUTS_DIR = config['paths']['inputs_dir']
+MODEL_PATH = config['paths']['model_path']
+SETTINGS_PATH = config['paths']['settings_path']
+LAUNCHPAD_FILE = config['paths']['launchpad_file']
+TAG = config['run']['tag']
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Calculation and optimization keyword arguments
+calc_kwargs1 = {
+    'model_path': MODEL_PATH,
+    'settings_path': SETTINGS_PATH,
+    'hess_method': None,
+}
+calc_kwargs2 = {
+    'model_path': MODEL_PATH,
+    'settings_path': SETTINGS_PATH,
+    'hess_method': 'autograd',
+}
+
+def main():
+    try:
+        # Read reactant and product structures
+        reactant = read(os.path.join(INPUTS_DIR, 'R.xyz'))
+        product = read(os.path.join(INPUTS_DIR, 'P.xyz'))
+        logger.info("Successfully read reactant and product structures.")
+    except Exception as e:
+        logger.error(f"Error reading reactant and product structures: {e}")
+        return
+
+    try:
+        # Create NEB job
+        job1 = neb_job(reactant, product, calc_kwargs=calc_kwargs1)
+        job1.update_metadata({"tag": f'neb_{TAG}'})
+        logger.info("Created NEB job.")
+    except Exception as e:
+        logger.error(f"Error creating NEB job: {e}")
+        return
+
+    try:
+        # Create TS job with custom Hessian
+        job2 = ts_job(job1.output['neb_results']['highest_e_atoms'], use_custom_hessian=True, **calc_kwargs2)
+        job2.update_metadata({"tag": f'ts_hess_{TAG}'})
+        logger.info("Created TS job with custom Hessian.")
+    except Exception as e:
+        logger.error(f"Error creating TS job: {e}")
+        return
+
+    try:
+        # Create IRC job in forward direction
+        job3 = irc_job(job2.output['atoms'], direction='forward', **calc_kwargs1)
+        job3.update_metadata({"tag": f'ircf_hess_{TAG}'})
+        logger.info("Created IRC job in forward direction.")
+    except Exception as e:
+        logger.error(f"Error creating IRC job in forward direction: {e}")
+        return
+
+    try:
+        # Create IRC job in reverse direction
+        job4 = irc_job(job2.output['atoms'], direction='reverse', **calc_kwargs1)
+        job4.update_metadata({"tag": f'ircr_hess_{TAG}'})
+        logger.info("Created IRC job in reverse direction.")
+    except Exception as e:
+        logger.error(f"Error creating IRC job in reverse direction: {e}")
+        return
+
+    try:
+        # Combine jobs into a flow
+        jobs = [job1, job2, job3, job4]
+        flow = jf.Flow(jobs)
+        logger.info("Jobs combined into a flow.")
+    except Exception as e:
+        logger.error(f"Error combining jobs into a flow: {e}")
+        return
+
+    try:
+        # Execute the workflow locally
+        responses = jf.managers.local.run_locally(flow)
+        logger.info("Workflow executed successfully.")
+        logger.info(f"Responses: {responses}")
+    except Exception as e:
+        logger.error(f"Error executing workflow: {e}")
+        return
+
+
+if __name__ == "__main__":
+    main()