{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This test file is an example of a script which uses the generic scattering calculator functionality.\r\n",
    "\r\n",
    "The generic scattering calculator can create orientation averages for purely nuclear systems and ideally\r\n",
    "a script should be able to recreate the same results - and extend the idea to magnetic systems.\r\n",
    "\r\n",
    "The default sld data in the calculator describes a 60x60x60 angstrom cube of 10x10x10 pixels, each with\r\n",
    "sld 6.97e-06 angstrom^-2. This fits very well in the fitting tool to a parallelepiped model.\r\n",
    "\r\n",
    "This script intends to recreate this data using an average of multiple calculations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First of all make the required imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\r\n",
    "\r\n",
    "import matplotlib.pyplot as plt\r\n",
    "import numpy as np\r\n",
    "from IPython.display import clear_output\r\n",
    "from scipy.spatial.transform import Rotation\r\n",
    "\r\n",
    "from sas.sascalc.calculator import sas_gen\r"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we need to set some constants. We want to define the size of the sample - a cube with n pixels in each direction, and a fixed stepsize between each node, as well as the discretisation of rotation space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "STEPSIZE=6\r\n",
    "NODES=10\r\n",
    "ANGLESTEP = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now define a function to create a discretised version of 3D rotation space.\r\n",
    "\r\n",
    "3D rotation space can be defined by three parameters. If the rotation takes the cartesian basis\r\n",
    "{x,y,z} to {x',y',z'} then we define the euler angles:\r\n",
    "\r\n",
    "- $\\phi$: a rotation about the z axis\r\n",
    "- $\\theta$: a rotation about the new y axis\r\n",
    "- $\\chi$: a rotation about the new z axis\r\n",
    "\r\n",
    "That is to say that $\\theta$ and $\\phi$ define the direction of the z' axis in polar coordinates from the {x,y,z} basis, and $\\chi$ \r\n",
    "defines the additional rotation about the z' axis of the x' and y' axes.\r\n",
    "\r\n",
    "The discretised space is returned as a (N,4) numpy array with N elements and columns representing:\r\n",
    "\r\n",
    "- 0: The $\\theta$ value of the element\r\n",
    "- 1: The $\\phi$ value of the element\r\n",
    "- 2: The $\\chi$ value of the element\r\n",
    "- 3: The relative weighting of the element (should sum to 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_rotation_grid(theta_step, phi_step, chi_step):\r\n",
    "    # create spherical grid\r\n",
    "    theta = np.linspace(0, 180, theta_step+1)\r\n",
    "    phi = np.linspace(0, 360, phi_step+1)\r\n",
    "    chi = np.linspace(0, 360, chi_step+1)\r\n",
    "    # first index on theta, second on phi - use \"ij\"\r\n",
    "    THETA, PHI, CHI = np.meshgrid(theta, phi, chi, indexing=\"ij\")\r\n",
    "    centre_theta = np.zeros((THETA.shape[0]-1, THETA.shape[1]-1, THETA.shape[2]-1))\r\n",
    "    centre_phi = np.zeros((PHI.shape[0]-1, PHI.shape[1]-1, PHI.shape[2]-1))\r\n",
    "    centre_chi = np.zeros((CHI.shape[0]-1, CHI.shape[1]-1, CHI.shape[2]-1))\r\n",
    "    solid_angle = np.zeros((THETA.shape[0]-1, THETA.shape[1]-1, THETA.shape[2]-1))\r\n",
    "    for i in range(THETA.shape[0]-1):\r\n",
    "        for j in range(THETA.shape[1]-1):\r\n",
    "            for k in range(THETA.shape[2]-1):\r\n",
    "                centre_theta[i,j,k] = (THETA[i, j, k] + THETA[i+1,j, k])/2.0\r\n",
    "                centre_phi[i,j,k] = (PHI[i, j, k] + PHI[i,j+1, k])/2.0\r\n",
    "                centre_chi[i,j,k] = (CHI[i, j, k] + CHI[i,j, k+1])/2.0\r\n",
    "                # solid angle = sin(theta)d(theta)d(phi)\r\n",
    "                # chi is evenly distributed so don't include in relative weighting\r\n",
    "                solid_angle[i,j, k] = math.sin(np.radians(centre_theta[i,j, k])) \\\r\n",
    "                                    *  np.radians(THETA[i+1, j, k] - THETA[i,j, k]) \\\r\n",
    "                                    *  np.radians(PHI[i, j+1, k] - PHI[i,j, k])\r\n",
    "    return np.column_stack((centre_theta.flatten(), centre_phi.flatten(), \\\r\n",
    "                            centre_chi.flatten(), solid_angle.flatten()/(np.sum(solid_angle))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now define the range of Q values at which we want to calculate the average. Here we will choose 30 values from 0 to 0.3 (measured in per angstroms, the sasview default). We also create the data describing the sample, our discretisation of rotation space, and a numpy array to store the output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "Qs = np.linspace(3e-4, 0.3, 30)\r\n",
    "points = np.linspace(0, STEPSIZE*NODES, NODES, endpoint=False)\r\n",
    "pos_x, pos_y, pos_z = np.meshgrid(points, points, points)\r\n",
    "pos_x = pos_x.flatten()\r\n",
    "pos_y = pos_y.flatten()\r\n",
    "pos_z = pos_z.flatten()\r\n",
    "data = sas_gen.MagSLD(pos_x, pos_y, pos_z, np.full_like(pos_x, 6.97e-06))\r\n",
    "model = sas_gen.GenSAS()\r\n",
    "model.set_sld_data(data)\r\n",
    "angles = create_rotation_grid(ANGLESTEP, 2*ANGLESTEP, 2*ANGLESTEP)\r\n",
    "output = np.zeros_like(Qs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now carry out the actual averaging. We loop over the rotation space and sum together the weighted instensities in each orientation. We take the Q values along the $Q_x$ axis each time. As the sample is rotated this axis stays fixed in beamline coordinates - and so it sweeps through all the orientations of the sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for printing progress - the next fraction of completeness at which to print\r\n",
    "next_frac = 0.01\r\n",
    "for i in range(len(angles)):\r\n",
    "    r = Rotation.from_euler(\"ZYZ\", [np.radians(angles[i, 1]), np.radians(angles[i, 0]), np.radians(angles[i, 2])])\r\n",
    "    model.set_rotations(xyz_to_UVW=r)\r\n",
    "    # weights already normalised no need to divide by sum of weights\r\n",
    "    output += model.runXY([Qs, np.zeros_like(Qs)]) * angles[i, 3]\r\n",
    "    if i/len(angles) >= next_frac:\r\n",
    "        print(f\"{next_frac*100: .1f}%\")\r\n",
    "        next_frac += 0.01\r\n",
    "clear_output()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because our sample is grid type data - with only a nuclear SLD, we can also use the inbuilt debye full average function - and compare the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 19532 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "model.set_rotations() # reset to identity\r\n",
    "debye_output = model.run([Qs, []])\r\n",
    "plt.plot(Qs, output, \"x\", label=\"script generated\")\r\n",
    "plt.plot(Qs, debye_output, \"x\", label=\"debye average\")\r\n",
    "plt.xlabel(\"Q / A^-1\")\r\n",
    "plt.ylabel(\"Intensity / cm^-1\")\r\n",
    "plt.yscale(\"log\")\r\n",
    "plt.title(\"A plot of orientationally averaged scattering intensity at different Q values\")\r\n",
    "plt.legend()\r\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In fact a comparison with the analytical result from the fitting calculator in Sasview would show that both of these results differ signifcantly from the correct value at higher Q values on the graph. This is because the cube has only been discretised in a 10x10x10 grid in real space. by changing the initial paramters to give a better discretisation for the same shape, such as:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "NODES = 20\r\n",
    "STEPSIZE = 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We would create a significantly more accurate result."
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "2c2de9287569066ce7117e224f7f9badc698e8d46797fb4236b3d8947d000fab"
  },
  "kernelspec": {
   "display_name": "Python 3.6.13 64-bit ('sasview_env': conda)",
   "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.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}