04-SrFit-two-datasets.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SrFit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The diffpy.srfit package provides a top-level framework for setting up multi-component fits.  The SrFit recipes tune parameters in the optimized complexes and link them to refinable variables, which can then be tied to functional constraints or restraints or fixed at a constant values.\n",
    "\n",
    "This notebook will demonstrate these features by setting up a SrFit recipe to manage a two-component fit.  To start with we'll generate a single data stream with some noisy linear data and pass it to a *Profile* object, a SrFit class for representing observed data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "import numpy as np\n",
    "A = 3\n",
    "B = 5\n",
    "noise_level = 3.0\n",
    "xobs0 = np.linspace(0, 10.0, 21)\n",
    "yobs0 = A * xobs0 + B + noise_level * np.random.randn(xobs0.size)\n",
    "dyobs0 = noise_level * np.ones_like(xobs0)\n",
    "from diffpy.srfit.fitbase import Profile\n",
    "linedata = Profile()\n",
    "linedata.setObservedProfile(xobs0, yobs0, dyobs0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next step is to create a fit contribution object, which will contain both the measured data and a\n",
    "function calculator to simulate the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffpy.srfit.fitbase import FitContribution\n",
    "linefit = FitContribution('linefit')\n",
    "linefit.setProfile(linedata)\n",
    "linefit.setEquation(\"A * x + B\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By defining the function calculator with the setEquation method, SrFit automatically recognizes A and B\n",
    "as parameters that will be used to model the data. We now create a fit recipe to hold our fit contribution\n",
    "and tell the recipe that A and B are the variables we want to refine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffpy.srfit.fitbase import FitRecipe\n",
    "recipe = FitRecipe()\n",
    "recipe.addContribution(linefit);\n",
    "recipe.addVar(recipe.linefit.A);\n",
    "recipe.addVar(recipe.linefit.B);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once defined, the fit recipe acts like a simple python function that takes the values of the variables as input \n",
    "and returns the fit residual as output. Thus, it can be plugged into a variety of available \n",
    "optimization programs. We'll use the least squares optimizer from the scipy.optimize module with some\n",
    "reasonable starting values for A and B."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "recipe.A =  2.8407684402216278\n",
      "recipe.B =  4.905964647496833\n"
     ]
    }
   ],
   "source": [
    "from scipy.optimize import leastsq\n",
    "recipe.A = 1.0\n",
    "recipe.B = 2.0\n",
    "recipe.clearFitHooks()\n",
    "leastsq(recipe.residual, recipe.values)\n",
    "\n",
    "print(\"recipe.A = \", recipe.A.value)\n",
    "print(\"recipe.B = \", recipe.B.value)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The fit converged and gave us values of A and B that are somewhat close to our original starting values. Here's\n",
    "where things get interesting. Suppose we have a second data set and model that also depend on A and B. We can\n",
    "add another fit contribution to the recipe and run a co-refinement. Let's now generate more data--this time\n",
    "we'll use a a quadratic form and, just for fun, we'll add another parameter, C."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "C = 8.0\n",
    "xobs1 = np.linspace(-3.0, 3.0, 31)\n",
    "yobs1 = A * xobs1 ** 2 + B * xobs1 + C + noise_level * np.random.randn(xobs1.size)\n",
    "dyobs1 = noise_level * np.ones_like(xobs0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As before we create a fit contribution object to hold our measured data and the corresponding model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "quadfit = FitContribution('quadfit')\n",
    "quaddata = Profile()\n",
    "quaddata.setObservedProfile(xobs1, yobs1)\n",
    "quadfit.setProfile(quaddata)\n",
    "quadfit.setEquation(\"A * x**2 + B * x + C\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, creating a co-refinement is as simple as adding this second fit contribution to our existing recipe and\n",
    "constraining the parameters A and B in our new equation to the existing variables A and B in our recipe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "recipe.addContribution(quadfit)\n",
    "recipe.constrain(recipe.quadfit.A, recipe.A)\n",
    "recipe.constrain(recipe.quadfit.B, recipe.B)\n",
    "recipe.addVar(recipe.quadfit.C);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once again we use scipy's least squares optimization engine to fit our data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "recipe.A =  2.8529436829871115\n",
      "recipe.B =  5.03015354415433\n",
      "recipe.C =  7.859159101901894\n"
     ]
    }
   ],
   "source": [
    "recipe.C = 1.0\n",
    "leastsq(recipe.residual, recipe.values)\n",
    "print(\"recipe.A = \", recipe.A.value)\n",
    "print(\"recipe.B = \", recipe.B.value)\n",
    "print(\"recipe.C = \", recipe.C.value)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the added data we achieve a better agreement between our original parameters and the fit to the model."
   ]
  }
 ],
 "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}