Add a simple sampling strategy (#34)

No rotations or other nonesense. Just bump each atom

notebooks/1_explore-sampling-methods/4_simple-unfirom.ipynb

@@ -0,0 +1,276 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "8233267b-e98b-44be-b9aa-116d0e67a94b",
+   "metadata": {},
+   "source": [
+    "# Compute Energies of Random Offsets\n",
+    "Vary every coordinate uniformly between -step_size and step_size"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c6a28419-6831-4197-8973-00c5591e19cb",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "from jitterbug.utils import make_calculator\n",
+    "from ase.io import write, read\n",
+    "from ase.db import connect\n",
+    "from ase import Atoms\n",
+    "from pathlib import Path\n",
+    "from tqdm import tqdm \n",
+    "import numpy as np\n",
+    "import os"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cec456a7-3c13-4b00-936a-abc31c898262",
+   "metadata": {},
+   "source": [
+    "Configuration"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c6be56c5-a460-4acd-9b89-8c3d9c812f5f",
+   "metadata": {
+    "tags": [
+     "parameters"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "starting_geometry = '../data/exact/caffeine_pm7_None.xyz'\n",
+    "method = 'hf/def2-svpd'\n",
+    "threads = min(os.cpu_count(), 12)\n",
+    "step_size: float = 0.005 # Perturbation amount, used as maximum L2 norm"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7010df09-73b2-4d58-be03-15a5f0d04b4c",
+   "metadata": {},
+   "source": [
+    "Derived"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0b6794cd-477f-45a1-b96f-2332804ddb20",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "relax_name = Path(starting_geometry).name[:-4]\n",
+    "name, relax_method, relax_basis = relax_name.split(\"_\")\n",
+    "method, basis = method.split(\"/\")\n",
+    "run_name = f'{name}_{method}_{basis}_at_{relax_method}_{relax_basis}'\n",
+    "print(f'Run name: {run_name}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cf9ff792-6b5b-46ce-9a78-78912e372912",
+   "metadata": {},
+   "source": [
+    "## Load in the Relaxed Structure\n",
+    "We generated a relaxed structure in the previous notebook"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ad9fd725-b1ba-4fec-ae41-959be0e540b3",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "atoms = read(starting_geometry)\n",
+    "atoms"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2284056b-ddf2-4a3b-88ca-b1c6dc84a2d5",
+   "metadata": {},
+   "source": [
+    "## Compute many random energies\n",
+    "Compute $3N + 3N(3N+1)/2 + 1$ energies with displacements sampled [on the unit sphere](https://mathoverflow.net/questions/24688/efficiently-sampling-points-uniformly-from-the-surface-of-an-n-sphere). This is enough to fit the Jacobian and Hessian exactly plus a little more"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ad4c5d8e-96d4-4bb6-9bf2-6474d7563448",
+   "metadata": {},
+   "source": [
+    "Prepare the output directory"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "23502eea-0974-4248-8f19-e85447069c61",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "out_dir = Path('data') / 'simple-uniform'\n",
+    "out_dir.mkdir(exist_ok=True, parents=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bf1366fc-d9a7-4a98-b9c9-cb3a0209b406",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "db_path = out_dir / f'{run_name}_d={step_size:.2e}.db'"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "004158dc-3fe9-47a6-99dd-268aa69bb27b",
+   "metadata": {},
+   "source": [
+    "Add the relaxed geometry if needed"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d4f21e81-5ec3-4877-a4d1-402077be2ee8",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "if not db_path.is_file():\n",
+    "    with connect(db_path) as db:\n",
+    "        db.write(atoms)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "56ebf431-75a0-44d5-8e18-43f2898d6dab",
+   "metadata": {},
+   "source": [
+    "Make the calculator"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0915595d-133a-43df-84fc-4ff6a3b538ea",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "calc = make_calculator(method, basis, num_threads=threads)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8e9e5ff2-3728-459b-b3d3-09acba0f71bc",
+   "metadata": {},
+   "source": [
+    "Generate the energies"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e2a28593-2634-4bb7-ae5b-8f557937bda1",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "n_atoms = len(atoms)\n",
+    "to_compute = 3 * n_atoms + 3 * n_atoms * (3 * n_atoms + 1) // 2 + 1\n",
+    "print(f'Need to run {to_compute} calculations for full accuracy.')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8bf40523-dcaa-4046-a9c6-74e35178e87f",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "with connect(db_path) as db:\n",
+    "    done = len(db)\n",
+    "print(f'Already done {done}. {to_compute - done} left to do.')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a6fa1b33-defc-4b35-895d-052eb64453fb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pbar = tqdm(total=to_compute)\n",
+    "pbar.update(done)\n",
+    "for i in range(to_compute - done):\n",
+    "    # Sample a perturbation\n",
+    "    disp = np.random.normal(-step_size, step_size, size=(n_atoms, 3))\n",
+    "\n",
+    "    # Make the new atoms\n",
+    "    new_atoms = atoms.copy()\n",
+    "    new_atoms.positions += disp\n",
+    "\n",
+    "    # Compute the energy and store in the db\n",
+    "    new_atoms.calc = calc\n",
+    "    new_atoms.get_potential_energy()\n",
+    "    with connect(db_path) as db:\n",
+    "        db.write(new_atoms)\n",
+    "\n",
+    "    pbar.update(1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "785add47-39b5-4d7e-9d92-0375c8128171",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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.9.18"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

notebooks/1_explore-sampling-methods/run-all-methods.sh

@@ -3,23 +3,23 @@
 xyz=../data/exact/caffeine_pm7_None.xyz
 method='pm7/None'
 
-for step_size in 0.02; do
+for step_size in 0.005 0.01 0.02; do
     # Do the randomized methods
-    for notebook in 0_random-directions-same-distance.ipynb 1_random-directions-variable-distance.ipynb; do
+    for notebook in 0_random-directions-same-distance.ipynb 1_random-directions-variable-distance.ipynb 4_simple-unfirom.ipynb; do
         papermill -p starting_geometry $xyz -p method $method -p step_size $step_size $notebook last.ipynb
     done
 
     # Test with different reductions for "along axes"
     notebook=2_displace-along-axes.ipynb
-    for n in 2 4 8; do
+    for n in 4; do
         papermill -p starting_geometry $xyz -p method $method -p perturbs_per_evaluation $n -p step_size $step_size $notebook last.ipynb
     done
 done
 
 # Test with the vibrational modes
 notebook=3_displace-along-vibrational-modes.ipynb
-for step_size in 0.001 0.002; do  # These step sizes are energy scales in eV, not distances in Angstrom as above
-    for n in 16 32 64; do
+for step_size in 0.0025; do  # These step sizes are energy scales in eV, not distances in Angstrom as above
+    for n in 32; do
         papermill -p starting_geometry $xyz -p method $method -p perturbs_per_evaluation $n -p step_size $step_size $notebook last.ipynb
     done
 done