BAGEL engine with support for CASPT2

examples/1-simple-examples/benzene_bagel_casscf/README.md

@@ -0,0 +1,2 @@
+An example optimisation of the electronic ground state of benzene using 
+CASSCF as inplemented in Bagel. No parallelisation is used here.

examples/1-simple-examples/benzene_bagel_casscf/benzene.json

@@ -0,0 +1,35 @@
+{ "bagel": [
+
+{
+    "title": "molecule",
+    "basis": "/home/miaoxin/software/bagel-1.2.2/obj/share/cc-pvdz.json",
+    "df_basis": "/home/miaoxin/software/bagel-1.2.2/obj/share/cc-pvdz-jkfit.json",
+    "angstrom": true,
+    "geometry": [
+        { "atom": "C", "xyz": [  1.39000000,  0.00000000,  0.00000000] },
+        { "atom": "C", "xyz": [  0.69500000,  1.20377500,  0.00000000] },
+        { "atom": "C", "xyz": [ -0.69500000,  1.20377500,  0.00000000] },
+        { "atom": "C", "xyz": [ -1.39000000,  0.00000000,  0.00000000] },
+        { "atom": "C", "xyz": [ -0.69500000, -1.20377500,  0.00000000] },
+        { "atom": "C", "xyz": [  0.69500000, -1.20377500,  0.00000000] },
+        { "atom": "H", "xyz": [  2.46000000,  0.00000000,  0.00000000] },
+        { "atom": "H", "xyz": [  1.23000000,  2.13042200,  0.00000000] },
+        { "atom": "H", "xyz": [ -1.23000000,  2.13042200,  0.00000000] },
+        { "atom": "H", "xyz": [ -2.46000000,  0.00000000,  0.00000000] },
+        { "atom": "H", "xyz": [ -1.23000000, -2.13042200,  0.00000000] },
+        { "atom": "H", "xyz": [  1.23000000, -2.13042200,  0.00000000] }
+    ]
+},
+
+{
+    "title": "force",
+    "method": [ {
+        "title": "casscf",
+        "nact": 6,
+        "nclosed": 18,
+        "active": [17, 20, 21, 22, 23, 30],
+        "thresh": 1e-8
+    } ]
+}
+
+]}

examples/1-simple-examples/benzene_bagel_casscf/command.sh

@@ -0,0 +1 @@
+geometric-optimize --engine bagel benzene.json > opt.log

examples/1-simple-examples/ethene_bagel_caspt2/README.md

@@ -0,0 +1,9 @@
+An example optimisation of the $\pi-\pi^\*$ state of ethene using XMS-CASPT2 
+as implemented in Bagel. 
+Here, Bagel is executed in parallel using MPI managed by slurm
+with 4 MPI processes and 2 Bagel threads per process.
+
+This optimisation will not converge using default parameters, since 
+the minimum is located at a conical intersection where the 
+adiabatic state energies are not differentiable. Therefore, the gradients 
+evaluated by Bagel are not reliable and will not vanish at the minimum.

examples/1-simple-examples/ethene_bagel_caspt2/command.sh

@@ -0,0 +1 @@
+geometric-optimize --engine bagel --nt 2 --bagelexe "srun --ntasks-per-node=4 --mpi=pmi2 BAGEL" ethene.json > opt.log

examples/1-simple-examples/ethene_bagel_caspt2/ethene.json

@@ -0,0 +1,47 @@
+{ "bagel" : [
+
+{
+    "title" : "molecule",
+    "basis": "/home/miaoxin/software/bagel-1.2.2/obj/share/cc-pvdz.json",
+    "df_basis": "/home/miaoxin/software/bagel-1.2.2/obj/share/cc-pvdz-jkfit.json",
+    "angstrom" : true,
+    "geometry" : [
+        { "atom": "C", "xyz": [  0.6579, -0.0045,  0.0649] },
+        { "atom": "C", "xyz": [ -0.6579,  0.0045, -0.0639] },
+        { "atom": "H", "xyz": [  1.1610,  0.0661,  1.0238] },
+        { "atom": "H", "xyz": [  1.3352, -0.0830, -0.7815] },
+        { "atom": "H", "xyz": [ -1.3355,  0.0830,  0.7812] },
+        { "atom": "H", "xyz": [ -1.1608, -0.0661, -1.0239] }]
+},
+
+{
+    "title": "casscf",
+    "nstate": 3,
+    "nact": 2,
+    "nclosed": 7,
+    "active": [8, 9],
+    "maxiter": 200,
+    "thresh": 1e-9
+},
+
+{
+    "title": "force",
+    "target": 1,
+    "method": [ {
+        "title": "caspt2",
+        "smith": {
+            "method": "caspt2",
+            "ms": true,
+            "xms": true,
+            "sssr": true,
+            "shift": 0.2,
+            "maxiter": 200
+        }, 
+        "nstate": 3,
+        "nact" : 2,
+        "nclosed" : 7,
+        "thresh" : 1e-8
+    } ]
+}
+
+]}

geometric/engine.py

@@ -42,11 +42,13 @@
 import numpy as np
 import re
 import os
+import json
 
 from .molecule import Molecule, format_xyz_coord
 from .nifty import bak, au2ev, eqcgmx, fqcgmx, bohr2ang, logger, getWorkQueue, queue_up_src_dest, rootdir, copy_tree_over
 from .errors import EngineError, CheckCoordError, Psi4EngineError, QChemEngineError, TeraChemEngineError, \
-    ConicalIntersectionEngineError, OpenMMEngineError, GromacsEngineError, MolproEngineError, QCEngineAPIEngineError, GaussianEngineError, QUICKEngineError, CFOUREngineError
+    ConicalIntersectionEngineError, OpenMMEngineError, GromacsEngineError, MolproEngineError, \
+    QCEngineAPIEngineError, GaussianEngineError, QUICKEngineError, CFOUREngineError, BagelEngineError
 from .xml_helper import read_coors_from_xml, write_coors_to_xml
 
 # Strings matching common DFT functionals
@@ -1838,6 +1840,233 @@ def detect_dft(self):
                     return True
         return False
 
+class Bagel(Engine):
+    """
+    Run a Bagel energy and gradient calculation.
+    """
+    def __init__(self, molecule=None, exe=None, threads=None):
+        """
+        Initialize the Bagel engine.
+
+        Parameters
+        ----------
+        molecule : Molecule
+            The molecule object to be used in the calculation.
+            Default is None, in which case the molecule is created from the input file.
+        exe : str
+            The path to the Bagel executable.
+            Default is "BAGEL". For parallel execution, this parameter should
+            be set e.g. to "mpirun BAGEL" or "srun --mpi=pmi2 BAGEL", etc.,
+            with specification of the number of threads if not set by
+            environment variables.
+        threads : int
+            The number of threads to start every Bagel process.
+            Default is 1. This parameter is used to set the environment variables
+            BAGEL_NUM_THREADS and MKL_NUM_THREADS, i.e.
+            $ export BAGEL_NUM_THREADS=threads
+            $ export MKL_NUM_THREADS=threads
+        """
+
+        if molecule is None:
+            # create a fake molecule
+            molecule = Molecule()
+            molecule.elem = ['H']
+            molecule.xyzs = [[[0,0,0]]]
+        super(Bagel, self).__init__(molecule)
+
+        if exe is None:
+            self.bagel_exe = 'BAGEL'
+        else:
+            if 'BAGEL' in exe:
+                self.bagel_exe = exe
+            else:
+                raise ValueError('Bagel executable must be named BAGEL')
+
+        self.threads = threads if threads is not None else 1
+
+    def load_bagel_input(self, bagel_input):
+        """
+        Bagel .json input parser. Supports both angstroms and bohr.
+        """
+        with open(bagel_input,'r') as bageljson:
+            bagel_dict = json.load(bageljson, object_pairs_hook=OrderedDict)
+
+        for i in range(0, len(bagel_dict['bagel'])):
+            if 'title' in bagel_dict['bagel'][i]:
+                bagel_dict['bagel'][i]['title'] = bagel_dict['bagel'][i]['title'].lower()
+
+        bagel_titles = [dct['title'] for dct in bagel_dict['bagel']]
+        if bagel_titles[0] != 'molecule':
+            raise RuntimeError(
+                f'The first dictionary in Bagel json file {bagel_input} '
+                'must have the title "molecule".'
+            )
+        if not 'force' in bagel_titles:
+            raise RuntimeError(
+                f'The Bagel json file {bagel_input} must have a dictionary '
+                'with the title "force".'
+            )
+        self.mol_idx = bagel_titles.index('molecule')
+        force_idx = bagel_titles.index('force')
+        self.force_method = bagel_dict['bagel'][force_idx]['method'][0]['title'].lower()
+
+        if self.force_method in ['fci', 'ras', 'nevpt2', 'zfci', 'zcasscf', 'relsmith']:
+            raise NotImplementedError(
+                'FCI, RASCI, NEVPT2 and relativistic methods '
+                'except for Dirac-HF are not currently supported '
+                'due to the lack of analytic gradients.'
+            )
+        elif self.force_method in ['casscf', 'caspt2']:
+            if self.force_method == 'caspt2':
+                # Check if CASSCF section is present
+                if 'casscf' not in bagel_titles:
+                    raise RuntimeError(
+                        'You requested a CASPT2 calculation, therefore the '
+                        f'Bagel json file {bagel_input} must have a CASSCF section.'
+                    )
+                else:
+                    if bagel_titles.index('casscf') > force_idx:
+                        raise RuntimeError(
+                            'The CASSCF section must be before the force section.'
+                        )
+
+            # Check if a guess wavefunction is provided
+            if 'load_ref' in bagel_titles:
+                self.ref_idx = bagel_titles.index('load_ref')
+                self.ref_file = bagel_dict['bagel'][self.ref_idx]['file']
+                bagel_dict['bagel'][self.ref_idx]['continue_geom'] = False
+            else:
+                self.ref_file = None
+            # Export CASSCF wavefunction for the next step
+            if 'save_ref' in bagel_titles:
+                save_idx = bagel_titles.index('save_ref')
+                bagel_dict['bagel'][save_idx]['file'] = 'run'
+            else:
+                bagel_dict['bagel'].append({
+                    'title': 'save_ref',
+                    'file': 'run',
+                })
+
+        # Add export keyword to the "force" dictionary if it is not present
+        if 'export' not in [key.lower() for key in bagel_dict['bagel'][force_idx]]:
+            bagel_dict['bagel'][force_idx]['export'] = True
+        # else:
+        #     bagel_dict['bagel'][force_idx].move_to_end('export')
+
+        elems = []
+        coords = []
+        readang = False
+
+        # Bagel by default reads coordinates in Bohr, check to see which input is used.
+        molecule_lower_keys = {
+            k.lower(): k for k in bagel_dict['bagel'][self.mol_idx]
+        }
+        if 'angstrom' in molecule_lower_keys:
+            readang = bagel_dict['bagel'][self.mol_idx][molecule_lower_keys['angstrom']]
+
+        for atom_dict in bagel_dict['bagel'][self.mol_idx]['geometry']:
+            elems.append(atom_dict['atom'])  # extract elements
+            coords.append(atom_dict['xyz'])  # extract coordinates
+
+        self.M = Molecule()
+        self.M.elem = elems
+
+        if readang:
+            self.M.xyzs = [np.array(coords, dtype=np.float64)]
+        elif not readang:  # save coordinates in molecule object as angstroms
+            self.M.xyzs = [np.array(coords, dtype=np.float64) * bohr2ang]
+            # for consistency geomeTRIC will create Bagel files that use angstrom
+            bagel_dict['bagel'][self.mol_idx]['angstrom'] = True
+
+        self.bagel_temp = bagel_dict
+
+    def calc_new(self, coords, dirname):
+        if not os.path.exists(dirname):
+            os.makedirs(dirname)
+        # Convert coordinates back to the xyz file
+        self.M.xyzs[0] = coords.reshape(-1, 3) * bohr2ang
+
+        # Change coordinates in bagel_temp dictionary
+        for atom in range(self.M.xyzs[0].shape[0]):
+            self.bagel_temp['bagel'][0]['geometry'][atom]['xyz'] = list(self.M.xyzs[0][atom])
+
+        # Handle CASSCF guess wavefunction
+        if self.force_method in ['casscf', 'caspt2']:
+            if self.ref_file is None:
+                if os.path.exists(os.path.join(dirname, 'run.archive')):
+                    if 'load_ref' not in [dct['title'].lower() for dct in self.bagel_temp['bagel']]:
+                        self.bagel_temp['bagel'].insert(
+                            self.mol_idx + 1,  # insert after the molecule section
+                            {
+                                'title': 'load_ref',
+                                'file': 'run',
+                                'continue_geom': False,
+                            },
+                        )
+                        # Remove the custom active space from the CASSCF section
+                        # since the active orbitals in the reference wavefunction
+                        # are already in the middle.
+                        if self.force_method == 'casscf':
+                            force_idx = [dct['title'] for dct in self.bagel_temp['bagel']].index('force')
+                            self.bagel_temp['bagel'][force_idx]['method'][0].pop('active', None)
+                        elif self.force_method == 'caspt2':
+                            casscf_idx = [dct['title'] for dct in self.bagel_temp['bagel']].index('casscf')
+                            self.bagel_temp['bagel'][casscf_idx].pop('active', None)
+            else:
+                if os.path.exists(os.path.join(dirname, 'run.archive')):
+                    self.bagel_temp['bagel'][self.ref_idx]['file'] = 'run'
+                else:
+                    shutil.copy(
+                        f'{self.ref_file}.archive',
+                        os.path.join(dirname, f'{self.ref_file}.archive'),
+                    )
+
+        # Write Json file
+        with open(os.path.join(dirname, 'run.json'), 'w') as outfile:
+            outfile.write(json.dumps(self.bagel_temp, indent=2))
+
+        try:
+            # Run Bagel
+            bagel_env = os.environ.copy()
+            bagel_env['BAGEL_NUM_THREADS'] = str(self.threads)
+            bagel_env['MKL_NUM_THREADS'] = str(self.threads)
+            subprocess.check_call(
+                f'{self.bagel_exe} run.json 1> run.out 2> run.err', 
+                cwd=dirname, shell=True, env=bagel_env,
+            )
+            # Read energy and gradients from Bagel output
+            result = self.read_result(dirname)
+        except (OSError, IOError, RuntimeError, subprocess.CalledProcessError):
+            raise BagelEngineError
+        return result
+
+
+    def read_result(self, dirname, check_coord=None):
+        """ read an output file from Bagel"""
+        if check_coord is not None:
+            raise CheckCoordError("Coordinate check not implemented")
+
+        force_idx = [dct['title'] for dct in self.bagel_temp['bagel']].index('force')
+        istate = self.bagel_temp['bagel'][force_idx].get('target', 0)
+
+        energy = np.loadtxt(os.path.join(dirname, 'ENERGY.out'))
+        if energy.ndim == 0:
+            energy = float(energy)
+        else:
+            energy = energy[istate]
+
+        gradient = np.loadtxt(
+            os.path.join(dirname, f'FORCE_{istate}.out'),
+            skiprows=1,
+            usecols=(1, 2, 3),
+        ).flatten()
+
+        # Remove exported files
+        os.remove(os.path.join(dirname, 'ENERGY.out'))
+        os.remove(os.path.join(dirname, f'FORCE_{istate}.out'))
+
+        return {'energy': energy, 'gradient': gradient}
+
 class QCEngineAPI(Engine):
     def __init__(self, schema, program):
         try:

geometric/errors.py

@@ -106,6 +106,9 @@ class GaussianEngineError(EngineError):
 class QCEngineAPIEngineError(EngineError):
     pass
 
+class BagelEngineError(EngineError):
+    pass
+
 class ConicalIntersectionEngineError(EngineError):
     pass
 

geometric/params.py

@@ -331,7 +331,7 @@ def parse_optimizer_args(*args):
                           '"psi4" = Psi4                       "openmm" = OpenMM (pass a force field or XML input file)\n'
                           '"molpro" = Molpro                   "gmx" = Gromacs (pass conf.gro; requires topol.top and shot.mdp\n '
                           '"gaussian" = Gaussian09/16          "ase" = ASE calculator, use --ase-class/--ase-kwargs\n '
-                          '"quick" = QUICK\n')
+                          '"quick" = QUICK                     "bagel" = Bagel\n')
     grp_univ.add_argument('--nt', type=int, help='Specify number of threads for running in parallel\n(for TeraChem this should be number of GPUs)')
 
     grp_jobtype = parser.add_argument_group('jobtype', 'Control the type of optimization job')
@@ -410,6 +410,7 @@ def parse_optimizer_args(*args):
     grp_software.add_argument('--molcnv', type=str2bool, help='Provide "yes" to use Molpro style convergence criteria instead of the default.\n ')
     grp_software.add_argument('--qcdir', type=str, help='Provide an initial Q-Chem scratch folder (e.g. supplied initial guess).\n ')
     grp_software.add_argument('--qccnv', type=str2bool, help='Provide "yes" to Use Q-Chem style convergence criteria instead of the default.\n ')
+    grp_software.add_argument('--bagelexe', type=str, help='Specify how to run Bagel, e.g. "mpirun -np 4 BAGEL".\n ')
 
     grp_software.add_argument(
         '--ase-class',