{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# This cell is added by sphinx-gallery\n\n%matplotlib inline\n\nimport mrsimulator\nprint(f'You are using mrsimulator v{mrsimulator.__version__}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Simulating site disorder (crystalline)\n\n87Rb (I=3/2) 3QMAS simulation with site disorder.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The following example illustrates an NMR simulation of a crystalline solid with site\ndisorders. We model such disorders with Extended Czjzek distribution. The following\ncase study shows an $^{87}\\text{Rb}$ 3QMAS simulation of RbNO3.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib as mpl\nimport numpy as np\nfrom mrsimulator import Simulator\nfrom mrsimulator.methods import ThreeQ_VAS\nimport matplotlib.pyplot as plt\n\nfrom mrsimulator.models import ExtCzjzekDistribution\nfrom mrsimulator.utils.collection import single_site_system_generator\nfrom scipy.stats import multivariate_normal\n\n# global plot configuration\nmpl.rcParams[\"figure.figsize\"] = [4.5, 3.0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Generate probability distribution\nCreate three extended Czjzek distributions for the three sites in RbNO3 about their\nrespective mean tensors.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# The range of isotropic chemical shifts, the quadrupolar coupling constant, and\n# asymmetry parameters used in generating a 3D grid.\niso_r = np.arange(101) / 6.5 - 35  # in ppm\nCq_r = np.arange(100) / 100 + 1.25  # in MHz\neta_r = np.arange(11) / 10\n\n# The 3D mesh grid over which the distribution amplitudes are evaluated.\niso, Cq, eta = np.meshgrid(iso_r, Cq_r, eta_r, indexing=\"ij\")\n\n\ndef get_prob_dist(iso, Cq, eta, eps, cov):\n    pdf = 0\n    for i in range(len(iso)):\n        # The 2D amplitudes for Cq and eta is sampled from the extended Czjzek model.\n        avg_tensor = {\"Cq\": Cq[i], \"eta\": eta[i]}\n        _, _, amp = ExtCzjzekDistribution(avg_tensor, eps=eps[i]).pdf(pos=[Cq_r, eta_r])\n\n        # The 1D amplitudes for isotropic chemical shifts is sampled as a Gaussian.\n        iso_amp = multivariate_normal(mean=iso[i], cov=[cov[i]]).pdf(iso_r)\n\n        # The 3D amplitude grid is generated as an uncorrelated distribution of the\n        # above two distribution, which is the product of the two distributions.\n        pdf_t = np.repeat(amp, iso_r.size).reshape(eta_r.size, Cq_r.size, iso_r.size)\n        pdf_t *= iso_amp\n        pdf += pdf_t\n    return pdf\n\n\niso_0 = [-27.4, -28.5, -31.3]  # isotropic chemical shifts for the three sites in ppm\nCq_0 = [1.68, 1.94, 1.72]  # Cq in MHz for the three sites\neta_0 = [0.2, 1, 0.5]  # eta for the three sites\neps_0 = [0.02, 0.02, 0.02]  # perturbation fractions for extended Czjzek distribution.\nvar_0 = [0.1, 0.1, 0.1]  # variance for the isotropic chemical shifts in ppm^2.\n\npdf = get_prob_dist(iso_0, Cq_0, eta_0, eps_0, var_0).T"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The two-dimensional projections from this three-dimensional distribution are shown\nbelow.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "_, ax = plt.subplots(1, 3, figsize=(9, 3))\n\n# isotropic shift v.s. quadrupolar coupling constant\nax[0].contourf(Cq_r, iso_r, pdf.sum(axis=2))\nax[0].set_xlabel(\"Cq / MHz\")\nax[0].set_ylabel(\"isotropic chemical shift / ppm\")\n\n# isotropic shift v.s. quadrupolar asymmetry\nax[1].contourf(eta_r, iso_r, pdf.sum(axis=1))\nax[1].set_xlabel(r\"quadrupolar asymmetry, $\\eta$\")\nax[1].set_ylabel(\"isotropic chemical shift / ppm\")\n\n# quadrupolar coupling constant v.s. quadrupolar asymmetry\nax[2].contourf(eta_r, Cq_r, pdf.sum(axis=0))\nax[2].set_xlabel(r\"quadrupolar asymmetry, $\\eta$\")\nax[2].set_ylabel(\"Cq / MHz\")\nplt.tight_layout()\nplt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Simulation setup\nGenerate spin systems from the above probability distribution.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "spin_systems = single_site_system_generator(\n    isotopes=\"87Rb\",\n    isotropic_chemical_shifts=iso,\n    quadrupolar={\"Cq\": Cq * 1e6, \"eta\": eta},  # Cq in Hz\n    abundance=pdf,\n)\nlen(spin_systems)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Simulate a $^{27}\\text{Al}$ 3Q-MAS spectrum by using the `ThreeQ_MAS` method.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "method = ThreeQ_VAS(\n    channels=[\"87Rb\"],\n    magnetic_flux_density=9.4,  # in T\n    rotor_angle=54.735 * np.pi / 180,\n    spectral_dimensions=[\n        {\n            \"count\": 96,\n            \"spectral_width\": 7e3,  # in Hz\n            \"reference_offset\": -7e3,  # in Hz\n            \"label\": \"Isotropic dimension\",\n        },\n        {\n            \"count\": 256,\n            \"spectral_width\": 1e4,  # in Hz\n            \"reference_offset\": -4e3,  # in Hz\n            \"label\": \"MAS dimension\",\n        },\n    ],\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Create the simulator object, add the spin systems and method, and run the simulation.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "sim = Simulator()\nsim.spin_systems = spin_systems  # add the spin systems\nsim.methods = [method]  # add the method\nsim.config.number_of_sidebands = 1\nsim.run()\n\ndata = sim.methods[0].simulation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The plot of the corresponding spectrum.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ax = plt.subplot(projection=\"csdm\")\ncb = ax.imshow(data / data.max(), cmap=\"gist_ncar_r\", aspect=\"auto\")\nax.set_ylim(-40, -70)\nax.set_xlim(-20, -60)\nplt.colorbar(cb)\nplt.tight_layout()\nplt.show()"
      ]
    }
  ],
  "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": 0
}