Ep/ensemble calc

.gitignore

@@ -126,3 +126,4 @@ interfaces/lammps/examples/*/deployed_model
 
 # batchwise optimizer examples
 examples/howtos/howto_batchwise_relaxations_outputs/*
+examples/howtos/howto_ensemble_relaxations_outputs/*

examples/howtos/01_howto_ensemble_calculations.ipynb

@@ -0,0 +1,144 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e16f3450",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import torch\n",
+    "import random\n",
+    "import shutil\n",
+    "import time\n",
+    "import numpy as np\n",
+    "from tqdm import tqdm \n",
+    "from matplotlib import pyplot as plt\n",
+    "from typing import Optional, List, Union, Dict\n",
+    "\n",
+    "import ase\n",
+    "from ase.io import read, write\n",
+    "from ase import Atoms\n",
+    "from ase.optimize.lbfgs import LBFGS\n",
+    "from ase.optimize import QuasiNewton\n",
+    "\n",
+    "import schnetpack as spk\n",
+    "from schnetpack import properties\n",
+    "from schnetpack.interfaces.ase_interface import SpkEnsembleCalculator, SpkCalculator\n",
+    "from schnetpack.interfaces.batchwise_optimization import ASEBatchwiseLBFGS, BatchwiseCalculator"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ea8ec623",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# TODO: Markdown comments, uncertainty evaluation example"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "58a19166",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_path_0 = \"/home/jonas/Documents/schnetpack/tests/testdata/md_ethanol.model\"\n",
+    "model_path_1 = \"/home/jonas/Documents/schnetpack/tests/testdata/md_ethanol.model\"\n",
+    "\n",
+    "# set device\n",
+    "device = torch.device(\"cuda\")\n",
+    "\n",
+    "# define neighbor list\n",
+    "cutoff = 5.0\n",
+    "nbh_list=spk.transform.MatScipyNeighborList(cutoff=cutoff)\n",
+    "\n",
+    "# Set up calculator\n",
+    "calculator = SpkEnsembleCalculator(\n",
+    "    model=[model_path_0, model_path_1],\n",
+    "    neighbor_list=nbh_list,\n",
+    "    energy_unit=\"kcal/mol\",\n",
+    "    position_unit=\"Angstrom\",\n",
+    "    device=device,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "34d4bce9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "input_structure_file = \"../../tests/testdata/md_ethanol.xyz\"\n",
+    "random.seed(42)\n",
+    "\n",
+    "# Setup directory\n",
+    "if not os.path.exists('howto_ensemble_relaxations_outputs'):\n",
+    "    os.makedirs('howto_ensemble_relaxations_outputs')\n",
+    "working_dir = \"howto_ensemble_relaxations_outputs/relax\"\n",
+    "if os.path.exists(working_dir):\n",
+    "    shutil.rmtree(working_dir)\n",
+    "os.makedirs(working_dir)\n",
+    "\n",
+    "# load initial structure\n",
+    "mol = read(input_structure_file)\n",
+    "pos = mol.get_positions()\n",
+    "# distort the structure\n",
+    "for n in range(pos.shape[0]):\n",
+    "    pos[n] = pos[n] * random.uniform(0.95,1.05)\n",
+    "molecule = Atoms(positions=pos, numbers=mol.get_atomic_numbers())\n",
+    "molecule.calc = calculator"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c6ebb93d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# run relaxation\n",
+    "optimize_file = os.path.join(working_dir, \"optimization\")\n",
+    "optimizer = LBFGS(\n",
+    "    molecule,\n",
+    "    trajectory=\"{:s}.traj\".format(optimize_file),\n",
+    "    restart=\"{:s}.pkl\".format(optimize_file),\n",
+    ")\n",
+    "optimizer.run(fmax=0.001, steps=1000)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "790e15aa",
+   "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.8.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}