From f3e6406dddf2b14bebb5cb9bdd8484d3c08e8a25 Mon Sep 17 00:00:00 2001
From: Logan Ward <ward.logan.t@gmail.com>
Date: Tue, 3 Oct 2023 09:20:26 -0400
Subject: [PATCH] Add utility for making calculator

Avoid copying code between notebooks
---
 jitterbug/utils.py  | 36 ++++++++++++++++++++++++++++++++++++
 tests/test_utils.py | 11 +++++++++++
 2 files changed, 47 insertions(+)
 create mode 100644 jitterbug/utils.py
 create mode 100644 tests/test_utils.py

diff --git a/jitterbug/utils.py b/jitterbug/utils.py
new file mode 100644
index 0000000..149ffaf
--- /dev/null
+++ b/jitterbug/utils.py
@@ -0,0 +1,36 @@
+"""Utility functions"""
+from typing import Optional
+
+
+from ase.calculators.calculator import Calculator
+from ase.calculators.mopac import MOPAC
+from ase.calculators.psi4 import Psi4
+
+mopac_methods = ['pm7']
+"""List of methods for which we will use MOPAC"""
+
+
+def make_calculator(method: str, basis: Optional[str], **kwargs) -> Calculator:
+    """Make an ASE calculator that implements a desired method.
+
+    This function will select the appropriate quantum chemistry code depending
+    on the method name using the following rules:
+
+    1. Use MOPAC if the method is PM7.
+    2. Use Psi4 otherwise
+
+    Any keyword arguments are passed to the calculator
+
+    Args:
+        method: Name of the quantum chemistry method
+        basis: Basis set name, if appropriate
+    Returns:
+        Calculator defined according to the user's settings
+    """
+
+    if method in mopac_methods:
+        if basis is not None:
+            raise ValueError(f'Basis must be none for method: {method}')
+        return MOPAC(method=method, command='mopac PREFIX.mop > /dev/null')
+    else:
+        return Psi4(method=method, basis=basis, **kwargs)
diff --git a/tests/test_utils.py b/tests/test_utils.py
new file mode 100644
index 0000000..8f1f484
--- /dev/null
+++ b/tests/test_utils.py
@@ -0,0 +1,11 @@
+from pytest import raises
+
+from jitterbug.utils import make_calculator
+
+
+def test_make_calculator():
+    assert make_calculator('b3lyp', 'def2_svpd').name == 'psi4'
+    assert make_calculator('pm7', None).name == 'mopac'
+
+    with raises(ValueError):
+        make_calculator('pm7', 'wrong')