Skip to content

Commit

Permalink
Start notebook for displacing along axes
Browse files Browse the repository at this point in the history
  • Loading branch information
@WardLT
WardLT committed Oct 4, 2023
1 parent 93dc2ab commit 12a6e2d
Showing 1 changed file with 360 additions and 0 deletions.
360 changes: 360 additions & 0 deletions notebooks/1_explore-sampling-methods/2_displace-along-axes.ipynb
View file
Original file line number Diff line number Diff line change
@@ -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
}

0 comments on commit 12a6e2d

Please sign in to comment.