[WIP] Add least-squares minimization capabiliy for Czjzek and Extended Czjzek

CHANGELOG

@@ -14,13 +14,18 @@ What's new
 - Support for user-defined isotopes using the ``Isotope.register()`` method. See the simulation gallery for use cases and examples.
 - Shortcuts for frequency contributions, such as ``Shielding``, ``Isotropic``, and ``cross``.
   Sets of contributions can also be excluded by placing an exclamation mark in front of the string, for example ``"!Shielding"`` excludes shielding interactions.
+- New functions for fitting Czjzek and Extended Czjzek tensor distribution models to experimental spectra. See (TODO: ADD) in the examples gallery for more information.
 - Adds ``DelayEvent`` to the events library.
 - Support for python 3.11
 
 **Simulator**
 
 - New instance method for the `Simulator` class -- `.optimize()` -- which pre-computes transition pathways before least-squares minimization. This improves the efficiency of least-squares minimizations.
 
+
+**Czjzek and Extended Czjzek**
+- The Czjzek model now uses an analytical expression for calculating the probability distribution greatly improving quality and calculation speed.
+
 **Bug Fixes**
 
 - Fixed bug where ``MixingEnum`` class had no attribute ``json`` (Issue [#260](https://github.com/deepanshs/mrsimulator/issues/260))
@@ -31,6 +36,7 @@ What's new
 - Fix bug when calculating ppm scale for large reference offsets.
 
 
+
 v0.7.0
 ------
 

docs/user_guide/spin_system_distributions/czjzek.rst

@@ -3,13 +3,21 @@
 Czjzek distribution
 -------------------
 
-The Czjzek distribution models random variations of a second-rank traceless
-symmetric tensors about zero, i.e., a tensor with zeta of zero. See :ref:`czjzek_model`
-for a mathematical description of the model as well as references to examples using the Czjzek
-distribution at the bottom of this page.
+The Czjzek distribution models random variations of second-rank traceless
+symmetric tensors about zero, i.e., a tensor with zeta of zero. An analytical expression
+for the Czjzek distribution exists (cite) which follows
 
-Czjzek distribution of symmetric shielding tensors
-''''''''''''''''''''''''''''''''''''''''''''''''''
+.. math::
+    f(\zeta, \eta, \sigma) = \eta \left(1-\frac{\eta^2}{9}\right)\frac{\zeta^4}{32\sigma^5 \sqrt{2 \pi}} \times \exp\left(-\frac{\zeta^2}{8\sigma^2}\left(1+\frac{\eta^2}{3}\right)\right),
+
+where :math:`\zeta` and :math:`\eta` are the Haberlen components of the tensor and :math:`\sigma` is the Czjzek width parameter. See :ref:`czjzek_model` for a further mathematical description of the model.
+
+The remainder of this page quickly describes how to generate Czjzek distributions and generate
+:py:class:`~mrsimulator.spin_system.SpinSystem` objects from these distributions. Also, look at the
+gallery examples using the Czjzek distribution listed at the bottom of this page.
+
+Creating and sampling a Czjzek distribution
+'''''''''''''''''''''''''''''''''''''''''''
 
 To generate a Czjzek distribution, use the :py:class:`~mrsimulator.models.CzjzekDistribution`
 class as follows.
@@ -21,7 +29,7 @@ class as follows.
 
     cz_model = CzjzekDistribution(sigma=0.8)
 
-The **CzjzekDistribution** class accepts a single argument, ``sigma``, which is the standard
+The **CzjzekDistribution** class accepts the argument, ``sigma``, which is the standard
 deviation of the second-rank traceless symmetric tensor parameters. In the above example,
 we create ``cz_model`` as an instance of the CzjzekDistribution class with
 :math:`\sigma=0.8`.
@@ -36,16 +44,13 @@ function. Let's first draw points from this distribution, using the
 
     zeta_dist, eta_dist = cz_model.rvs(size=50000)
 
-In the above example, we draw *50000* random points of the distribution. The output
-``zeta_dist`` and ``eta_dist`` hold the tensor parameter coordinates of the points, defined
-in the Haeberlen convention.
-The scatter plot of these coordinates is shown below.
+In the above example, we draw *50000* random points of the distribution. The output ``zeta_dist`` and ``eta_dist`` hold the tensor parameter coordinates of the points, defined in the Haeberlen convention. It is further assumed that the points in ``zeta_dist`` are in units of ``ppm`` while ``eta_dist`` has values since :math:`\eta` is dimensionless. The scatter plot of these coordinates is shown below.
 
 .. skip: next
 
 .. plot::
     :context: close-figs
-    :caption: Czjzek Distribution of shielding tensors.
+    :caption: Random sampling of Czjzek distribution of shielding tensors.
 
     import matplotlib.pyplot as plt
 
@@ -57,52 +62,211 @@ The scatter plot of these coordinates is shown below.
     plt.tight_layout()
     plt.show()
 
+Creating and sampling a Czjzek distribution in polar coordinates
+''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+
+The :py:class:`~mrsimulator.models.czjzek.CzjzekDistribution` class also supports sampling tensors in polar coordinates. The logic behind transforming from a :math:`\zeta`-:math:`\eta` Cartesian grid is further described in mrinversion (cite), and the following definitions are used
+
+.. math::
+
+    \begin{split}r_\zeta = \left| \zeta \right| ~~~~\text{and}~~~~
+    \theta = \left\{ \begin{array}{l r}
+                \frac{\pi}{4} \eta      &: \zeta \le 0, \\
+                \frac{\pi}{2} \left(1 - \frac{\eta}{2} \right) &: \zeta > 0.
+            \end{array}
+            \right.\end{split}
+
+Because Cartesian grids are more manageable in computation, the above polar piece-wise grid is re-express as the x-y Cartesian grid following,
+
+.. math::
+
+    x = r_\zeta \cos\theta ~~~~\text{and}~~~~ y = r_\zeta \sin\theta.
+
+Below, we create another instance of the :py:class:`~mrsimulator.models.czjzek.CzjzekDistribution`
+class with the same value of :math:`sigma=0.8`, but we now also include the argument ``polar=True``
+which means the :py:meth:`~mrsimulator.models.CzjzekDistribution.rvs` will sample x and y points.
+
+.. skip: next
+
+.. plot::
+    :context: close-figs
+    :caption: Random sampling of Czjzek distribution of shielding tensors in polar coordinates.
+
+    cz_model_polar = CzjzekDistribution(sigma=0.8, polar=True)
+
+    # Sample (x, y) points
+    x_dist, y_dist = cz_model_polar.rvs(size=50000)
+
+    # Plot the distribution
+    plt.figure(figsize=(4, 4))
+    plt.scatter(x_dist, y_dist, s=4, alpha=0.02)
+    plt.xlabel("$x$ / ppm")
+    plt.ylabel("$y$ / ppm")
+    plt.xlim(0, 8)
+    plt.ylim(0, 8)
+    plt.tight_layout()
+    plt.show()
+
+
 ----
 
-Czjzek distribution of symmetric quadrupolar tensors
-''''''''''''''''''''''''''''''''''''''''''''''''''''
+Generating probability distribution functions from a Czjzek model
+'''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+
+The :py:meth:`~mrsimulator.models.CzjzekDistribution.pdf` instance method will generate a
+probability distribution function on the supplied grid using the analytical function defined above.
+The provided grid -- passed to the ``pos`` keyword argument -- needs to be defined in either
+Cartesian or polar coordinates depending on if the
+:py:attr:`~mrsimulator.models.CzjzekDistribution.polar` attribute is ``True`` or ``False``.
 
-The Czjzek distribution of symmetric quadrupolar tensors follows a similar setup as the
-Czjzek distribution of symmetric shielding tensors, except we assign the outputs to Cq
-and :math:`\eta_q`. In the following example, we generate the probability distribution
-function using the :py:meth:`~mrsimulator.models.CzjzekDistribution.pdf` method.
+Below, we generate and plot a probability distribution on a :math:`\zeta`-:math:`\eta` Cartesian
+grid where ``zeta_range`` and ``eta_range`` define the desired coordinates in each dimension of the
+grid system.
 
 .. plot::
     :context: close-figs
 
     import numpy as np
 
-    Cq_range = np.arange(100) * 0.3 - 15  # pre-defined Cq range in MHz.
-    eta_range = np.arange(21) / 20  # pre-defined eta range.
-    Cq, eta, amp = cz_model.pdf(pos=[Cq_range, eta_range])
+    cz_model = CzjzekDistribution(sigma=1.2, polar=False)  # sample in (zeta, eta)
 
-To generate a probability distribution, we need to define a grid system over which the
-distribution probabilities will be evaluated. We do so by defining the range of coordinates
-along the two dimensions. In the above example, ``Cq_range`` and ``eta_range`` are the
-range of :math:`\text{Cq}` and :math:`\eta_q` coordinates, which is then given as the
-argument to the :py:meth:`~mrsimulator.models.CzjzekDistribution.pdf` method. The output
-``Cq``, ``eta``, and ``amp`` hold the two coordinates and amplitude, respectively.
+    zeta_range = np.linspace(-12, 12, num=200)  # pre-defined zeta range in units of ppm
+    eta_range = np.linspace(0, 1, num=50)       # pre-defined eta range.
+    zeta_grid, eta_grid, amp = cz_model.pdf(pos=[zeta_range, eta_range])
 
-The plot of the Czjzek probability distribution is shown below.
+Here, ``zeta_grid`` and ``eta_grid`` are numpy arrays defining a set of pair-wise points on the
+grid system, and ``amp`` is another numpy array holding the probability density at each point
+on the grid. Below, the distribution is plotted
 
 .. skip: next
 
 .. plot::
     :context: close-figs
-    :caption: Czjzek Distribution of EFG tensors.
+    :caption: Czjzek Distribution of shielding tensors.
 
-    import matplotlib.pyplot as plt
 
-    plt.contourf(Cq, eta, amp, levels=10)
-    plt.xlabel("$C_q$ / MHz")
+    plt.contourf(zeta_grid, eta_grid, amp, levels=10)
+    plt.xlabel("$\zeta$ / ppm")
     plt.ylabel("$\eta$")
     plt.tight_layout()
     plt.show()
 
-.. note::
-    The ``pdf`` method of the instance generates the probability distribution function
-    by first drawing random points from the distribution and then binning it
-    onto a pre-defined grid.
+---
+
+The probability distribution function can also be generated in polar coordinates. The workflow
+is the same, except we now define an (x, y) grid system using the variables ``x_range``
+and ``y_range``. The code to generate and plot the polar Czjzek distribution is shown below.
+
+.. skip: next
+
+.. plot::
+    :context: close-figs
+    :caption: Czjzek Distribution of shielding tensors in polar coordinates.
+
+    cz_model_polar = CzjzekDistribution(sigma=1.2, polar=True)  # sample in (x, y)
+
+    x_range = np.linspace(0, 10, num=150)
+    y_range = np.linspace(0, 10, num=150)
+    x_grid, y_grid, amp = cz_model_polar.pdf(pos=[x_range, y_range])
+
+    plt.figure(figsize=(4, 4))
+    plt.contourf(x_grid, y_grid, amp, levels=10)
+    plt.xlabel("$x$ / ppm")
+    plt.ylabel("$y$ / ppm")
+    plt.tight_layout()
+    plt.show()
+
+
+Distributions of shielding and quadrupolar tensors and a note on units
+''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+
+The :py:class:`~mrsimulator.models.CzjzekDistribution` class can be used to generate
+distributions for both symmetric chemical shielding tensors and electric field gradient
+tensors. It is important to note the Czjzek model is only aware of the Haberlen components
+of the tensors and not the units of the tensor. In the above code cells, we generated
+distributions for symmetric shielding tensors and assumed all units for :math:`\zeta` were
+in ppm.
+
+Quadrupolar tensors, defined using values of :math:`C_q` in MHz and unitless :math:`\eta`,
+can also be drawn from the Czjzek distribution in the same manner; however, the dimensions
+are assumed to be in units of MHz. The following code draws a distribution of quadrupolar
+tensor parameters.
+
+.. skip: next
+
+.. plot::
+    :context: close-figs
+
+    Cq_range = np.linspace(-12, 12, num=200)  # pre-defined Cq range in units of MHz
+    eta_range = np.linspace(0, 1, num=50)     # pre-defined eta range.
+    Cq_grid, eta_grid, amp = cz_model.pdf(pos=[Cq_range, eta_range])
+
+
+the units for ``Cq_range`` and ``Cq_grid`` are assumed in MHz. Similarly, x and y are assumed to
+be in units of MHz when sampling quadrupolar tensors in polar coordinates.
+
+.. skip: next
+
+.. plot::
+    :context: close-figs
+
+    x_range = np.linspace(0, 10, num=150)  # pre-defined x grid in units of MHz
+    y_range = np.linspace(0, 10, num=150)  # pre-defined y grid in units of MHz
+    x_grid, y_grid, amp = cz_model_polar.pdf(pos=[x_range, y_range])
+
+Generating a list of SpinSystem objects from a Czjzek model
+'''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
+
+The utility function :py:meth:`~mrsimulator.utils.collection.single_site_system_generator`, further
+described in :ref:`single_site_system_generator_documentation`, can be used in conjunction with
+the :py:class:`~mrsimulator.models.CzjzekDistribution` class to generate a set of spin systems whose
+tensor parameters follow the Czjzek distribution.
+
+.. plot::
+    :context: close-figs
+
+    from mrsimulator.utils.collection import single_site_system_generator
+
+    # Distribution of quadrupolar tensors
+    cz_model = CzjzekDistribution(sigma=0.7)
+    Cq_range = np.linspace(0, 10, num=100)
+    eta_range = np.linspace(0, 1, num=50)
+
+    # Create (Cq, eta) grid points and amplitude
+    Cq_grid, eta_grid = np.meshgrid(Cq_range, eta_range)
+    _, _, amp = cz_model.pdf(pos=[Cq_range, eta_range])
+
+    sys = single_site_system_generator(
+        isotope="27Al",
+        quadrupolar={"Cq": Cq_grid * 1e6, "eta": eta_grid},  # Cq argument in units of Hz
+        abundance=amp,
+    )
+
+A spin system will be generated for each point on the :math:`\zeta`-:math:`\eta` grid, and the
+abundance of each spin system matches the amplitude of the Czjzek distribution. When working in
+polar coordinates, the set of :math:`\left(x, y\right)` coordinates needs to be transformed into
+a set of :math:`\left(\zeta, \eta\right)` coordinates before being passed to the
+:py:meth:`~mrsimulator.utils.collection.single_site_system_generator` function. The utility
+function :py:meth:`~mrsimulator.utils.x_y_to_zeta_eta` performs this transformation, as shown
+below.
+
+.. plot::
+    :context: close-figs
+
+    from mrsimulator.models.utils import x_y_to_zeta_eta
+
+    # Sample distribution of shielding tensors in polar coords
+    cz_model_polar = CzjzekDistribution(sigma=0.7, polar=True)
+    x_range = np.linspace(0, 10, num=150)
+    y_range = np.linspace(0, 10, num=150)
+
+    # Create (x, y) grid points and amplitude
+    x_grid, y_grid, amp = cz_model_polar.pdf(pos=[x_range, y_range])
+
+    # To transformation (x, y) -> (zeta, eta)
+    zeta_grid, eta_grid = x_y_to_zeta_eta(x_grid, y_grid)
+
+---
 
 .. minigallery:: mrsimulator.models.CzjzekDistribution
     :add-heading: Mini-gallery using the Czjzek distributions