diff --git a/notebooks/1_explore-sampling-methods/0_random-directions-same-distance.ipynb b/notebooks/1_explore-sampling-methods/0_random-directions-same-distance.ipynb index 78fba34..ea8d7b8 100644 --- a/notebooks/1_explore-sampling-methods/0_random-directions-same-distance.ipynb +++ b/notebooks/1_explore-sampling-methods/0_random-directions-same-distance.ipynb @@ -272,7 +272,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.17" + "version": "3.9.18" } }, "nbformat": 4, diff --git a/notebooks/1_explore-sampling-methods/3_displace-along-vibrational-modes.ipynb b/notebooks/1_explore-sampling-methods/3_displace-along-vibrational-modes.ipynb new file mode 100644 index 0000000..097b41f --- /dev/null +++ b/notebooks/1_explore-sampling-methods/3_displace-along-vibrational-modes.ipynb @@ -0,0 +1,434 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "8233267b-e98b-44be-b9aa-116d0e67a94b", + "metadata": {}, + "source": [ + "# Compute Energies of Displacements Along Vibrational Models\n", + "Compute the vibrational modes using a lower level of theory then displace along those axes. Options include:\n", + "- Energy scale of vibrations. Set the anticipated energy increase based on eigenvalue of the Hessians\n", + "- Number of modes to vibrate at a time.\n", + "- Maximum length of displacement vector." + ] + }, + { + "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.vibrations import Vibrations\n", + "from ase.io import write, read\n", + "from ase.db import connect\n", + "from ase import Atoms, units\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": "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.002 # Target energy increase (units: eV)\n", + "perturbs_per_evaluation: int = 16 # Number of perturbations to perform at once\n", + "max_perturb = 0.06 # Maximum length of displacement vector\n", + "lower_level: tuple[str, str] = ('xtb', None)" + ] + }, + { + "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": [ + "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": 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": [ + "## Assemble the List of Perturbations\n", + "Start by computing the vibrational modes using a lower level of theory, then pick the vibrational modes that are large enough" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "218adae4-a2dc-4bdb-846b-5f2cab34320f", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "lower_calc = make_calculator(*lower_level)" + ] + }, + { + "cell_type": "markdown", + "id": "bd0db689-6f9e-4753-a866-8bca09aee4de", + "metadata": {}, + "source": [ + "Compute the vibrational modes" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "08fd5ca2-aacc-4f10-8b6b-2246dd78e59b", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "%%time\n", + "atoms.calc = lower_calc\n", + "vib = Vibrations(atoms)\n", + "vib.run()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "20c58921-bf90-47a2-b6c5-0b4f778312c9", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "vib_data = vib.get_vibrations()" + ] + }, + { + "cell_type": "markdown", + "id": "c618cf41-946e-44d5-8c09-a4c80f504256", + "metadata": {}, + "source": [ + "Get the eigenvalues and eigenvectors." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "30c21da4-5733-4b0e-9ebb-2ec60b861f09", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "evalues, emodes = np.linalg.eigh(vib_data.get_hessian_2d())" + ] + }, + { + "cell_type": "markdown", + "id": "78563dd5-c6da-46d8-9b99-4928187023da", + "metadata": { + "tags": [] + }, + "source": [ + "Scale the eigenmodes such that a perturbation of length 1 should increase the energy by the target `step_size`.\n", + "\n", + "The value of the eigenmode is the a second derivative wrt atom positions. So, $\\Delta E = 0.5 \\lambda_i \\delta x_i^2$ and $x_i = \\sqrt{2 \\Delta E / \\lambda_i}$\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "715822cf-4ced-488a-9894-fda13aefde87", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "scale_mag = np.clip(np.sign(evalues) * np.sqrt(2 * step_size / np.abs(evalues)),\n", + " a_min=-max_perturb,\n", + " a_max=max_perturb)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "54a2a550-7920-4477-815c-fdabfc84c913", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "scaled_modes = emodes * scale_mag[:, None]" + ] + }, + { + "cell_type": "markdown", + "id": "38c876d2-9571-4c14-80df-d47f3d067a2e", + "metadata": {}, + "source": [ + "Check my math" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9a030307-951f-49fe-a9c5-aedac301f3df", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "perturbed_atoms = atoms.copy()\n", + "perturbed_atoms.positions += scaled_modes[0, :].reshape((-1, 3))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1458cdeb-f734-4c30-8cca-6e1fa7e3e5a8", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "perturbed_energy = lower_calc.get_potential_energy(perturbed_atoms) - lower_calc.get_potential_energy(atoms)" + ] + }, + { + "cell_type": "markdown", + "id": "3dbfc92b-a561-4b8f-9fab-98b36029653e", + "metadata": {}, + "source": [ + "## Run the Perturbations\n", + "Run as many as we should need to fill all degrees of freedom for the Hessian." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "db25d8b1-112a-4ba9-b4f4-850f13796132", + "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": "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') / 'along-vibrational-modes'\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}-N={perturbs_per_evaluation}-maxstep={max_perturb:.2e}-lower={\"+\".join(map(str, lower_level))}.db'\n", + "print(f'Writing to {db_path}')" + ] + }, + { + "cell_type": "markdown", + "id": "39477a6b-c8ba-4ef7-bb66-40f80b2603d2", + "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)\n", + "atoms.calc = calc" + ] + }, + { + "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", + " atoms.get_potential_energy()\n", + " with connect(db_path) as db:\n", + " db.write(atoms, name='initial')" + ] + }, + { + "cell_type": "markdown", + "id": "8e9e5ff2-3728-459b-b3d3-09acba0f71bc", + "metadata": {}, + "source": [ + "Generate the energies" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4f9ac37d-6c2f-4ece-bb8c-3d3eb7fbd8a1", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "with connect(db_path) as db:\n", + " num_done = len(db) - 1\n", + "print(f'We have finished {num_done} perturbations already. {to_compute - num_done} left to do.')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a6fa1b33-defc-4b35-895d-052eb64453fb", + "metadata": {}, + "outputs": [], + "source": [ + "rng = np.random.RandomState(1)\n", + "iterator = tqdm(range(to_compute - num_done))\n", + "for i in iterator:\n", + " # Choose a number of perturbation vectors\n", + " to_disp = rng.choice(scaled_modes.shape[0], size=(perturbs_per_evaluation,))\n", + " \n", + " # Pick a random magnitude for each\n", + " disp = (scaled_modes[to_disp, :] * rng.uniform(-1, 1, size=(perturbs_per_evaluation, 1))).sum(axis=0)\n", + " \n", + " # Make the new atoms\n", + " new_atoms = atoms.copy()\n", + " new_atoms.positions += disp.reshape((-1, 3))\n", + " \n", + " # Make the name for the computation\n", + " name = f\"perturb-{i}_modes-{','.join(map(str, to_disp))}\"\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.18" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/1_explore-sampling-methods/run-all-methods.sh b/notebooks/1_explore-sampling-methods/run-all-methods.sh index 311dcdc..55b6c58 100644 --- a/notebooks/1_explore-sampling-methods/run-all-methods.sh +++ b/notebooks/1_explore-sampling-methods/run-all-methods.sh @@ -13,3 +13,11 @@ for step_size in 0.02; do papermill -p starting_geometry $xyz -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 + for n in 8 16 32; do + papermill -p starting_geometry $xyz -p perturbs_per_evaluation $n -p step_size $step_size $notebook last.ipynb + done +done