{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2ebf008e",
   "metadata": {},
   "source": [
    "## <span style=\"color:red\">Batchwise structure optimization is deprecated and will be upgraded soon</span>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d100d71c",
   "metadata": {},
   "source": [
    "# Batch-wise Structure Relaxation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd51dfe1",
   "metadata": {},
   "source": [
    "In this tutorial, we show how to use the ``ASEBatchwiseLBFGS``. It enables relaxation of structures in a batch-wise manner, i.e. it optimizes multiple structures in parallel. This is particularly useful, when many relatively similar structures (--> similar time until convergence) should be relaxed while requiring possibly short simulation time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68581ba5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import os\n",
    "# import shutil\n",
    "\n",
    "# import torch\n",
    "# from ase.io import read\n",
    "\n",
    "# import schnetpack as spk\n",
    "# from schnetpack import properties\n",
    "# from schnetpack.interfaces.ase_interface import AtomsConverter\n",
    "# from schnetpack.interfaces.batchwise_optimization import ASEBatchwiseLBFGS, BatchwiseCalculator"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6339e784",
   "metadata": {},
   "source": [
    "First, we load the force field model that provides the forces for the relaxation process. Furthermore, we define the atoms converter, which is used to convert ase Atoms objects to SchNetPack input. Eventually the calculator is initialized. The latter provides the necessary functionality to load a model and calculates forces and energy for the respective structures. Please note that running batchwise relaxations is significantly faster on a cuda device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f1bd733",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model_path = \"../../tests/testdata/md_ethanol.model\"\n",
    "\n",
    "## set device\n",
    "# device = torch.device(\"cpu\")\n",
    "\n",
    "## load model\n",
    "# model = torch.load(model_path, map_location=device, weights_only=False)\n",
    "\n",
    "## define neighbor list\n",
    "# cutoff = model.representation.cutoff.item()\n",
    "# nbh_list=spk.transform.MatScipyNeighborList(cutoff=cutoff)\n",
    "\n",
    "## build atoms converter\n",
    "# atoms_converter = AtomsConverter(\n",
    "#    neighbor_list=nbh_list,\n",
    "#    device=device,\n",
    "# )\n",
    "\n",
    "## build calculator\n",
    "# calculator = BatchwiseCalculator(\n",
    "#    model=model_path,\n",
    "#    atoms_converter=atoms_converter,\n",
    "#    device=device,\n",
    "#    energy_unit=\"kcal/mol\",\n",
    "#    position_unit=\"Ang\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66f87158",
   "metadata": {},
   "source": [
    "Subsequently, we load the batch of initial structures utilizing ASE (supports xyz, db and more)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3ebc6e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# input_structure_file = \"../../tests/testdata/md_ethanol.xyz\"\n",
    "\n",
    "## load initial structures\n",
    "# ats = read(input_structure_file, index=\":\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5323eaa6",
   "metadata": {},
   "source": [
    "For some systems it helps to fix the positions of certain atoms during the relaxation. This can be achieved by providing a mask of boolean entries to ``ASEBatchwiseLBFGS``. The mask is a list of $n_\\text{atoms}$ entries, indicating atoms, which positions are fixed during the relaxation. Here, we do not fix any atoms. Hence, the mask only contains ``True``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28e377f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "## define structure mask for optimization (True for fixed, False for non-fixed)\n",
    "# n_atoms = len(ats[0].get_atomic_numbers())\n",
    "# single_structure_mask = [False for _ in range(n_atoms)]\n",
    "## expand mask by number of input structures (fixed atoms are equivalent for all input structures)\n",
    "# mask = single_structure_mask * len(ats)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "252b22e8",
   "metadata": {},
   "source": [
    "Finally, we run the optimization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2532bb4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# results_dir = \"./howto_batchwise_relaxations_outputs\"\n",
    "# if not os.path.exists(results_dir):\n",
    "#    os.makedirs(results_dir)\n",
    "\n",
    "## Initialize optimizer\n",
    "# optimizer = ASEBatchwiseLBFGS(\n",
    "#    calculator=calculator,\n",
    "#    atoms=ats,\n",
    "#    trajectory=\"./howto_batchwise_relaxations_outputs/relax_traj\",\n",
    "# )\n",
    "\n",
    "## run optimization\n",
    "# optimizer.run(fmax=0.0005, steps=1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb369782",
   "metadata": {},
   "outputs": [],
   "source": [
    "# if os.path.exists(results_dir):\n",
    "#    shutil.rmtree(results_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f579c83",
   "metadata": {},
   "source": [
    "Optimzed structures (in the form of ASE `Atoms`) and properties can be obtained with the `get_relaxation_results` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78f81235",
   "metadata": {},
   "outputs": [],
   "source": [
    "## get list of optimized structures and properties\n",
    "# opt_atoms, opt_props = optimizer.get_relaxation_results()\n",
    "\n",
    "# for oatoms in opt_atoms:\n",
    "#    print(oatoms.get_positions())\n",
    "\n",
    "# print(opt_props)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11776e3c",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}