{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "# Parameter transformation for the Stillinger-Weber potential\n",
    "\n",
    "Parameters in the empirical interatomic potential are often restricted by some physical\n",
    "constraints. As an example, in the Stillinger-Weber (SW) potential, the energy scaling\n",
    "parameters (e.g., ``A`` and ``B``) and the length scaling parameters (e.g., ``sigma`` and\n",
    "``gamma``) are constrained to be positive.\n",
    "\n",
    "Due to these constraints, we might want to work with the log of the parameters, i.e.,\n",
    "``log(A)``, ``log(B)``, ``log(sigma)``, and ``log(gamma)`` when doing the optimization.\n",
    "After the optimization, we can transform them back to the original parameter space using\n",
    "an exponential function, which will guarantee the positiveness of the parameters.\n",
    "\n",
    "In this tutorial, we show how to apply parameter transformation to the SW potential for\n",
    "silicon that is archived on OpenKIM_. Compare this with [Si_training_set.tar.gz](https://raw.githubusercontent.com/openkim/kliff/master/examples/Si_training_set.tar.gz).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To start, let's first install the SW model::\n\n   $ kim-api-collections-management install user SW_StillingerWeber_1985_Si__MO_405512056662_006\n\n.. seealso::\n   This installs the model and its driver into the ``User Collection``. See\n   `install_model` for more information about installing KIM models.\n\nThis is\n\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from kliff.legacy.calculators import Calculator\n",
    "from kliff.dataset import Dataset\n",
    "from kliff.dataset.weight import Weight\n",
    "from kliff.legacy.loss import Loss\n",
    "from kliff.models import KIMModel\n",
    "from kliff.models.parameter import Parameter\n",
    "from kliff.transforms.parameter_transforms import LogParameterTransform\n",
    "from kliff.utils import download_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before creating a KIM model for the SW potential, we first instantiate the parameter\ntransformation class that we want to use. ``kliff`` has a built-in log-transformation;\nhowever, extending it to other parameter transformation can be done by creating a\nsubclass of :class:`~kliff.models.parameter_transform.ParameterTransform`.\n\nTo make a direct comparison to `tut_kim_sw`, in this tutorial we will apply\nlog-transformation to parameters ``A``, ``B``, ``sigma``, and ``gamma``, which\ncorrespond to energy and length scales.\n\n\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "transform = LogParameterTransform(param_names=[\"A\", \"B\", \"sigma\", \"gamma\"])\nmodel = KIMModel(\n    model_name=\"SW_StillingerWeber_1985_Si__MO_405512056662_006\",\n    params_transform=transform,\n)\nmodel.echo_model_params(params_space=\"original\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "``model.echo_model_params(params_space=\"original\")`` above will print out parameter\nvalues in the original, untransformed space, i.e., the original parameterization of\nthe model. If we supply the argument ``params_space=\"transformed\"``, then the printed\nparameter values are given in the transformed space, e.g., log space (below). The\nvalues of the other parameters are not changed.\n\n\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "model.echo_model_params(params_space=\"original\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare the output of ``params_space=\"transformed\"`` and ``params_space=\"original\"``,\nyou can see that the values of ``A``, ``B``, ``sigma``, and ``gamma`` are in the log\nspace after the transformation.\n\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we will set up the initial guess of the parameters to optimize. A value of\n``\"default\"`` means the initial guess will be directly taken from the value already in\nthe model.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>The parameter values we initialize, as well as the lower and upper bounds, are in\n   transformed space (i.e. log space here).</p></div>\n\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "model.set_opt_params(\n    A=[[np.log(5.0), np.log(1.0), np.log(20)]],\n    B=[[\"default\"]],\n    sigma=[[np.log(2.0951), \"fix\"]],\n    gamma=[[np.log(1.5)]],\n)\nmodel.echo_opt_params()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can show the parameters we’ve just set by ``model.echo_opt_params()``.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>``model.echo_opt_params()`` always displays the parameter values in the transformed\n   space. And it only shows all the parameters specified to optimize. To show all\n   the parameters, do ``model.echo_model_params(params_space=\"transformed\")``.</p></div>\n\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once we set the model and the parameter transformation scheme, then further\ncalculations, e.g., training the model, will be performed using the transformed space\nand can be done in the same way as in `tut_kim_sw`.\n\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Training set\ndataset_path = download_dataset(dataset_name=\"Si_training_set\")\nweight = Weight(energy_weight=1.0, forces_weight=0.1)\ntset = Dataset(dataset_path, weight)\nconfigs = tset.get_configs()\n\n# Calculator\ncalc = Calculator(model)\n_ = calc.create(configs)\n\n# Loss function and model training\nsteps = 100\nloss = Loss(calc, nprocs=2)\nloss.minimize(method=\"L-BFGS-B\", options={\"disp\": True, \"maxiter\": steps})\n\nmodel.echo_model_params(params_space=\"original\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The optimized parameter values from this model training are very close, if not the\nsame, as in `tut_kim_sw`. This is expected for the simple tutorial example\nconsidered. But for more complex models, training in a transformed space can make it\nmuch easier for the optimizer to navigate the parameter space.\n\n\n\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}