diff --git a/notebooks/1_explore-sampling-methods/2_displace-along-axes.ipynb b/notebooks/1_explore-sampling-methods/2_displace-along-axes.ipynb new file mode 100644 index 0000000..00657ef --- /dev/null +++ b/notebooks/1_explore-sampling-methods/2_displace-along-axes.ipynb @@ -0,0 +1,360 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "8233267b-e98b-44be-b9aa-116d0e67a94b", + "metadata": {}, + "source": [ + "# Compute Energies of Displacements Along Coordinate Systems\n", + "Displace each atom along the axes, in the same way that we would when computing finite differences. The code is similar to what is used in [ASE's vibration analysis](https://databases.fysik.dtu.dk/ase/ase/vibrations/vibrations.html) with a few differences:\n", + "\n", + "- We perturb every pair of coordinates to compute numerical second derivatives. ASE compute the first derivatives of force to access the Hessian.\n", + "- Optionally, perturb more than pair at a time." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "c6a28419-6831-4197-8973-00c5591e19cb", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "from jitterbug.utils import make_calculator\n", + "from itertools import permutations\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": 2, + "id": "52252ee2-315c-48bb-8cba-07620e6e2faa", + "metadata": { + "tags": [ + "parameters" + ] + }, + "outputs": [], + "source": [ + "starting_geometry = '../data/exact/caffeine_pm7_None.xyz'\n", + "threads = min(os.cpu_count(), 12)\n", + "step_size: float = 0.005 # Lambda parameter for an expontential distribution for the Perturbation amount\n", + "perturbs_per_evaluation: int = 1 # Number of perturbations to perform at once" + ] + }, + { + "cell_type": "markdown", + "id": "7010df09-73b2-4d58-be03-15a5f0d04b4c", + "metadata": {}, + "source": [ + "Derived" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "0b6794cd-477f-45a1-b96f-2332804ddb20", + "metadata": {}, + "outputs": [], + "source": [ + "run_name = Path(starting_geometry).name[:-4]\n", + "name, method, basis = run_name.split(\"_\")" + ] + }, + { + "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": 4, + "id": "ad9fd725-b1ba-4fec-ae41-959be0e540b3", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "data": { + "text/plain": [ + "Atoms(symbols='O2N4C8H10', pbc=False, forces=..., calculator=SinglePointCalculator(...))" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "atoms = read(starting_geometry)\n", + "atoms" + ] + }, + { + "cell_type": "markdown", + "id": "2284056b-ddf2-4a3b-88ca-b1c6dc84a2d5", + "metadata": {}, + "source": [ + "## Assemble the List of Perturbations\n", + "We are going to set up all of the perturbations needed for numerical partial second derivatives, which include [perturbing only one coordinate for the diagonal terms and every pair of coordinates for the off-diagonal](https://en.wikipedia.org/wiki/Finite_difference#Multivariate_finite_differences)." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "c50272bb-4c71-4851-8bf8-86514d7e690c", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Collected 144 diagonal terms\n", + "Collected 10368 with off-diagonal terms\n" + ] + } + ], + "source": [ + "n_coords = len(atoms) * 3\n", + "perturbations = [\n", + " (d * i,) for d in [-1, 1] for i in range(1, n_coords + 1) # Start from 1 so we can encode direction with the sign (-0 == 0)\n", + "]\n", + "print(f'Collected {len(perturbations)} diagonal terms')\n", + "perturbations.extend(\n", + " (d * i, d * j) for d in [-1, 1] for i, j in permutations(range(1, n_coords + 1), 2) # Assumes we re-use the data from diagonal terms for the off-diagonal\n", + ")\n", + "print(f'Collected {len(perturbations)} with off-diagonal terms')" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "47932272-de88-49ee-a6cf-e28c619b6dbe", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "assert perturbs_per_evaluation == 1, 'Have not implemented combining them yet.'" + ] + }, + { + "cell_type": "markdown", + "id": "b3c9b069-7fad-434c-8ece-bb5185167d14", + "metadata": {}, + "source": [ + "TODO: Combine them. We must ensure that + and - displacements along the same axis are not mixed and the random order is repeatable" + ] + }, + { + "cell_type": "markdown", + "id": "3dbfc92b-a561-4b8f-9fab-98b36029653e", + "metadata": {}, + "source": [ + "## Run the Perturbations" + ] + }, + { + "cell_type": "markdown", + "id": "ad4c5d8e-96d4-4bb6-9bf2-6474d7563448", + "metadata": {}, + "source": [ + "Prepare the output directory" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "23502eea-0974-4248-8f19-e85447069c61", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "out_dir = Path('data') / 'along-axes'\n", + "out_dir.mkdir(exist_ok=True, parents=True)" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "bf1366fc-d9a7-4a98-b9c9-cb3a0209b406", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Writing to data/along-axes/caffeine_pm7_None_d=5.00e-03-N=1.db\n" + ] + } + ], + "source": [ + "db_path = out_dir / f'{run_name}_d={step_size:.2e}-N={perturbs_per_evaluation}.db'\n", + "print(f'Writing to {db_path}')" + ] + }, + { + "cell_type": "markdown", + "id": "004158dc-3fe9-47a6-99dd-268aa69bb27b", + "metadata": {}, + "source": [ + "Add the relaxed geometry if needed" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "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, name='initial')" + ] + }, + { + "cell_type": "markdown", + "id": "56ebf431-75a0-44d5-8e18-43f2898d6dab", + "metadata": {}, + "source": [ + "Make the calculator" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "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": 11, + "id": "4f9ac37d-6c2f-4ece-bb8c-3d3eb7fbd8a1", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "We have finished 0 perturbations already. 10368 left to do.\n" + ] + } + ], + "source": [ + "with connect(db_path) as db:\n", + " num_done = len(db) - 1\n", + "print(f'We have finished {num_done} perturbations already. {len(perturbations) - num_done} left to do.')" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "a6fa1b33-defc-4b35-895d-052eb64453fb", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "d+71+70: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10368/10368 [05:53<00:00, 29.30it/s]\n" + ] + } + ], + "source": [ + "iterator = tqdm(perturbations[num_done:])\n", + "for perturb in iterator:\n", + " # Create the perturbation vector\n", + " disp = np.zeros((n_coords,))\n", + " for d in perturb:\n", + " disp[abs(d) - 1] = (1 if abs(d) > 0 else -1) * step_size\n", + " disp = disp.reshape((-1, 3))\n", + " \n", + " # Make the new atoms\n", + " new_atoms = atoms.copy()\n", + " new_atoms.positions += disp\n", + " \n", + " # Make the name for the computation\n", + " name = \"d\" + \"\".join(f'{\"+\" if d > 0 else \"-\"}{abs(d)-1}' for d in perturb)\n", + " iterator.set_description(name)\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, name=name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b0850be6-bd52-450e-991a-46a20649c98a", + "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.17" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}