doc/inequality_constraints.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inequality constraints with *refnx*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Simple equality constraints can use the mechanisms outlined in this notebook, but are better expressed using the `Parameter.constraint` mechanism, or by sharing `Parameter` objects. It is sometimes also possible to implement different parameterisation of the model to use physically relevant values.\n",
    "\n",
    "The following processes can be used to make inequality constraints with *refnx*. The dataset is reflectivity from a clean silicon wafer with a native oxide layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import os.path\n",
    "import numpy as np\n",
    "\n",
    "import refnx\n",
    "from refnx.dataset import ReflectDataset\n",
    "from refnx.reflect import SLD, MaterialSLD, ReflectModel\n",
    "from refnx.analysis import Objective, CurveFitter\n",
    "\n",
    "np.random.seed(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pth = os.path.dirname(refnx.__file__)\n",
    "DATASET_NAME = 'c_PLP0000708.dat'\n",
    "file_path = os.path.join(pth, 'dataset', 'test', DATASET_NAME)\n",
    "\n",
    "data = ReflectDataset(file_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "air = SLD(0)\n",
    "sio2 = MaterialSLD('SiO2', 2.2)\n",
    "si = MaterialSLD('Si', 2.33)\n",
    "s = air | sio2(15, 3) | si(0, 3)\n",
    "\n",
    "model = ReflectModel(s, bkg=3e-8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.bkg.setp(vary=True, bounds=(0, 1e-6))\n",
    "# model.scale.setp(vary=True, bounds=(0.9, 1.1))\n",
    "\n",
    "# sio2 layer\n",
    "s[1].rough.setp(vary=True, bounds=(0, 10))\n",
    "s[1].thick.setp(vary=True, bounds=(0, 20))\n",
    "\n",
    "# si/sio2 interface\n",
    "s[-1].rough.setp(vary=True, bounds=(0, 10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "objective = Objective(model, data)\n",
    "fitter = CurveFitter(objective)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fitter.fit('differential_evolution')\n",
    "print(s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inequality constraints with `differential_evolution`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Simple equality constraints can use the following mechanism, but are better expressed using the `Parameter.constraint` mechanism, or by sharing `Parameter` objects. It is sometimes also possible to implement different parameterisation of the model to use physically relevant values.*\n",
    "\n",
    "We see that the thickness of the SiO2 layer is 12.45 and the roughness of the air/SiO2 interface is 4.77. Let's make a constraint that the roughness can't be more than a quarter of the layer thickness. In optimisation such constraints are expressed as inequalities:\n",
    "\n",
    "\\begin{equation*}\n",
    "t > 4\\sigma\n",
    "\\end{equation*}\n",
    "\n",
    "We need to rearrange so that all variables are on one side, we do the rearrangement like this so there is no divide by 0:\n",
    "\n",
    "\\begin{equation*}\n",
    "t - 4\\sigma > 0\n",
    "\\end{equation*}\n",
    "\n",
    "Now we create a callable object (has the `__call__` magic method) that encodes this inequality. We're going to create the object with the parameters we want to constrain (`pars`), so we can refer to them later. We'll also store the objective because we'll need to update it with the fitting parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DEC(object):\n",
    "    def __init__(self, pars, objective):\n",
    "        # we'll store the parameters and objective in this object\n",
    "        # this will be necessary for pickling in the future\n",
    "        self.pars = pars\n",
    "        self.objective = objective\n",
    "\n",
    "    def __call__(self, x):\n",
    "        # we need to update the varying parameters in the\n",
    "        # objective first\n",
    "        self.objective.setp(x)\n",
    "        return float(self.pars[0] - 4*self.pars[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now lets create an instance of that object, using the parameters we want to constrain. Following that we set up a `scipy.optimize.NonlinearConstraint` for use with `differential_evolution`. Note that we want the constraint calculation to be greater than 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pars = (s[1].thick, s[1].rough)\n",
    "dec = DEC(pars, objective)\n",
    "\n",
    "from scipy.optimize import NonlinearConstraint\n",
    "constraint = NonlinearConstraint(dec, 0, np.inf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now do the fit with the added constraint. Note that you can have more than one constraint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(1)\n",
    "fitter.fit('differential_evolution', constraints=(constraint,))\n",
    "print(s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inequality constraints during MCMC sampling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we want to implement that inequality constraint during sampling we can add an extra log-probability term to the `Objective`. This log-probability term will return 0 if the inequality is satisfied, but `-np.inf` if not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LogpExtra(object):\n",
    "    def __init__(self, pars):\n",
    "        # we'll store the parameters and objective in this object\n",
    "        # this will be necessary for pickling in the future\n",
    "        self.pars = pars\n",
    "\n",
    "    def __call__(self, model, data):\n",
    "        if float(self.pars[0] - 4*self.pars[1]) > 0:\n",
    "            return 0\n",
    "        return -np.inf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lpe = LogpExtra(pars)\n",
    "\n",
    "# set the log_extra attribute of the Objective with our extra log-probability term.\n",
    "objective.logp_extra = lpe"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets check what happens to the probabilities with the specified inequality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(s[1].thick)\n",
    "print(s[1].rough)\n",
    "print(objective.logpost())\n",
    "\n",
    "print(\"Now exceed the inequality\")\n",
    "s[1].rough.value = 5.\n",
    "print(objective.logpost())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's MCMC sample the system. There will be a user warning because some walkers have initial starting points which disobey the inequality. Normally one would sample for a far longer time, and thin more appropriately. However, the purpose of the following is to produce a corner plot that demonstrates the inequality - note the sharp dropoff in the probability distribution for the roughness. The roughness doesn't like to go much higher than ~2.5, which is around a quarter of the optimal layer thickness of ~10."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s[1].rough.value = 2.\n",
    "fitter.initialise('covar')\n",
    "fitter.sample(20, nthin=10, pool=1);\n",
    "objective.corner();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}