{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Jscatter: Lattices\n",
    "\n",
    "This Notebook aims to give a short introduction about usage of lattices also for 2D fitting of scattering images.\n",
    "\n",
    "### Content\n",
    "- Ewald sphere of a simple cubic lattice\n",
    "- 2D Fitting: 2D pattern of a lattice\n",
    "\n",
    "![Jscatter](https://jscatter.readthedocs.io/en/latest/_images/Jscatter.jpeg \"Jscatter\")\n",
    "\n",
    "#### Prerequisite "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "!{sys.executable} -m pip install jscatter\n",
    "%matplotlib notebook\n",
    "\n",
    "import jscatter as js\n",
    "import numpy as np\n",
    "js.usempl(True)   # force matplotlib, not needed in Linux"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ewald sphere of a simple cubic lattice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jscatter as js\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib.pyplot as pyplot\n",
    "from matplotlib import cm, colors\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "\n",
    "phi, theta = np.meshgrid(np.r_[0:np.pi:45j], np.r_[0:2 * np.pi:90j])\n",
    "\n",
    "# The Cartesian coordinates of the Ewald sphere\n",
    "q = 3\n",
    "x = q * (np.sin(phi) * np.cos(theta) + 1)\n",
    "y = q * np.sin(phi) * np.sin(theta)\n",
    "z = q * np.cos(phi)\n",
    "qxzy = np.asarray([x, y, z]).reshape(3, -1).T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create basic lattice "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sclattice = js.lattice.scLattice(2.1, 5)\n",
    "sclattice.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate the scattering pattern of an oriented sclattice and plot the corresponding Ewald sphere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ffe = js.ff.orientedCloudScattering(qxzy, sclattice.XYZ, coneangle=5, nCone=20, rms=0.02)\n",
    "\n",
    "# log scale for colors\n",
    "ffey = np.log(ffe.Y)\n",
    "fmax, fmin = ffey.max(), ffey.min()\n",
    "ffeY = (np.reshape(ffey, x.shape) - fmin) / (fmax - fmin)\n",
    "\n",
    "# Set the aspect ratio to 1 so our sphere looks spherical\n",
    "fig = pyplot.figure(figsize=pyplot.figaspect(1.))\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "ax.plot_surface(x, y, z, rstride=1, cstride=1, facecolors=cm.seismic(ffeY))\n",
    "fig.suptitle('Ewald sphere of simple cubic lattice')\n",
    "\n",
    "# Turn off the axis planes\n",
    "ax.set_axis_off()\n",
    "pyplot.show(block=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Show Ewald sphere and scattering plane (at detector) of a  simple cubic lattice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "phi, theta = np.meshgrid(np.r_[0:np.pi:90j], np.r_[0:1 * np.pi:90j])\n",
    "# The Cartesian coordinates of the unit sphere\n",
    "q = 8\n",
    "x = q * (np.sin(phi) * np.cos(theta) + 1)\n",
    "y = q * np.sin(phi) * np.sin(theta)\n",
    "z = q * np.cos(phi)\n",
    "qxzy = np.asarray([x, y, z]).reshape(3, -1).T\n",
    "\n",
    "# Set the aspect ratio to 1 so our sphere looks spherical\n",
    "fig = pyplot.figure(figsize=pyplot.figaspect(1.))\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "\n",
    "# create lattice and show it scaled\n",
    "sclattice = js.lattice.scLattice(2.1, 1)\n",
    "sclattice.rotatehkl2Vector([1, 1, 1], [0, 0, 1])\n",
    "gg = ax.scatter(sclattice.X / 3 + q, sclattice.Y / 3, sclattice.Z / 3, c='k', alpha=0.9)\n",
    "gg.set_visible(True)\n",
    "\n",
    "ds = 15\n",
    "fpi = np.pi / 180.\n",
    "ffs = js.sf.orientedLatticeStructureFactor(qxzy, sclattice, rotation=None, domainsize=ds, rmsd=0.1, hklmax=2, nGauss=23)\n",
    "# show scattering in log scale on Ewald sphere\n",
    "ffsy = np.log(ffs.Y)\n",
    "fmax, fmin = ffsy.max(), ffsy.min()\n",
    "ffsY = (np.reshape(ffsy, x.shape) - fmin) / (fmax - fmin)\n",
    "\n",
    "ax.plot_surface(x, y, z, rstride=1, cstride=1, facecolors=cm.hsv(ffsY), alpha=0.4)\n",
    "ax.plot([q, 2 * q], [0, 0], [0, 0], color='k')\n",
    "pyplot.show(block=False)\n",
    "\n",
    "#  ad flat detector xy plane\n",
    "xzw = np.mgrid[-8:8:80j, -8:8:80j]\n",
    "qxzw = np.stack([np.zeros_like(xzw[0]), xzw[0], xzw[1]], axis=0)\n",
    "\n",
    "ff2 = js.sf.orientedLatticeStructureFactor(qxzw.reshape(3, -1).T, sclattice, rotation=None, domainsize=ds, rmsd=0.1,\n",
    "                                           hklmax=2, nGauss=23)\n",
    "ffs2 = np.log(ff2.Y)\n",
    "fmax, fmin = ffs2.max(), ffs2.min()\n",
    "ff2Y = (np.reshape(ffs2, xzw[0].shape) - fmin) / (fmax - fmin)\n",
    "ax.plot_surface(qxzw[0], qxzw[1], qxzw[2], rstride=1, cstride=1, facecolors=cm.gist_ncar(ff2Y), alpha=0.3)\n",
    "ax.set_xlabel('x axis')\n",
    "ax.set_ylabel('y axis')\n",
    "ax.set_zlabel('z axis')\n",
    "\n",
    "fig.suptitle('Scattering plane and Ewald sphere of lattice \\n matplotlib is limited in 3D')\n",
    "pyplot.show(block=False)\n",
    "# matplotlib is not real 3D rendering\n",
    "# maybe use later something else"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2D Fitting: 2D pattern of a lattice\n",
    "The loaded image is similar to SAXS images from ordered particles.\n",
    "\n",
    "Looking in 3D along xy lane one can see at high q the plateau due to diffusive scattering which vanishes at low q. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jscatter as js\n",
    "import numpy as np\n",
    "\n",
    "# load a image with hexagonal peaks (synthetic data) and set center\n",
    "image = js.sas.sasImage(js.examples.datapath + '/2Dhex.tiff')\n",
    "image.setPlaneCenter([51, 51])\n",
    "# transform to dataarray with X,Z as wavevectors\n",
    "# The fit algorithm works also with masked areas\n",
    "hexa = image.asdataArray(masked=0)\n",
    "\n",
    "fig = js.mpl.surface(hexa.X, hexa.Z, hexa.Y,shape=image.shape)\n",
    "fig.suptitle('original')\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def latticemodel(qx, qz, R, ds, rmsd, hklmax, bgr, I0):\n",
    "    # a hexagonal grid with background, domain size and Debye Waller-factor (rmsd)\n",
    "    # 3D wavevector\n",
    "    qxyz = np.c_[qx, qz, np.zeros_like(qx)]\n",
    "    # define lattice\n",
    "    grid = js.sf.hcpLattice(R, [3, 3, 3])\n",
    "    # here one may rotate the crystal\n",
    "\n",
    "    # calc scattering\n",
    "    lattice = js.sf.orientedLatticeStructureFactor(qxyz, grid, domainsize=ds, rmsd=rmsd, hklmax=hklmax)\n",
    "    # add I0 and background\n",
    "    lattice.Y = I0 * lattice.Y + bgr\n",
    "    return lattice"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### First *differential evolution* then polish\n",
    "Because of the high plateau in the chi2 landscape we first need to use a algorithm finding a global minimum. The 'differential evolution'needs around 2300 evaluations and some time.\n",
    "\n",
    "Alternatively one may find in other ways good starting parameters. in this case one should first evaluate in 1D the peak positions and fit later the orientation. As a third step one can fit all together. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if 0:\n",
    "    # Please do this only if you have time to wait\n",
    "    hexa.fit(latticemodel, {'R': 2, 'ds': 10, 'rmsd': 0.1, 'bgr': 1, 'I0': 10},\n",
    "             {'hklmax': 4, }, {'qx': 'X', 'qz': 'Z'}, method='differential_evolution')\n",
    "    #\n",
    "    fig = js.mpl.surface(hexa.X, hexa.Z, hexa.Y, shape=image.shape)\n",
    "    fig = js.mpl.surface(hexa.lastfit.X, hexa.lastfit.Z, hexa.lastfit.Y)\n",
    "    \n",
    "# We use as starting parameters the result of the previous fit\n",
    "# We use LevenbergMarquardt algorithm to polish the result\n",
    "hexa.fit(latticemodel, {'R': 3.02, 'ds': 8.7, 'rmsd': 0.193, 'bgr': 3, 'I0': 31},\n",
    "         {'hklmax': 4, }, {'qx': 'X', 'qz': 'Z'},output=0)\n",
    "\n",
    "\n",
    "fig2 = js.mpl.surface(hexa.lastfit.X, hexa.lastfit.Z, hexa.lastfit.Y ,shape=image.shape)\n",
    "fig2.suptitle('fit result')\n",
    "fig2.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## For more examples please visit the Jscatter documentation."
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
