{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Interfacing with QC Oracles (ORCA)\n\nThis tutorial demonstrates how to use the ``OrcaEngine`` to perform\nquantum chemistry calculations.\n\nThe Oracle module acts as a bridge between your ML methods and\nexternal computational chemistry software. We will explore two\ndistinct ways to set up and execute an ORCA calculation:\n1. Defining fidelity parameters directly in Python.\n2. Using a pre-defined ORCA input template file. (recommended for general use)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Setting Up the Environment\nFirst, let's create a temporary working directory and generate a simple\nXYZ file for a water molecule that we can use for our calculations.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import os\nimport shutil\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom mfml_qc.oracles import OrcaEngine\n\n# this will be the working directory\nworkspace = \"tutorial_workspace\"\nos.makedirs(workspace, exist_ok=True)\n\nsingle_mol_xyz = os.path.join(workspace, \"test_water.xyz\")\n\n# Write a dummy water molecule to an XYZ file\nwith open(single_mol_xyz, \"w\") as f:\n    f.write(\"3\\nWater test molecule\\nO 0.000 0.000 0.000\\n\")\n    f.write(\"H 0.000 0.757 0.587\\nH 0.000 -0.757 0.587\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Method 1: Dynamic Parameter Generation\nThe first way to run a calculation is to pass all your desired ORCA keywords\ndirectly as a Python dictionary. The ``OrcaEngine`` will automatically\nconstruct the `.inp` file for you.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fidelity_params_dynamic = {\n    \"method\": \"CAM-B3LYP\",\n    \"basis\": \"def2-TZVP\",\n    \"charge\": 0,\n    \"multiplicity\": 1,\n    \"nprocs\": 2,  # Request 2 cores\n    \"maxcore\": 2000,  # Allocate max 2000 MB memory per core\n    \"optional_tags\": [\"TightSCF\", \"RIJCOSX\"],\n    \"EnGrad\": None,  # we will not compute gradients (Forces)\n    # We can inject multi-line blocks for complex tasks like TD-DFT:\n    \"custom_blocks\": \"%tddft\\n  ETol 1e-6\\n  RTol 1e-6\\n  nroots 10\\nend\",\n}\n\n# Initialize the engine.\n# NOTE: Update 'orca_path' to point to your actual ORCA binary,\n# e.g., '/opt/orca_5_0_3_linux_x86-64_openmpi411/orca'\norca_engine_1 = OrcaEngine(\n    orca_path=\"/../../../../opt/orca_5_0_3_linux_x86-64_openmpi411/orca\"\n)\n\n\ntry:\n    # We set clean=False so we can inspect the generated files later.\n    # return_outs=False means the engine will run but won't parse it automatically.\n    orca_engine_1.evaluate(\n        geometry=single_mol_xyz,\n        fidelity_params=fidelity_params_dynamic,\n        work_dir=os.path.join(workspace, \"orca_dynamic\"),\n        return_outs=False,\n        clean=False,\n    )\n\n    # We can manually define what to extract and parse the output file afterward:\n    orca_engine_1.properties_to_extract = {\"Total Energy\": \"FINAL SINGLE POINT ENERGY\"}\n    results_1 = orca_engine_1.parse_output(\n        os.path.join(workspace, \"orca_dynamic\", \"calc.out\")\n    )\n    print(\"Extracted Properties:\", results_1)\n\n    # let's also take a look at the input file generated:\n    with open(os.path.join(workspace, \"orca_dynamic\", \"calc.inp\"), \"r\") as f:\n        print(\"\\nGenerated Input File:\\n\", f.read())\n\n\n# in case you do not have ORCA installed you can still see the input file generated\nexcept FileNotFoundError:\n    print(\"ORCA executable not found in PATH! Skipping calculation execution.\")\n    print(\"However, the engine successfully generated the input file!\")\n\n    # Let's peek at the file it generated for us:\n    with open(os.path.join(workspace, \"orca_dynamic\", \"calc.inp\"), \"r\") as f:\n        print(\"\\nGenerated Input File:\\n\", f.read())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Method 2: Using an Input Template\nIf you have complex calculations with extensive blocks in your input files, you can simply\nwrite an ORCA template file on your local machine. The ``OrcaEngine`` will read\nthis template and dynamically append the molecular geometry block at the bottom.\n\nFor example, you would create a file named ``TZVP_template.inp`` in your\ndirectory with the following contents:\n\n```text\n! CAM-B3LYP def2-TZVP  def2/J TightSCF RIJCOSX EnGrad\n#%maxcore 2000\n%PAL nproc 2 end\n%tddft\n   ETol 1e-6\n   RTol 1e-6\n   nroots 10\n   maxdim 100\n   tprint 1E-10\n   triplets false\nend\n```\nFor the sake of this automated tutorial, we will write this string to a temporary file for you:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "template_file = os.path.join(workspace, \"TZVP_template.inp\")\n\n# notice that we will now ask for the energy gradients (with EnGrad keyword)\ntemplate_content = \"\"\"! CAM-B3LYP def2-TZVP  def2/J TightSCF RIJCOSX EnGrad\n#%maxcore 2000\n%PAL nproc 2 end\n%tddft\n   ETol 1e-6\n   RTol 1e-6\n   nroots 10\n   maxdim 100\n   tprint 1E-10\n   triplets false\nend\n\"\"\"\nwith open(template_file, \"w\") as f:\n    f.write(template_content)\n\nfidelity_params_template = {\n    \"template_file\": template_file,\n    \"charge\": 0,\n    \"multiplicity\": 1,\n}\n\n# Define the exact text labels to search for in the ORCA output file\nadvanced_properties = {\n    \"e_scf\": \"E(SCF)  =\",\n    \"de_cis\": \"DE(CIS) =\",\n    \"e_tot\": \"E(tot)  =\",\n    \"e_final\": \"FINAL SINGLE POINT ENERGY\",\n}\n\n# Initialize a new engine, passing the properties we want to extract upfront\norca_engine_2 = OrcaEngine(\n    orca_path=\"/../../../../opt/orca_5_0_3_linux_x86-64_openmpi411/orca\",\n    properties_to_extract=advanced_properties,\n)\n\norca_ran_successfully = False\n\ntry:\n    # By setting return_outs=True, evaluate() automatically parses the output!\n    results_2 = orca_engine_2.evaluate(\n        geometry=single_mol_xyz,\n        fidelity_params=fidelity_params_template,\n        work_dir=os.path.join(workspace, \"orca_template\"),\n        clean=False,\n        return_outs=True,\n    )\n    print(\"Parsed Results:\", results_2)\n\n    # %%\n    # Parsing Spectra and Gradients\n    # -----------------------------\n    # Because we included 'EnGrad' and '%tddft' in our template, ORCA generates\n    # gradients and absorption spectra. We can manually call the parser on the\n    # generated output file to extract these advanced arrays!\n\n    output_path = os.path.join(workspace, \"orca_template\", \"calc.out\")\n    results_advanced = orca_engine_2.parse_output(\n        output_file=output_path, parse_gradients=True, parse_spectra=True\n    )\n\n    if \"gradients\" in results_advanced:\n        print(\"First atom gradients (x, y, z):\", results_advanced[\"gradients\"][0])\n\n    if \"tddft_spectrum\" in results_advanced:\n        # The spectrum is returned as a list of state dictionaries\n        state_1 = results_advanced[\"tddft_spectrum\"][0]\n        print(\"\\nTD-DFT State 1:\")\n        for key, val in state_1.items():\n            print(f\"  {key}: {val}\")\n\n    orca_ran_successfully = True\n\nexcept FileNotFoundError:\n    print(\"ORCA executable not found in PATH! Skipping calculation execution.\")\n    with open(os.path.join(workspace, \"orca_template\", \"calc.inp\"), \"r\") as f:\n        print(\"\\nGenerated Template Input File:\\n\", f.read())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Visualizing the Oracle Results\nLet's visualize the target water molecule. If ORCA successfully executed on\nyour machine, we will overlay the computed properties (Energy, Gradients,\nand TD-DFT spectra) directly onto the plot!\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fig = plt.figure(figsize=(6, 6))\nax = fig.add_subplot(111, projection=\"3d\")\n\n# coordinates of our test water molecule\ncoords = np.array(\n    [\n        [0.000, 0.000, 0.000],  # Oxygen\n        [0.000, 0.757, 0.587],  # Hydrogen\n        [0.000, -0.757, 0.587],  # Hydrogen\n    ]\n)\natoms = [\"O\", \"H\", \"H\"]\ncolors = [\"red\", \"lightgrey\", \"lightgrey\"]\nsizes = [600, 250, 250]\n\n# Plot the atoms\nfor i in range(3):\n    ax.scatter(\n        coords[i, 0],\n        coords[i, 1],\n        coords[i, 2],\n        c=colors[i],\n        s=sizes[i],\n        edgecolors=\"k\",\n        depthshade=False,\n        zorder=5,\n    )\n    # Add small text labels next to atoms\n    ax.text(\n        coords[i, 0],\n        coords[i, 1],\n        coords[i, 2] + 0.1,\n        atoms[i],\n        fontsize=12,\n        ha=\"center\",\n        weight=\"bold\",\n        zorder=10,\n    )\n\n# Draw bonds (O-H)\nax.plot(\n    [coords[0, 0], coords[1, 0]],\n    [coords[0, 1], coords[1, 1]],\n    [coords[0, 2], coords[1, 2]],\n    \"k-\",\n    lw=3,\n    zorder=1,\n)\nax.plot(\n    [coords[0, 0], coords[2, 0]],\n    [coords[0, 1], coords[2, 1]],\n    [coords[0, 2], coords[2, 2]],\n    \"k-\",\n    lw=3,\n    zorder=1,\n)\n\nax.set_axis_off()\nax.set_title(\"Properties from the ORCA Oracle\", fontsize=14, weight=\"bold\", pad=20)\nax.view_init(elev=20, azim=45)\n\n# text box of properties\ntext_str = \"\"\nif orca_ran_successfully and results_2.get(\"success\", False):\n    energy = results_2.get(\"e_final\", \"N/A\")\n    text_str += f\"Total Energy:\\n  {energy:.3f} Eh\\n\\n\"\n\n    if \"gradients\" in results_advanced:\n        max_grad = np.max(np.abs(results_advanced[\"gradients\"]))\n        text_str += f\"Max Gradient:\\n  {max_grad:.6f} Eh/bohr\\n\\n\"\n\n    if \"tddft_spectrum\" in results_advanced:\n        text_str += \"First 3 Excited States:\\n\"\n        for state in results_advanced[\"tddft_spectrum\"][:3]:\n            text_str += f\"  S{state['state']}: {state['energy_cm1']:>8.1f} cm\u207b\u00b9\\n\"\nelse:\n    text_str += \"ORCA Binary Not Found.\\nDisplaying Geometry Only.\"\n\nprops = dict(\n    boxstyle=\"round,pad=0.8\", facecolor=\"whitesmoke\", alpha=0.9, edgecolor=\"gray\"\n)\nax.text2D(\n    0.05,\n    0.2,\n    text_str,\n    transform=ax.transAxes,\n    fontsize=11,\n    verticalalignment=\"top\",\n    bbox=props,\n    family=\"monospace\",\n)\n\nplt.tight_layout()\nplt.show()\n\n# Clean up our tutorial workspace\n# comment this line if you wish to see what ORCA produced\nshutil.rmtree(workspace)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}