TFY4255_Lab2_PC_2023.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Template Code for Lab 2 - Computer Lab (Code written by Tina Bergh)\n",
    "TFY4255 Materials Physics 2023\n",
    "\n",
    "Lecturer: Antonius T.J. van Helvoort\n",
    "\n",
    "Lab assistant: Syam Sadan"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a template Python code (Python 3.6) written in a Jupyter notebook for lab exercise 2 in materials physics. \n",
    "The instructions are given in the lab text. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Jupyter notebook bascis: \n",
    "- To run the cell where your cursor is, press **Shift and Enter**\n",
    "- In [*] indicates that a cell is running, and a number appears after completion\n",
    "\n",
    "To modify cells, click at the [ ]-icon outside of text edit mode, and\n",
    "- a inserts a cell above\n",
    "- b inserts a cell below\n",
    "- m changes the cell to markdown, text mode\n",
    "- y changes it back to code\n",
    "- Shift and Ctrl and - splits the cell at the position of your cursor\n",
    "- Shift and Tab inside the parathesis after a function name shows you the docstring, where you easily can find input variables and a description of the function\n",
    "\n",
    "For tips on how to better use Jupyter notebook, there are many online resources, see e.g. https://www.dataquest.io/blog/jupyter-notebook-tips-tricks-shortcuts/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import the required libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib widget\n",
    "#%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "#from scipy.constants import pi\n",
    "from numpy import pi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 1: Diffraction from 1D crystals"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 Model the electron density of 1D atomic chains"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define a function that returns the x-axis. Useful e.g. if you want to customize a plot of the electron density. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_x_axis(a_0=2., num_atoms=50, x_steplength=0.1):   \n",
    "    \n",
    "    \"\"\"Returns an x-axis with the given steplength within the bounds given by num_atoms.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x_steplength : float\n",
    "        steplength between each value on the x-axis.\n",
    "    a_0: float\n",
    "        lattice constant; equilibrium position of atoms in [Å]. \n",
    "    num_atoms: float\n",
    "        Number of atoms\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    x_axis : ndarray\n",
    "        x_axis corresponding to the x_axis where the electron density is calculated by get_electron_density.\n",
    "    \"\"\"\n",
    "    \n",
    "    return np.arange(start=-(num_atoms)*a_0/2,\n",
    "                     stop=(num_atoms)*a_0/2, \n",
    "                     step=x_steplength)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define a function that calculates the total electron density from a customisable 1D atomic chain. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_electron_density(x_steplength = 0.1,\n",
    "                         model = 'Monoatomic',\n",
    "                         sigma_1 = .1,\n",
    "                         sigma_2 = .2,\n",
    "                         a_0 = 2.,\n",
    "                         thermal_disorder = False,\n",
    "                         num_atoms = 50,\n",
    "                         epsilon = 0.2,\n",
    "                         epsilon_glass = 0.2,\n",
    "                         amplitude_1 = 1.,\n",
    "                         amplitude_2 = 2.,\n",
    "                         plot_electron_density = True,\n",
    "                         return_x_axis = False):\n",
    "    \n",
    "    \"\"\"Calculates the total electron density of a 1D atomic chain.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x_steplength : float\n",
    "        steplength between each value on the x-axis: the value where the electron\n",
    "        density is calculated.   \n",
    "    a_0: float\n",
    "        lattice constant; equilibrium position of atoms. \n",
    "    num_atoms: float\n",
    "        Number of atoms.\n",
    "    model : string\n",
    "        Model type: 'Monoatomic', 'Diatomic' or 'Glass', determines the expression\n",
    "        for electron density.\n",
    "    amplitude_1, amplitude_2: float\n",
    "        Amplitude; factor multiplied with the Gaussian at each atom position.\n",
    "        amplitude_2 is used in the diatomic model, otherwise only amplitude_1.\n",
    "    sigma_1, sigma_2: float\n",
    "        Width of the Gaussians in the expression for electron density.\n",
    "        sigma_2 is used in diatomic model, otherwise only sigma_1.\n",
    "    thermal_disorder: bool\n",
    "        True if thermal disorder is to be included in the model, and False (default) \n",
    "        for no thermal disorder.\n",
    "    epsilon: float\n",
    "        Only used for thermal_disorder=True. \n",
    "        Limit of the random number (epsilon * random_number, where random_number\n",
    "        is taken from the interval (-1,1)) that is added to the atom position. \n",
    "        Only used for thermal_disorder=True. \n",
    "    epsilon_glass: float\n",
    "        Only used for model='glass'. \n",
    "        Limit of the random number (epsilon * random_number, where random_number\n",
    "        is taken from the interval (-1,1)) that is added to the atom position. \n",
    "        Only used for thermal_disorder=True. \n",
    "    plot_electron_density: bool\n",
    "        Determines if electron density should be plotted.\n",
    "    return_x_axis : bool\n",
    "        If True (False is default), the x_axis is returned in addition to the electron_density.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    electron_density : ndarray\n",
    "        Electron density at positions defined by x_axis.\n",
    "    \"\"\"\n",
    "    \n",
    "    #Spatial discretisation. Create an x-axis by using the function get_x_axis.\n",
    "    x_axis = get_x_axis(a_0=a_0,\n",
    "                        num_atoms=num_atoms,\n",
    "                        x_steplength=x_steplength) \n",
    "    #Initiate variables\n",
    "    electron_density = np.zeros(np.shape(x_axis))\n",
    "    counter = 0\n",
    "    # Initial atom position\n",
    "    atom_position = - (num_atoms - 1 ) * a_0 / 2 - a_0\n",
    "    \n",
    "    for i in range(num_atoms):\n",
    "        if thermal_disorder: \n",
    "            atom_position = -(num_atoms - 1) * a_0 / 2 + a_0*(i) + epsilon*(2 * np.random.random() - 1)\n",
    "        else:\n",
    "            atom_position = atom_position + a_0\n",
    "        if model == 'Monoatomic':\n",
    "            this_atoms_electron_density = amplitude_1 * np.exp(-(x_axis-atom_position)**2 / ( 2 * sigma_1**2) )\n",
    "        elif model == 'Diatomic':\n",
    "            if np.fmod(counter,2):\n",
    "                sigma_i = sigma_1\n",
    "                amplitude_i = amplitude_1\n",
    "            else:\n",
    "                sigma_i = sigma_2\n",
    "                amplitude_i = amplitude_2\n",
    "            this_atoms_electron_density = amplitude_i * np.exp(-(x_axis-atom_position)**2 / ( 2 * sigma_i**2) )\n",
    "            counter = counter+1\n",
    "        elif model == 'Glass':\n",
    "            atom_position = atom_position + epsilon_glass * (2 * np.random.random() - 1)\n",
    "            this_atoms_electron_density = amplitude_1 * np.exp(-(x_axis-atom_position)**2 / ( 2 * sigma_1**2) )\n",
    "        else: \n",
    "            print('Error: Must specify valid model type in string format.')\n",
    "            \n",
    "        electron_density = electron_density + this_atoms_electron_density\n",
    "    \n",
    "    if plot_electron_density:\n",
    "        plt.figure()\n",
    "        plt.plot(x_axis,\n",
    "                 electron_density, 'k')\n",
    "        plt.ylabel('Electron density [arb.]')\n",
    "        plt.xlabel('Distance [$\\AA$]')\n",
    "        if thermal_disorder:\n",
    "            plt.title(model+' 1D model with thermal disorder and with '+str(num_atoms)+' atoms.')\n",
    "        else:\n",
    "            plt.title(model+' 1D model with '+str(num_atoms)+' atoms.')\n",
    "    if return_x_axis:\n",
    "        return x_axis, electron_density\n",
    "    else:\n",
    "        return electron_density"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Explore different models and parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Read the documentation in the function above to see what the different parameters are. Select parameters that will be given to the function get_electron_density. \n",
    "- Explore different parameters and different models. Note the differences between the models and how the different parameters affect the results. \n",
    "\n",
    "<i>(Yes, this can also be specified directly into the call of the function. Here it is defined as global variables, so that it can be reused.)</i>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Steplength between each value on the x-axis in [Å]. \n",
    "x_steplength = 0.001\n",
    "\n",
    "# Model type: 'Monoatomic', 'Diatomic' or 'Glass', determines the expression\n",
    "# for electron density.\n",
    "model='Monoatomic'\n",
    "\n",
    "# Number of atoms. \n",
    "num_atoms = 100\n",
    "\n",
    "# Distance between atoms in [Å].\n",
    "a_0 = 2. \n",
    "\n",
    "# The distribution width of the electron density of each atom, \n",
    "# i.e. the width of the Gaussians in [Å].\n",
    "sigma_1 = .1 \n",
    "\n",
    "# Magnitude of the electron density of each atom, \n",
    "# i.e. factor multiplied with the Gaussians at each atom position.\n",
    "amplitude_1 = 1\n",
    "\n",
    "# Only used for thermal_disorder=True. \n",
    "# Limit of the random number (epsilon * random_number, where random_number\n",
    "# is taken from the interval (-1,1)) that is added to the atom position. \n",
    "epsilon = 0.2\n",
    "\n",
    "# Width of the Gaussians in the expression for electron density.\n",
    "# sigma_2 is used in diatomic model.\n",
    "sigma_2 = 0.15\n",
    "\n",
    "# Amplitude; factor multiplied with the Gaussian at each atom position.\n",
    "# amplitude_2 is used in the diatomic model.\n",
    "amplitude_2 = 1.3\n",
    "\n",
    "# Only used for model='glass'. \n",
    "# Limit of the random number (epsilon * random_number, where random_number\n",
    "# is taken from the interval (-1,1)) that is added to the atom position. \n",
    "epsilon_glass = 0.5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate and plot the electron density. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "electron_density_mono_atomic = get_electron_density(x_steplength=x_steplength, \n",
    "                                                    model=model, \n",
    "                                                    sigma_1=sigma_1, \n",
    "                                                    a_0=a_0, \n",
    "                                                    num_atoms=num_atoms, \n",
    "                                                    amplitude_1=amplitude_1,\n",
    "                                                    thermal_disorder=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot several 1D models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'Electron densities of 1D models')"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_axis = get_x_axis(a_0=a_0,\n",
    "                    num_atoms=num_atoms,\n",
    "                    x_steplength=x_steplength) \n",
    "\n",
    "electron_density_mono_atomic = get_electron_density(x_steplength=x_steplength,\n",
    "                                                    sigma_1=sigma_1,\n",
    "                                                    a_0=a_0,\n",
    "                                                    num_atoms=num_atoms,\n",
    "                                                    amplitude_1=amplitude_1, \n",
    "                                                    thermal_disorder=False,\n",
    "                                                    plot_electron_density=False)\n",
    "\n",
    "electron_density_diatomic = get_electron_density(x_steplength=x_steplength,\n",
    "                                                 model='Diatomic',\n",
    "                                                 sigma_1=sigma_1,\n",
    "                                                 sigma_2=sigma_2,\n",
    "                                                 a_0=a_0,\n",
    "                                                 num_atoms=num_atoms,\n",
    "                                                 amplitude_1=amplitude_1,\n",
    "                                                 amplitude_2=amplitude_2,\n",
    "                                                 thermal_disorder=False,\n",
    "                                                 plot_electron_density=False)\n",
    "\n",
    "electron_density_mono_atomic_thermal = get_electron_density(x_steplength=x_steplength,\n",
    "                                                            sigma_1=sigma_1,\n",
    "                                                            a_0=a_0,\n",
    "                                                            num_atoms=num_atoms,\n",
    "                                                            amplitude_1=amplitude_1, \n",
    "                                                            thermal_disorder=True,\n",
    "                                                            epsilon = epsilon,\n",
    "                                                            plot_electron_density=False)\n",
    "\n",
    "electron_density_glass = get_electron_density(x_steplength=x_steplength,\n",
    "                                              model='Glass',\n",
    "                                              sigma_1=sigma_1,\n",
    "                                              a_0=a_0,\n",
    "                                              num_atoms=num_atoms,\n",
    "                                              epsilon_glass=epsilon_glass,\n",
    "                                              amplitude_1=amplitude_1,\n",
    "                                              plot_electron_density=False)\n",
    "\n",
    "\n",
    "\n",
    "plt.figure(figsize=(12,5))\n",
    "plt.plot(x_axis,\n",
    "         electron_density_mono_atomic, 'orange', label=('mono-atomic'))\n",
    "plt.plot(x_axis,\n",
    "         electron_density_mono_atomic_thermal, 'b--', label=('mono-atomic thermal'))\n",
    "plt.plot(x_axis,\n",
    "         electron_density_glass, 'r', label=('glass'), linestyle='dashdot')\n",
    "plt.plot(x_axis,\n",
    "         electron_density_diatomic, 'g', label=('diatomic'), linestyle='dotted')\n",
    "plt.ylabel('Electron density [arb.]')\n",
    "plt.xlabel('Distance [Å]')\n",
    "axes = plt.gca()\n",
    "#axes.set_xlim([-(num_atoms+2)*a_0/2,\n",
    "#               (num_atoms+2)*a_0/2])\n",
    "plt.legend(fontsize = 'x-small')\n",
    "plt.title('Electron densities of 1D models')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.2 Fast Fourier Transforms of the Electron Density"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.signal.windows import tukey"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Tukey window can be useful when computing the FFTs. \n",
    "Take a look at the Tukey (tapered cosine) window, calcualte and plot it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\helvoort\\Anaconda3\\envs\\pyXem014\\lib\\site-packages\\matplotlib\\cbook\\__init__.py:1298: ComplexWarning: Casting complex values to real discards the imaginary part\n",
      "  return np.asarray(x, float)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'tukey window * electron_density_mono_atomic')"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "alpha=0.5\n",
    "\n",
    "x_axis = get_x_axis() \n",
    "\n",
    "tukey_window = tukey(len(x_axis),alpha=alpha)\n",
    "FFT_tukey_window = np.fft.fftshift(np.fft.fft(tukey_window))\n",
    "Q_axis = (np.arange(len(FFT_tukey_window)) - len(FFT_tukey_window)/2) / (len(FFT_tukey_window)/2 * x_steplength / pi)\n",
    "\n",
    "electron_density_mono_atomic_tw = get_electron_density(plot_electron_density=False) * tukey_window\n",
    "\n",
    "fig, axes = plt.subplots(ncols=3, figsize=(18, 6), sharex=False, sharey=False)\n",
    "ax = axes.ravel()\n",
    "ax[0].plot(x_axis,tukey_window, 'k')\n",
    "ax[0].set_title('tukey window')\n",
    "ax[1].plot(Q_axis,FFT_tukey_window, 'k')\n",
    "ax[1].set_title('FFT of tukey window (note xlim)')\n",
    "#ax[1].set_xlim([-1,1])\n",
    "ax[1].set_xlim([x_axis[0], x_axis[-1]])\n",
    "ax[2].plot(x_axis,electron_density_mono_atomic_tw, 'k')\n",
    "ax[2].set_title('tukey window * electron_density_mono_atomic')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define a function that calculates the FFT and the intensity of an electron density. \n",
    "\n",
    "<i>NB: Here intensity is simply given by the FFT multiplied with its complex conjugate, even though the theory says that this is only a proportionality. </i>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_intenisty_of_electron_density(x_steplength=.1,\n",
    "                                      model='Monoatomic',\n",
    "                                      sigma_1=.1,\n",
    "                                      sigma_2=.2,\n",
    "                                      a_0=2.,\n",
    "                                      thermal_disorder=False,\n",
    "                                      num_atoms=50,\n",
    "                                      epsilon=0.2,\n",
    "                                      epsilon_glass=0.2,\n",
    "                                      amplitude_1=1.,\n",
    "                                      amplitude_2=2.,\n",
    "                                      tukey_window=True, \n",
    "                                      plot_on=True):\n",
    "    \n",
    "    \"\"\"Calculates the intenisty of the electron density.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x_steplength : float\n",
    "        steplength between each value on the x-axis: the value where the electron\n",
    "        density is calculated.   \n",
    "    a_0: float\n",
    "        lattice constant; equilibrium position of atoms. \n",
    "    num_atoms: float\n",
    "        Number of atoms.\n",
    "    model : string\n",
    "        Model type: 'Monoatomic', 'Diatomic' or 'Glass', determines the expression\n",
    "        for electron density.\n",
    "    amplitude_1, amplitude_2: float\n",
    "        Amplitude; factor multiplied with the Gaussian at each atom position.\n",
    "        amplitude_2 is used in the diatomic model, otherwise only amplitude_1.\n",
    "    sigma_1, sigma_2: float\n",
    "        Width of the Gaussians in the expression for electron density.\n",
    "        sigma_2 is used in diatomic model, otherwise only sigma_1.\n",
    "    thermal_disorder: bool\n",
    "        True if thermal disorder is to be included in the model, and False (default) \n",
    "        for no thermal disorder.\n",
    "    epsilon: float\n",
    "        Only used for thermal_disorder=True. \n",
    "        Limit of the random number (epsilon * random_number, where random_number\n",
    "        is taken from the interval (-1,1)) that is added to the atom position. \n",
    "    epsilon_glass: float\n",
    "        Only used for model='glass'. \n",
    "        Limit of the random number (epsilon * random_number, where random_number\n",
    "        is taken from the interval (-1,1)) that is added to the atom position. \n",
    "    tukey_window : bool\n",
    "        If True, Tukey window is applied, i.e. multiplied with the electron density.\n",
    "    plot_on : bool\n",
    "        If True, the electron density and its intenisty is plotted.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    x_axis : ndarray\n",
    "        x_axis corresponding to the x_axis where the electron density is calculated by get_electron_density.\n",
    "    \"\"\"\n",
    "    \n",
    "    x_axis, electron_density = get_electron_density(x_steplength=x_steplength,\n",
    "                                                    model=model,\n",
    "                                                    sigma_1=sigma_1,\n",
    "                                                    sigma_2=sigma_2,\n",
    "                                                    a_0=a_0,\n",
    "                                                    thermal_disorder=thermal_disorder,\n",
    "                                                    num_atoms=num_atoms,\n",
    "                                                    epsilon=epsilon,\n",
    "                                                    epsilon_glass=epsilon_glass,\n",
    "                                                    amplitude_1=amplitude_1,\n",
    "                                                    amplitude_2=amplitude_2,\n",
    "                                                    plot_electron_density=False,\n",
    "                                                    return_x_axis=True)\n",
    "    \n",
    "    if tukey_window:\n",
    "        tukey_window = tukey(len(x_axis))\n",
    "        electron_density = electron_density * tukey_window\n",
    "\n",
    "    FFT_of_electron_density = np.fft.fftshift(np.fft.fft(electron_density))\n",
    "    \n",
    "    Q_axis = (np.arange(len(FFT_of_electron_density)) -\n",
    "              len(FFT_of_electron_density)/2) /\\\n",
    "              (len(FFT_of_electron_density)/2 * x_steplength * pi)\n",
    "    \n",
    "    intensity_of_electron_density = FFT_of_electron_density *\\\n",
    "                                    FFT_of_electron_density.conjugate()\n",
    "\n",
    "    if plot_on: \n",
    "        fig, axes = plt.subplots(ncols=2, figsize=(12, 6), #dpi = 200,\n",
    "                                 sharex=False, sharey=False)\n",
    "        ax = axes.ravel()\n",
    "        ax[0].plot(x_axis,electron_density, 'k')\n",
    "        ax[0].set_title('Electron density')\n",
    "        ax[0].set_xlabel('Distance [Å]')\n",
    "        ax[0].set_ylabel('Intensity [a.u.]')\n",
    "        ax[1].plot(Q_axis,intensity_of_electron_density, 'k')\n",
    "        ax[1].set_title('Intensity of electron density')\n",
    "        ax[1].set_xlabel('Frequency [1/Å]')\n",
    "        ax[1].set_ylabel('Intensity [a.u.]')\n",
    "        plt.ticklabel_format(axis='y',style='sci',scilimits=(1,2))\n",
    "        \n",
    "    return intensity_of_electron_density"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Study the intensity. Again, investigate different models and paramters and understand how that influence the results. \n",
    "\n",
    "Note the limits on the q_x-axis in reciprocal space. You might need to zoom in."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Steplength between each value on the x-axis in [Å]. \n",
    "x_steplength = 0.001\n",
    "\n",
    "# Model type: 'Monoatomic', 'Diatomic' or 'Glass', determines the expression\n",
    "# for electron density.\n",
    "model='Diatomic'\n",
    "\n",
    "# Number of atoms. \n",
    "num_atoms = 50\n",
    "\n",
    "# Distance between atoms in [Å].\n",
    "a_0 = 2. \n",
    "\n",
    "# The distribution width of the electron density of each atom, \n",
    "# i.e. the width of the Gaussians in [Å].\n",
    "sigma_1 = .1 \n",
    "\n",
    "# Magnitude of the electron density of each atom, \n",
    "# i.e. factor multiplied with the Gaussians at each atom position.\n",
    "amplitude_1 = 1\n",
    "\n",
    "# True if thermal disorder is to be included in the model, \n",
    "# and False for no thermal disorder.\n",
    "thermal_disorder=False\n",
    "\n",
    "# Only used for thermal_disorder=True. \n",
    "# Limit of the random number (epsilon * random_number, where random_number\n",
    "# is taken from the interval (-1,1)) that is added to the atom position. \n",
    "epsilon = 0.2\n",
    "\n",
    "# Width of the Gaussians in the expression for electron density.\n",
    "# sigma_2 is used in diatomic model.\n",
    "sigma_2 = 0.15\n",
    "\n",
    "# Amplitude; factor multiplied with the Gaussian at each atom position.\n",
    "# amplitude_2 is used in the diatomic model.\n",
    "amplitude_2 = 1.3\n",
    "\n",
    "# Only used for model='glass'. \n",
    "# Limit of the random number (epsilon * random_number, where random_number\n",
    "# is taken from the interval (-1,1)) that is added to the atom position. \n",
    "epsilon_glass = 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\helvoort\\Anaconda3\\envs\\pyXem014\\lib\\site-packages\\ipykernel\\eventloops.py:107: UserWarning: Attempted to set non-positive bottom ylim on a log-scaled axis.\n",
      "Invalid limit will be ignored.\n",
      "  app.exec_()\n"
     ]
    }
   ],
   "source": [
    "intensity_of_electron_density_mono_atomic = get_intenisty_of_electron_density(\n",
    "    x_steplength=x_steplength,\n",
    "    model=model,\n",
    "    sigma_1=sigma_1,\n",
    "    sigma_2=sigma_2,\n",
    "    a_0=a_0,\n",
    "    thermal_disorder=thermal_disorder,\n",
    "    num_atoms=num_atoms,\n",
    "    epsilon=epsilon,\n",
    "    epsilon_glass=epsilon_glass,\n",
    "    amplitude_1=amplitude_1,\n",
    "    amplitude_2=amplitude_2,\n",
    "    plot_on=True,\n",
    "    tukey_window=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 2: Fast Fourier Transformations and Imaging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.image as mpimg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import an image. If it is not located in the same folder as your Jupyter notebook, you need to include the whole file directory. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = mpimg.imread('Example images/grid_1000.tif')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You probably know that the image is simply a large matrix, where each element corresponds to the value of a certain pixel. Plot the original image and take a look."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "plt.imshow(image)\n",
    "ax = fig.gca()\n",
    "ax.set_axis_off()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Noise\n",
    "Optionally add noise to the image, in order to study how that affects the simulated diffraction pattern, and in order to study noise-reduction and blurring. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_noise(image_temp, mu=0, sigma=0.1):\n",
    "    \"\"\"Adds Gaussian noise to an image. \n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    image_temp: ndarray\n",
    "        Color or grayscale image. \n",
    "    mu : float\n",
    "        Mean (centre) of the Gaussian distribution.\n",
    "    scale : float\n",
    "        Standard deviation (spread or width) of the distribution.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    noise : ndarray\n",
    "        The noise that is added to the image. \n",
    "    noisy_image : ndarray\n",
    "        Image with noise added. \n",
    "    \"\"\"\n",
    "    C = np.max(image_temp)\n",
    "    noise = (np.random.normal(loc=mu, scale=sigma, size=image_temp.shape)*C).astype('int')\n",
    "    noisy_image = image_temp + noise\n",
    "    noisy_image[noisy_image<0] = 0\n",
    "    noisy_image[noisy_image>C] = C\n",
    "    \n",
    "    return noise, noisy_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "noise, image = add_noise(image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).\n"
     ]
    }
   ],
   "source": [
    "plt.figure()\n",
    "plt.imshow(noise)\n",
    "\n",
    "fig = plt.figure()\n",
    "plt.imshow(image)\n",
    "ax = fig.gca()\n",
    "ax.set_axis_off()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define a function that calculates a grayscale image from a color image. This is useful, since it is easier to compute the FFT from a grayscale image (or an image with only one color channel). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_grayscale_image(image_temp): \n",
    "    \"\"\"Converts a grayscale image to a color image.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    image_temp: ndarray\n",
    "        Color image, i.e. image with several channels (multi-dimensional). \n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    gray_image : ndarray\n",
    "        Grayscale image, i.e. image with only one channel (two-dimensional).\n",
    "    \"\"\"\n",
    "    try: \n",
    "        if np.shape(image_temp[0,0])[0] > 2: \n",
    "            gray_image = image_temp[...,0]*0.299 + image_temp[...,1]*0.587 + image_temp[...,2]*0.114\n",
    "            if np.shape(image_temp[0,0])[0] > 3: \n",
    "                gray_image[np.where(image_temp[...,3] == 0.)] = 1.\n",
    "            return gray_image\n",
    "    except IndexError as e: \n",
    "        return image_temp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define a function that applies a Tukey window in 2D to an image, similar to the 1D case in part 1. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_window_to_image(image_temp, alpha = 1.0):\n",
    "    \"\"\"Multiplies a grayscale image with a Tukey window function.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    image_temp: ndarray\n",
    "        Grayscale image, i.e. image with one channel (two-dimensional). \n",
    "    alpha : float, optional\n",
    "        (From the scipy documentation) Shape parameter of the Tukey window, \n",
    "        representing the fraction of the window inside the cosine tapered region. \n",
    "        If zero, the Tukey window is equivalent to a rectangular window. \n",
    "        If one, the Tukey window is equivalent to a Hann window.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    image_temp * window_2D : ndarray\n",
    "        Grayscale image multiplied with the window function.\n",
    "    \n",
    "    \"\"\"\n",
    "    Y,X = np.indices((image_temp.shape[0], image_temp.shape[1]))\n",
    "    Radius = (np.sqrt((X-int(image_temp.shape[0]/2))**2+(Y-int(image_temp.shape[1]/2))**2)+0.5).astype(int)\n",
    "    \n",
    "    window_2D = np.array(tukey(\n",
    "        int(np.sqrt(2*image_temp.shape[0]**2)), \n",
    "        alpha = alpha)[Radius.ravel()].reshape((image_temp.shape[0], image_temp.shape[1]))*(-1)+1)\n",
    "\n",
    "    return image_temp * window_2D"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define a function that calculates the intensity of a 2D image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_intensity_of_image(image_temp):\n",
    "    \"\"\"Calculates the FFT of a grayscale image, and the intensity (as FFT times conjugate(FFT)). \n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    image_temp: ndarray\n",
    "        Grayscale or one-channel image. \n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    FFT_image : ndarray\n",
    "        The FFT of the image.\n",
    "    intensity_image : ndarray\n",
    "        Absolute value of the FFT times its complex conjugate. \n",
    "    \"\"\"\n",
    "    \n",
    "    FFT_image = np.fft.fftshift(np.fft.fft2(image_temp))\n",
    "    intensity_image = np.abs(FFT_image * FFT_image.conjugate())\n",
    "\n",
    "    return FFT_image, intensity_image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define a function that calculates the intensity of a 2D image and optionally plots it on a log-scale, after ensuring that the image is square and gray-scale, and optionally after applying a Tukey window function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_intensity_of_image(image_temp, apply_window=True, alpha=1.0, zoom_factor=0.2, \n",
    "                           plot_intensity=True, cmap='inferno'):\n",
    "    \"\"\"Calculates the the intensity of a grayscale image (optionally multiplied by a window function), \n",
    "    as FFT times conjugate(FFT) and plots the image, the intensity and the central part of the intensity. \n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    image_temp: ndarray\n",
    "        Grayscale or one-channel image. \n",
    "    apply_window : bool, optional \n",
    "        If True (default), the iamge is multiplied by a window function.\n",
    "    alpha : float\n",
    "        (From the scipy documentation) Shape parameter of the Tukey window, \n",
    "        representing the fraction of the window inside the cosine tapered region. \n",
    "        If zero, the Tukey window is equivalent to a rectangular window. \n",
    "        If one, the Tukey window is equivalent to a Hann window.\n",
    "    zoom_factor : float, optional\n",
    "        Factor representing the fraction of the image that should be dispalyed.\n",
    "        Takes values (1,0).\n",
    "    plot_intensity : bool, optional\n",
    "        If True (default), the image and the intensity of the image will be plotted. \n",
    "    cmap : matplotlib.cm\n",
    "        Matplotlib colormap. See e.g. for the options: \n",
    "        https://matplotlib.org/3.1.0/api/cm_api.html#matplotlib.cm.get_cmap\n",
    "        Tips: by adding '_r', you invert the colormap!  \n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    image_temp : ndarray\n",
    "        The input image as a grayscale, square image, optionally with a Tukey window applied.\n",
    "    FFT_image : ndarray\n",
    "        The FFT of the image.\n",
    "    intensity : ndarray\n",
    "        Absolute value of the FFT of the image times its complex conjugate. \n",
    "    \"\"\"\n",
    "    # Convert to a grayscale image, since then it is easier to compute the FFT. \n",
    "    image_temp = get_grayscale_image(image_temp)\n",
    "    \n",
    "    # If the image is not square, crop the image, since otherwise, \n",
    "    # this would give artifacts in the FFT. \n",
    "    if image_temp.shape[0] != image_temp.shape[1]:\n",
    "        min_length = np.min([image_temp.shape[0], image_temp.shape[1]])\n",
    "        image_temp = image_temp[:min_length,:min_length]\n",
    "\n",
    "    # Optionally apply a Tukey window to the image. \n",
    "    if apply_window:\n",
    "        image_temp = apply_window_to_image(image_temp, alpha = alpha)\n",
    "\n",
    "    FFT_image, intensity = calculate_intensity_of_image(image_temp)\n",
    "    \n",
    "    if plot_intensity:\n",
    "            \n",
    "        # Also, crop the image to view only the central part\n",
    "        half_shape = int(intensity.shape[0]/2)\n",
    "        limit = [int((intensity.shape[0]/2)*(1-zoom_factor)),\n",
    "                int((intensity.shape[0]/2)*(1+zoom_factor))]\n",
    "        intensity_cropped = intensity[limit[0]:limit[1], limit[0]:limit[1]]\n",
    "        \n",
    "        fig, axes = plt.subplots(ncols=3, figsize=(12,4))\n",
    "        ax = axes.ravel()\n",
    "        ax[0].imshow(image_temp, cmap=cmap)\n",
    "        ax[0].set_title('Image')\n",
    "        ax[0].set_axis_off()\n",
    "        ax[1].imshow(np.log(intensity), cmap=cmap)\n",
    "        ax[1].set_title('Intensity')\n",
    "        ax[1].set_axis_off()\n",
    "        ax[2].imshow(np.log(intensity_cropped), cmap=cmap)\n",
    "        ax[2].set_title('Intensity - Central Part')\n",
    "        ax[2].set_axis_off()\n",
    "    \n",
    "    return image_temp, FFT_image, intensity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "img, FFT_img, img_int = get_intensity_of_image(image)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define a function that enables masking of the intensity image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mask_FFT_image(FFT_image,\n",
    "                   intensity_image,\n",
    "                   mask_type='circular',\n",
    "                   mask_length=None,\n",
    "                   x_offset=0,\n",
    "                   y_offset=0,\n",
    "                   invert_mask=True,\n",
    "                   make_plot=True,\n",
    "                   cmap='inferno'): \n",
    "    \n",
    "    \"\"\"Masks an FFT and intensity, and calculates its inverse transform. \n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    FFT_image: ndarray\n",
    "        FFT of a one-channel image. \n",
    "    intensity_image : ndarray\n",
    "        Absolute value of the FFT times its complex conjugate. \n",
    "    mask_type: string\n",
    "        'circular' or 'rectangular': determines the shape of the mask. \n",
    "    mask_length: float\n",
    "        Length of the mask. If type is 'circular', this is the radius, and for 'rectangular,\n",
    "        it is half of the rectangle's length. \n",
    "    x_offset, y_offset : float\n",
    "        mask center offset in horisontal or vertical direction \n",
    "    invert_mask : bool\n",
    "        If True, the mask will keep only the region of the FFT that inside of the mask. \n",
    "        If False, only that outside of the mask. \n",
    "    make_plot : bool\n",
    "        If True, the FFT, the masked FFT and the inverse tranformed masked FFT will be plotted. \n",
    "    cmap : matplotlib.cm\n",
    "        Matplotlib colormap. See e.g. for the options: \n",
    "        https://matplotlib.org/3.1.0/api/cm_api.html#matplotlib.cm.get_cmap\n",
    "        Tips: by adding '_r', you invert the colormap!\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    FFT_image : ndarray\n",
    "        FFT of image.\n",
    "    intensity_image : ndarray\n",
    "        Absolute value of the FFT times its complex conjugate. \n",
    "    \"\"\"\n",
    "    \n",
    "    FFT_mask = np.ones(np.shape(FFT_image),dtype=bool)\n",
    "    cx = np.round(np.shape(FFT_image)[1]/2)\n",
    "    cy = np.round(np.shape(FFT_image)[0]/2)\n",
    "\n",
    "    if mask_length is None: \n",
    "        mask_length=1/10*np.min(np.shape(FFT_image)[:1])\n",
    "    \n",
    "    if mask_type == 'rectangular':\n",
    "        FFT_mask[int(cy+y_offset-mask_length):int(cy+y_offset+mask_length), \n",
    "                 int(cx+x_offset-mask_length):int(cx+x_offset+mask_length)] = False\n",
    "        \n",
    "    elif mask_type == 'circular':\n",
    "        Y,X=np.indices((np.shape(FFT_image)[0],np.shape(FFT_image)[1]))+0.5\n",
    "        Radius = np.sqrt((X-(cx+x_offset))**2+(Y-(cy+y_offset))**2)+0.5\n",
    "        Radius = Radius.astype(int)\n",
    "        Radius -= int(Radius.min())\n",
    "        FFT_mask[np.where(Radius<mask_length)]=False\n",
    "    \n",
    "    if invert_mask: \n",
    "        masked_FFT = np.logical_not(FFT_mask) * FFT_image\n",
    "    else: \n",
    "        masked_FFT = FFT_mask * FFT_image\n",
    "        \n",
    "    abs_masked_image = np.abs(np.fft.ifft2(masked_FFT))\n",
    "    \n",
    "    if make_plot:\n",
    "        fig, axes = plt.subplots(ncols=3, figsize=(18, 6), \n",
    "                                 sharex=True, sharey=True,\n",
    "                                 subplot_kw={'adjustable': 'box-forced'})\n",
    "        ax = axes.ravel()\n",
    "        ax[0].imshow(np.log(intensity_image.astype(np.float32)), cmap=cmap, interpolation='none')\n",
    "        ax[0].set_title('FFT of image')\n",
    "        ax[1].imshow(np.log(np.abs(masked_FFT)), cmap=cmap, interpolation='none')\n",
    "        ax[1].set_title('Masked FFT of image')\n",
    "        ax[2].imshow(abs_masked_image, cmap=cmap, interpolation='none')\n",
    "        ax[2].set_title('IFFT of masked FFT')\n",
    "        for a in ax:\n",
    "            a.set_axis_off()\n",
    "        fig.tight_layout()\n",
    "\n",
    "    return np.abs(masked_FFT), abs_masked_image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define a function that incorporates all the functions defined earlier: calculates the intensity of the image, masks it in the frequency domain, before the image is back-transformed. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def back_transform_image(image_temp,\n",
    "                         apply_window=True,\n",
    "                         alpha = 1.,\n",
    "                         mask_FFT_on = True,\n",
    "                         mask_type = 'circular',\n",
    "                         mask_length = None,\n",
    "                         x_offset = 0,\n",
    "                         y_offset = 0,\n",
    "                         invert_mask = True,\n",
    "                         make_FFT_mask_plot = False,\n",
    "                         make_back_transform_image_plot = True,\n",
    "                         cmap='inferno'):\n",
    "    \"\"\"Masks an FFT and intensity, and calculates its inverse transform. \n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    image: ndarray\n",
    "        Color or grayscale image. If color, it will be tranformed into grayscale. \n",
    "        If not square, it will be cropped. \n",
    "    apply_window : bool\n",
    "        If True, a Tukey window will be applied to the grayscale image. \n",
    "    alpha: float 0.-1.\n",
    "        Shape parameter of the Tukey window.\n",
    "    mask_FFT_on : bool\n",
    "        If True, a mask will be applied to the FFT of the image. \n",
    "    mask_type: string\n",
    "        'circular' or 'rectangular': determines the shape of the mask. \n",
    "    mask_length: float\n",
    "        Length of the mask. If type is 'circular', this is the radius, and for 'rectangular,\n",
    "        it is half of the rectangle's length. \n",
    "    x_offset, y_offset : float\n",
    "        Mask center offset in horisontal or vertical direction \n",
    "    invert_mask : bool\n",
    "        If True, the mask will keep only the region of the FFT that is inside of the mask. \n",
    "        If False, only that outside of the mask will be kept. \n",
    "    make_FFT_mask_plot : bool\n",
    "        If True, the FFT, the masked FFT and the inverse tranformed masked FFT will be plotted. \n",
    "    make_back_transform_image_plot : bool\n",
    "        If True, all image processing steps will be shown in a figure. \n",
    "    cmap : matplotlib.cm\n",
    "        Matplotlib colormap. See e.g. for the options: \n",
    "        https://matplotlib.org/3.1.0/api/cm_api.html#matplotlib.cm.get_cmap\n",
    "        Tips: by adding '_r', you invert the colormap!\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    FFT_image : ndarray\n",
    "        FFT of image.\n",
    "    intensity_image : ndarray\n",
    "        Absolute value of the FFT times its complex conjugate. \n",
    "    \"\"\"\n",
    "    \n",
    "    image, FFT_image, intensity_image = get_intensity_of_image(\n",
    "        image_temp, apply_window=apply_window, alpha=alpha, plot_intensity=False)\n",
    "    \n",
    "    if mask_FFT_on:\n",
    "        abs_masked_FFT, abs_IFFT = mask_FFT_image(FFT_image,\n",
    "                                                  intensity_image,\n",
    "                                                  mask_type=mask_type,\n",
    "                                                  mask_length=mask_length,\n",
    "                                                  x_offset=x_offset,\n",
    "                                                  y_offset=y_offset,\n",
    "                                                  invert_mask=invert_mask,\n",
    "                                                  make_plot=make_FFT_mask_plot,\n",
    "                                                  cmap=cmap) \n",
    "    else:\n",
    "        abs_IFFT = np.abs(np.fft.ifft2(FFT_image))\n",
    "        \n",
    "    if make_back_transform_image_plot:\n",
    "        if mask_FFT_on: \n",
    "            fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(6, 4))#, \n",
    "                                     #sharex=True, sharey=True, subplot_kw={'adjustable': 'box'})\n",
    "            ax = axes.ravel()\n",
    "            ax[1].imshow(image, cmap=cmap, interpolation='none')\n",
    "            ax[1].set_title('Image with window')\n",
    "            ax[-3].imshow(np.log(intensity_image.astype(np.float32)), cmap=cmap, interpolation='none')\n",
    "            ax[-3].set_title('FFT')\n",
    "            ax[-2].imshow(np.log(abs_masked_FFT), cmap=cmap, interpolation='none')\n",
    "            ax[-2].set_title('Masked FFT')\n",
    "        else: \n",
    "            fig, axes = plt.subplots(ncols=2, nrows=2, figsize=(6, 6))#, sharex=True, sharey=True, subplot_kw={'adjustable': 'box-forced'})\n",
    "            ax = axes.ravel()\n",
    "            ax[1].imshow(image, cmap=cmap, interpolation='none')\n",
    "            ax[1].set_title('Image with window')\n",
    "            ax[-2].imshow(np.log(intensity_image.astype(np.float32)), cmap=cmap, interpolation='none')\n",
    "            ax[-2].set_title('FFT')\n",
    "        ax[0].imshow(image_temp, cmap=cmap, interpolation='none')\n",
    "        ax[0].set_title('Image')\n",
    "        ax[-1].imshow(abs_IFFT, cmap=cmap, interpolation='none')\n",
    "        ax[-1].set_title('IFFT')\n",
    "        for a in ax:\n",
    "            a.set_axis_off()\n",
    "        fig.tight_layout(h_pad = 0.1, w_pad = 0.1)\n",
    "        fig.set_dpi(200)\n",
    "\n",
    "    return abs_IFFT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Adjust the input parameters to customize the image processing. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the shape of the image to get an idea of suitable values for mask_length, x_offset, y_offset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1920, 2560, 3)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\helvoort\\AppData\\Local\\Temp\\ipykernel_19276\\2316985931.py:78: RuntimeWarning: divide by zero encountered in log\n",
      "  ax[-2].imshow(np.log(abs_masked_FFT), cmap=cmap, interpolation='none')\n"
     ]
    }
   ],
   "source": [
    "back_transform = back_transform_image(image_temp=image,\n",
    "                                      apply_window = False,\n",
    "                                      alpha = 1.,\n",
    "                                      mask_FFT_on = True,\n",
    "                                      mask_type = 'circular',\n",
    "                                      mask_length = 100,\n",
    "                                      x_offset = 0,\n",
    "                                      y_offset = 0,\n",
    "                                      invert_mask = False,\n",
    "                                      make_FFT_mask_plot = False,\n",
    "                                      make_back_transform_image_plot = True,\n",
    "                                      cmap='inferno')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 3: Diffraction from 2D lattices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_electron_density_2D(steplength = .1,\n",
    "                            model = 'Monoatomic',\n",
    "                            sigma_1 = .1,\n",
    "                            sigma_2 = .2,\n",
    "                            a_0 = 2.,\n",
    "                            num_atoms_per_row = 5,\n",
    "                            amplitude_1 = 1.,\n",
    "                            amplitude_2 = 2.,\n",
    "                            thermal_disorder = False,\n",
    "                            epsilon = 0.2,\n",
    "                            plot_electron_density = True,\n",
    "                            cmap='inferno'):\n",
    "    '''Calculates the total electron density of a 2D crystal.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    steplength : float\n",
    "        steplength, e.g. spatial discretisation\n",
    "    model : string\n",
    "        Model type: 'Monoatomic' or 'Diatomic'\n",
    "    sigma_1, sigma_2: float\n",
    "        Width of the Gaussians in the expression for electron density.\n",
    "        sigma_2 is used in diatomic model, otherwise only sigma_1.\n",
    "    a_0: float\n",
    "        lattice constant in [Å].\n",
    "    num_atoms_per_row: int\n",
    "        Number of atoms per row. The total number of atoms is num_atoms_per_row**2.\n",
    "    amplitude_1, amplitude_2: float\n",
    "        Amplitude; factor multiplied with the Gaussian at each atom position.\n",
    "        amplitude_2 is used in the diatomic model, otherwise only amplitude_1.\n",
    "    thermal_disorder: bool\n",
    "        True if thermal disorder is to be included in the model, and False (default) \n",
    "        for no thermal disorder.\n",
    "    epsilon: float\n",
    "        Limits of the random number (epsilon * random_number, where random_number\n",
    "        is taken from the interval (-1,1)) that is added to the atom position. \n",
    "        Only used for thermal_disorder=True. \n",
    "    plot_electron_density: bool\n",
    "        Determines if the electron density should be plotted.\n",
    "    cmap : matplotlib.cm\n",
    "        Matplotlib colormap. See e.g. for the options: \n",
    "        https://matplotlib.org/3.1.0/api/cm_api.html#matplotlib.cm.get_cmap\n",
    "        Tips: by adding '_r', you invert the colormap!\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    electron_density_2D : ndarray\n",
    "        The electron density in 2D. \n",
    "    '''\n",
    "    \n",
    "    half_length = (num_atoms_per_row-1)*a_0/2\n",
    "    atom_positions = np.meshgrid(np.arange(-half_length,half_length+a_0,a_0),\n",
    "                                np.arange(-half_length,half_length+a_0,a_0))    \n",
    "    if model == 'Glass':\n",
    "        print('Glass model is not defined in 2D.')\n",
    "        return 0\n",
    "    \n",
    "    if thermal_disorder: \n",
    "        atom_positions[0] = atom_positions[0] + epsilon* (2 * np.random.random(size = np.shape(atom_positions[0])) -1)\n",
    "        atom_positions[1] = atom_positions[1] + epsilon* (2 * np.random.random(size = np.shape(atom_positions[1])) -1)\n",
    "            \n",
    "    if model == 'Diatomic':\n",
    "        if np.fmod(num_atoms_per_row,2): # odd:\n",
    "            amplitudes = np.reshape(np.tile(np.array([amplitude_1,amplitude_2]), \n",
    "                                            int(num_atoms_per_row**2/2+num_atoms_per_row/2)),\n",
    "                                    newshape = (-1,num_atoms_per_row))[:num_atoms_per_row,:]\n",
    "            sigmas = np.reshape(np.tile(np.array([sigma_1,sigma_2]), \n",
    "                                            int(num_atoms_per_row**2/2+num_atoms_per_row/2)),\n",
    "                                    newshape = (-1,num_atoms_per_row))[:num_atoms_per_row,:]\n",
    "        else: \n",
    "            amplitudes = np.reshape(np.tile(np.array([amplitude_1,amplitude_2]), \n",
    "                                int(num_atoms_per_row**2/2+num_atoms_per_row/2)),\n",
    "                        newshape = (num_atoms_per_row,-1))[:num_atoms_per_row,:num_atoms_per_row]\n",
    "            sigmas = np.reshape(np.tile(np.array([sigma_1,sigma_2]), \n",
    "                                int(num_atoms_per_row**2/2+num_atoms_per_row/2)),\n",
    "                        newshape = (num_atoms_per_row,-1))[:num_atoms_per_row,:num_atoms_per_row]\n",
    "    else:\n",
    "        amplitudes = np.zeros_like(atom_positions[0])+amplitude_1\n",
    "        sigmas = np.zeros_like(atom_positions[0])+sigma_1\n",
    "        \n",
    "    def calc_this_atoms_electron_density_2D(atom_position_x, atom_position_y, amplitude, sigma):    \n",
    "        length = int( (num_atoms_per_row) * a_0 / steplength + 1)\n",
    "        center = length/2 - 0.5  \n",
    "        this_atoms_electron_density_2D = np.zeros((length, length))\n",
    "        X,Y = (np.indices((length, length)) - center) * steplength \n",
    "        return amplitude * np.exp(-((X - atom_position_x)**2 + (Y - atom_position_y)**2) / ( 2 * sigma**2) )\n",
    "\n",
    "    electron_density_2D = np.sum(list(map(calc_this_atoms_electron_density_2D,\n",
    "                               np.reshape(atom_positions[0],(-1,1)),\n",
    "                               np.reshape(atom_positions[1],(-1,1)),\n",
    "                               np.reshape(amplitudes,(-1,1)),\n",
    "                               np.reshape(sigmas,(-1,1)))),\n",
    "                            axis=0)\n",
    "    \n",
    "    if plot_electron_density:\n",
    "        plt.figure(dpi = 200)\n",
    "        plt.imshow(electron_density_2D, cmap=cmap)\n",
    "        plt.colorbar()\n",
    "        plt.minorticks_on()\n",
    "        pos = np.linspace(start = 0,\n",
    "                          stop = int( (num_atoms_per_row) * a_0 / steplength), \n",
    "                          num = 5, dtype = int)\n",
    "        plt.xticks(pos, (pos*steplength-half_length-a_0/2))\n",
    "        plt.yticks(pos, (pos*steplength-half_length-a_0/2))\n",
    "        plt.xlabel('Distance [Å]')\n",
    "        if thermal_disorder:\n",
    "            plt.title(model+' 1D model with thermal disorder and '+str(num_atoms_per_row**2)+' atoms.')\n",
    "        else:\n",
    "            plt.title(model+' 1D model with '+str(num_atoms_per_row**2)+' atoms.')\n",
    "    \n",
    "    return electron_density_2D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Steplength between each value on the axes in [Å]. \n",
    "# PS A too big value here might give you a memory error or consume time...\n",
    "steplength = 0.05\n",
    "\n",
    "# Model type: 'Monoatomic', 'Diatomic' or 'Glass', determines the expression\n",
    "# for electron density.\n",
    "model='Diatomic'\n",
    "\n",
    "# Number of atoms. The total number of atoms is num_atoms_per_row**2.\n",
    "num_atoms_per_row = 10\n",
    "\n",
    "# Distance between atoms in [Å]; lattice constant.\n",
    "a_0 = 2. \n",
    "\n",
    "# Width of the Gaussians in the expression for electron density in [Å].\n",
    "# sigma_2 is used in diatomic model.\n",
    "sigma_1 = .2 \n",
    "sigma_2 = 0.3\n",
    "\n",
    "# Amplitude; factor multiplied with the Gaussian at each atom position.\n",
    "# amplitude_2 is used in the diatomic model.\n",
    "amplitude_1 = 3\n",
    "amplitude_2 = 5\n",
    "\n",
    "# True if thermal disorder is to be included in the model, \n",
    "# and False for no thermal disorder.\n",
    "thermal_disorder=False\n",
    "\n",
    "# Only used for thermal_disorder=True. \n",
    "# Limits of the random number (epsilon * random_number, where random_number\n",
    "# is taken from the interval (-1,1)) that is added to the atom position. \n",
    "# Only used for thermal_disorder=True. \n",
    "epsilon = 0.2\n",
    "\n",
    "# Only used for model='glass'. \n",
    "# Limit of the random number (epsilon * random_number, where random_number\n",
    "# is taken from the interval (-1,1)) that is added to the atom position. \n",
    "epsilon_glass = 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "electron_density_diatomic_2D = get_electron_density_2D(steplength=steplength,\n",
    "                                                       model='Diatomic',\n",
    "                                                       sigma_1=sigma_1,\n",
    "                                                       sigma_2=sigma_2,\n",
    "                                                       a_0=a_0,\n",
    "                                                       num_atoms_per_row=num_atoms_per_row,\n",
    "                                                       amplitude_1=amplitude_1,\n",
    "                                                       amplitude_2=amplitude_2,\n",
    "                                                       thermal_disorder=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Optionally add some noise to the electron density"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "electron_density_diatomic_2D_noisy = add_noise(electron_density_diatomic_2D, sigma=0.15)[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the same function that you used to calculate the FFT, mask and calculate the IFFT of images, defined in section 2, on the 2D crystal. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\helvoort\\AppData\\Local\\Temp\\ipykernel_19276\\2316985931.py:78: RuntimeWarning: divide by zero encountered in log\n",
      "  ax[-2].imshow(np.log(abs_masked_FFT), cmap=cmap, interpolation='none')\n"
     ]
    }
   ],
   "source": [
    "back_transform = back_transform_image(image_temp = electron_density_diatomic_2D_noisy,\n",
    "                                      apply_window = False,\n",
    "                                      mask_FFT_on = True,\n",
    "                                      mask_type = 'circular',\n",
    "                                      mask_length = 50,\n",
    "                                      x_offset = 0,\n",
    "                                      y_offset = 0,\n",
    "                                      invert_mask = True,\n",
    "                                      make_FFT_mask_plot = False,\n",
    "                                      make_back_transform_image_plot = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:tfy4255]",
   "language": "python",
   "name": "conda-env-tfy4255-py"
  },
  "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.11.5"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}