From 12df12bbec5e2237f00b9d3184cd98532b8ca6fd Mon Sep 17 00:00:00 2001 From: Matthew Giammar <43814525+mgiammar@users.noreply.github.com> Date: Sun, 9 Jul 2023 19:38:51 -0400 Subject: [PATCH 01/21] Use analytical expression in Czjzek distribution --- CHANGELOG | 3 ++ src/mrsimulator/models/czjzek.py | 53 +++++++++++++++++-- .../tests/test_czjzek_and_extended_czjzek.py | 6 +-- .../models/tests/test_model_utils.py | 6 +-- src/mrsimulator/models/utils.py | 11 ++-- 5 files changed, 66 insertions(+), 13 deletions(-) diff --git a/CHANGELOG b/CHANGELOG index d3553d399..0f1a5f0be 100644 --- a/CHANGELOG +++ b/CHANGELOG @@ -17,6 +17,9 @@ What's new - 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. + v0.7.0 ------ diff --git a/src/mrsimulator/models/czjzek.py b/src/mrsimulator/models/czjzek.py index bcdd3140f..f46e479b5 100644 --- a/src/mrsimulator/models/czjzek.py +++ b/src/mrsimulator/models/czjzek.py @@ -3,12 +3,34 @@ from .utils import get_Haeberlen_components from .utils import get_principal_components -from .utils import x_y_from_zeta_eta +from .utils import zeta_eta_to_x_y +from .utils import x_y_to_zeta_eta + __author__ = "Deepansh J. Srivastava" __email__ = "srivastava.89@osu.edu" +def _analytical_czjzek_pdf(zeta, eta, sigma): + """Computes the probability density on a (zeta, eta) grid point for a Czjzek + distribution for a given value of sigma. + + Arguments: + (float) zeta: Zeta value of probability to calculate + (float) eta: Eta value of probability to calculate + (float) sigma: The size of the noise for the Czjzek distribution + + Returns: + The normalized probability density function as a numpy array + """ + denom = (2 * np.pi) ** 0.5 * sigma**5 + res = (zeta**4 * eta) * (1 - eta**2 / 9) / denom + res *= np.exp(-(zeta**2 * (1 + (eta**2 / 3))) / (2 * sigma**2)) + res /= res.sum() # Normalize total probability to 1 + + return res + + def _czjzek_random_distribution_tensors(sigma, n): r"""Czjzek random distribution model. @@ -158,8 +180,33 @@ def rvs(self, size: int): tensors = _czjzek_random_distribution_tensors(self.sigma, size) if not self.polar: return get_Haeberlen_components(tensors) - return x_y_from_zeta_eta(*get_Haeberlen_components(tensors)) + return zeta_eta_to_x_y(*get_Haeberlen_components(tensors)) + + def pdf(self, pos, **kwargs): + """Overloaded function which uses the analytical expression for the Czjzek + distribution to calculate the probability density on the given grid. + Additional arguments such as size are ignored. + + Arguments: + (tuple) pos: Coordinates along the two dimensions given as numpy arrays to + calculate the probability density on + + Returns: + The me probability density on the grid given as a numpy array + """ + _sigma = 2 * self.sigma + VV, ee = np.meshgrid(pos[0], pos[1]) + + # If the polar attribute is true, pos is assumed to be (x, y) grid and must be + # converted to (zeta, eta) before passing to the analytical pdf function + if self.polar: + VV, ee = x_y_to_zeta_eta(x=VV, y=ee) + + amp = _analytical_czjzek_pdf(zeta=VV, eta=ee, sigma=_sigma) + amp = amp.reshape((pos[1].size, pos[0].size)) # Reshape into 2D grid array + # Meshgrid called again to handle the polar and cartesian case + return *np.meshgrid(pos[0], pos[1]), amp class ExtCzjzekDistribution(AbstractDistribution): r"""An extended Czjzek distribution distribution model. @@ -249,4 +296,4 @@ def rvs(self, size: int): if not self.polar: return get_Haeberlen_components(total_tensors) - return x_y_from_zeta_eta(*get_Haeberlen_components(total_tensors)) + return zeta_eta_to_x_y(*get_Haeberlen_components(total_tensors)) diff --git a/src/mrsimulator/models/tests/test_czjzek_and_extended_czjzek.py b/src/mrsimulator/models/tests/test_czjzek_and_extended_czjzek.py index 405aeb341..5b2ec4024 100644 --- a/src/mrsimulator/models/tests/test_czjzek_and_extended_czjzek.py +++ b/src/mrsimulator/models/tests/test_czjzek_and_extended_czjzek.py @@ -3,7 +3,7 @@ import numpy as np from mrsimulator.models import CzjzekDistribution from mrsimulator.models import ExtCzjzekDistribution -from mrsimulator.models.utils import x_y_from_zeta_eta +from mrsimulator.models.utils import zeta_eta_to_x_y from mrsimulator.models.utils import x_y_to_zeta_eta __author__ = "Deepansh J. Srivastava" @@ -29,7 +29,7 @@ def test_extended_czjzek_eta_distribution_1(): def test_extended_czjzek_polar(): S0 = {"zeta": 1, "eta": 0.1} x, y = ExtCzjzekDistribution(S0, eps=0.05, polar=True).rvs(size=COUNT) - x1, y1 = x_y_from_zeta_eta(*x_y_to_zeta_eta(x, y)) + x1, y1 = zeta_eta_to_x_y(*x_y_to_zeta_eta(x, y)) np.testing.assert_almost_equal(x, x1) np.testing.assert_almost_equal(y, y1) @@ -118,6 +118,6 @@ def test_czjzek_pdf(): def test_czjzek_polar(): x, y = CzjzekDistribution(sigma=0.5, polar=True).rvs(size=COUNT) - x1, y1 = x_y_from_zeta_eta(*x_y_to_zeta_eta(x, y)) + x1, y1 = zeta_eta_to_x_y(*x_y_to_zeta_eta(x, y)) np.testing.assert_almost_equal(x, x1) np.testing.assert_almost_equal(y, y1) diff --git a/src/mrsimulator/models/tests/test_model_utils.py b/src/mrsimulator/models/tests/test_model_utils.py index 570a372e4..1857b911b 100644 --- a/src/mrsimulator/models/tests/test_model_utils.py +++ b/src/mrsimulator/models/tests/test_model_utils.py @@ -1,7 +1,7 @@ import numpy as np from mrsimulator.models.utils import get_Haeberlen_components from mrsimulator.models.utils import get_principal_components -from mrsimulator.models.utils import x_y_from_zeta_eta +from mrsimulator.models.utils import zeta_eta_to_x_y from mrsimulator.models.utils import x_y_to_zeta_eta @@ -49,10 +49,10 @@ def test_get_Haeberlen_components(): assert components == (30, 0.5) -def test_x_y_from_zeta_eta(): +def test_zeta_eta_to_x_y(): zeta = np.random.rand(20) * 40 eta = np.random.rand(20) - x, y = x_y_from_zeta_eta(zeta, eta) + x, y = zeta_eta_to_x_y(zeta, eta) theta = np.pi * eta / 4.0 x_ = zeta * np.sin(theta) diff --git a/src/mrsimulator/models/utils.py b/src/mrsimulator/models/utils.py index 780323453..fc743092d 100644 --- a/src/mrsimulator/models/utils.py +++ b/src/mrsimulator/models/utils.py @@ -41,9 +41,10 @@ def get_Haeberlen_components(tensors): return zeta, eta -def x_y_from_zeta_eta(zeta, eta): - """Convert the zeta, eta coordinates from the Haeberlen convention to the - x-y notation.""" +def zeta_eta_to_x_y(zeta, eta): + """Convert a set of (zeta, eta) coordinates from the Haeberlen convention to a set + of (x, y) coordinates. + """ xa = np.empty(zeta.size) ya = np.empty(zeta.size) @@ -64,7 +65,9 @@ def x_y_from_zeta_eta(zeta, eta): def x_y_to_zeta_eta(x, y): - """Same as def x_y_to_zeta_eta, but for ndarrays.""" + """Convert a set of (x, y) coordinates defined by two numpy arrays into equivalent + (zeta, eta) coordinates. + """ x = np.abs(x) y = np.abs(y) zeta = np.sqrt(x**2 + y**2) # + offset From 87a8e9483c8782595205fa9ac4649e945abd89fe Mon Sep 17 00:00:00 2001 From: Matthew Giammar <43814525+mgiammar@users.noreply.github.com> Date: Mon, 10 Jul 2023 11:54:16 -0400 Subject: [PATCH 02/21] spectral fitting funcs for Czjzek and ext Czjzek --- CHANGELOG | 1 + src/mrsimulator/models/czjzek.py | 56 ++- src/mrsimulator/utils/spectral_fitting.py | 503 +++++++++++++++++++++- 3 files changed, 556 insertions(+), 4 deletions(-) diff --git a/CHANGELOG b/CHANGELOG index 0f1a5f0be..79d32ebb9 100644 --- a/CHANGELOG +++ b/CHANGELOG @@ -12,6 +12,7 @@ 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. **Simulator** diff --git a/src/mrsimulator/models/czjzek.py b/src/mrsimulator/models/czjzek.py index f46e479b5..c29af0a0a 100644 --- a/src/mrsimulator/models/czjzek.py +++ b/src/mrsimulator/models/czjzek.py @@ -3,8 +3,8 @@ from .utils import get_Haeberlen_components from .utils import get_principal_components -from .utils import zeta_eta_to_x_y from .utils import x_y_to_zeta_eta +from .utils import zeta_eta_to_x_y __author__ = "Deepansh J. Srivastava" @@ -14,7 +14,7 @@ def _analytical_czjzek_pdf(zeta, eta, sigma): """Computes the probability density on a (zeta, eta) grid point for a Czjzek distribution for a given value of sigma. - + Arguments: (float) zeta: Zeta value of probability to calculate (float) eta: Eta value of probability to calculate @@ -31,6 +31,57 @@ def _analytical_czjzek_pdf(zeta, eta, sigma): return res +def _extended_czjzek_pdf_from_tensors(pos, tensors, zeta0, eta0, epsilon, polar): + """Takes in a list of random noise tensors along with the parameters for + a given Extended Czjzek distribution, applies the requisite math to + the tensors, then diagonalizes the tensors to get the (zeta, eta) + distribution. This allows the same sampling points to be drawn between different + minimization steps + + Arguments: + (tuple) pos: Two np.ndarrays defining the grid on which to calculate the pdf + (np.ndarray) tensors: A np.ndarray with shape (n, 3, 3) representing + the random tensors in the Extended Czjzek model + (float) zeta0: The zeta value of the central tensor + (float) eta0: The eta value of the central tensor + (float) epsilon: The noise parameter for the Extended Czjzek distribution + (bool) polar: Weather the distribution should be sampled in polar coordinates. + + Returns: + (np.ndarray, np.ndarray) of the zeta and eta distributions, respectively + """ + T0 = np.asarray([zeta0 * (eta0 - 1) / 2, -zeta0 * (eta0 + 1) / 2, zeta0]) + + norm_T0 = np.linalg.norm(T0) + rho = epsilon * norm_T0 / np.sqrt(30) + + if rho != 0: + ext_tensors = (rho * tensors) + np.diag(T0) + else: + ext_tensors = tensors.copy() + + zeta_dist, eta_dist = get_Haeberlen_components(ext_tensors) + if polar: + zeta_dist, eta_dist = zeta_eta_to_x_y(zeta_dist, eta_dist) + + delta_0 = (pos[0][1] - pos[0][0]) / 2 + delta_1 = (pos[1][1] - pos[1][0]) / 2 + pts_0 = [pos[0][0] - delta_0, pos[0][-1] + delta_0] + pts_1 = [pos[1][0] - delta_1, pos[1][-1] + delta_1] + + size_0 = pos[0].size + size_1 = pos[1].size + + hist, _, _ = np.histogram2d( + zeta_dist, eta_dist, bins=[size_0, size_1], range=[pts_0, pts_1] + ) + + hist /= hist.sum() + xx, yy = np.meshgrid(pos[0], pos[1]) + + return xx, yy, hist.T + + def _czjzek_random_distribution_tensors(sigma, n): r"""Czjzek random distribution model. @@ -208,6 +259,7 @@ def pdf(self, pos, **kwargs): # Meshgrid called again to handle the polar and cartesian case return *np.meshgrid(pos[0], pos[1]), amp + class ExtCzjzekDistribution(AbstractDistribution): r"""An extended Czjzek distribution distribution model. diff --git a/src/mrsimulator/utils/spectral_fitting.py b/src/mrsimulator/utils/spectral_fitting.py index 1aeb1e8d8..2fea1315f 100644 --- a/src/mrsimulator/utils/spectral_fitting.py +++ b/src/mrsimulator/utils/spectral_fitting.py @@ -1,7 +1,11 @@ +import csdmpy as cp import mrsimulator.signal_processor as sp import numpy as np from lmfit import Parameters from mrsimulator import Simulator +from mrsimulator.models.czjzek import _extended_czjzek_pdf_from_tensors +from mrsimulator.models.czjzek import CzjzekDistribution +from mrsimulator.spin_system.tensors import SymmetricTensor __author__ = ["Maxwell C Venetos", "Deepansh Srivastava"] __email__ = ["maxvenetos@gmail.com", "srivastava.89@osu.edu"] @@ -142,7 +146,7 @@ def _traverse_dictionaries(instance, parent="spin_systems"): return [] -def _post_sim_LMFIT_params(params, process, index): +def _make_params_single_processor(params, process, index): """Creates a LMFIT Parameters object for SignalProcessor operations involved in spectrum fitting. @@ -179,7 +183,10 @@ def make_signal_processor_params(processors: list): raise ValueError("Expecting a list of `SignalProcessor` objects.") params = Parameters() - _ = [_post_sim_LMFIT_params(params, obj, i) for i, obj in enumerate(processors)] + _ = [ + _make_params_single_processor(params, proc, i) + for i, proc in enumerate(processors) + ] return params @@ -520,3 +527,495 @@ def residuals(sim: Simulator, processors: list = None): res.y[0].components[0] *= -1 return residual_ + + +def _apply_iso_shift(csdm_obj, iso_shift_ppm, larmor_freq_Hz): + """Apply isotropic chemical shift to a CSDM object using the FFT shift theorem.""" + csdm_obj = csdm_obj.fft() + time_coords = csdm_obj.x[0].coordinates.value + iso_shift_Hz = larmor_freq_Hz * iso_shift_ppm + csdm_obj.y[0].components[0] *= np.exp(-np.pi * 2j * time_coords * iso_shift_Hz) + csdm_obj = csdm_obj.fft() + + return csdm_obj + + +def make_LMFIT_params_Czjzek( + cz_models: list, + weights: list = None, + iso_shifts: list = None, + processor: sp.SignalProcessor = None, +) -> Parameters: + """Generates the LMfit Parameters object with the proper keywords for a Czjzek + distribution. Each distribution has the following parameters: + - czjzek_i_sigma + - czjzek_i_iso_shift + - czjzek_i_weight + + where *i* refers to the :math:i^\text{th} Czjzek distribution + + Arguments: + (list) cz_models: A list of :class:`~mrsimulator.models.CzjzekDistribution` + objects used to set initial values. + (list) weights: An optional list of float values defining the relative weight of + each Czjzek distribution. The default is (1 / n_dist), that is, all + distributions have the same weight + (list) iso_shifts: An optional list of float values defining the central + isotropic chemical shift for each distribution. + (sp.SignalProcessor) processor: An optional signal processor to apply + post-simulation signal processing to each guess spectrum. + + Returns: + A LMfit Parameters object + """ + n_dist = len(cz_models) + iso_shifts = [0] * n_dist if iso_shifts is None else iso_shifts + weights = np.ones(n_dist) if weights is None else np.asarray(weights) + weights /= weights.sum() # Normalize weights array to total of 1 + + params = Parameters() + for i in range(n_dist): + params.add(f"czjzek_{i}_sigma", value=cz_models[i].sigma, min=0) + params.add(f"czjzek_{i}_iso_shift", value=iso_shifts[i]) + params.add(f"czjzek_{i}_weight", value=weights[i], min=0, max=1) + + # Set last weight parameter as a non-free variable + expr = "-".join([f"czjzek_{i}_weight" for i in range(n_dist - 1)]) + expr = "1" if expr == "" else f"1-{expr}" + params[f"czjzek_{n_dist-1}_weight"].vary = False + params[f"czjzek_{n_dist-1}_weight"].expr = expr + + # Add SignalProcessor parameters, if requested + if processor is not None: + _make_params_single_processor(params, processor, 0) + + return params + + +def _make_spectrum_from_Czjzek_distribution_parameters( + params: Parameters, + exp_spectrum: cp.CSDM, + pos: tuple, + polar: bool, + kernel: np.ndarray, + n_dist: int, + larmor_freq_Hz: float, + processor: sp.SignalProcessor = None, +) -> cp.CSDM: + """Helper function for generating a spectrum from a set of LMfit Parameters and + and arguments for defining the grid, kernel, etc. The other functions used in least- + squares minimization use this function to reduce code overlap. + + Arguments: + (Parameters) params: The LMfit parameters object holding parameters used during + the least-squares minimization. + (cp.CSDM) exp_spectrum: The experimental spectrum to fit to. + (tuple) pos: A tuple of two np.ndarray objects defining the grid on which to + sample the distribution. + (bool) polar: True if the sample grid is in polar coordinates. False if the grid + is defined using the Haberlen components. + (np.ndarray) The pre-computed lineshape kernel. The kernel needs to be defined + on the same grid defined by pos. + (int) n_dist: The number of Czjzek distributions to fit for. + (float) larmor_freq_Hz: This value is used in conjunction with the FFT shift + theorem to apply an isotropic chemical shift to each distribution. + (sp.SignalProcessor) processor: A + :py:class:~`mrsimulator.signal_processor.Processor` object used to apply + post-simulation signal processing to the resulting spectrum. + TODO: expand processor to apply to individual distributions (as a list) + + Returns: + Guess spectrum as a cp.CSDM object + + """ + guess_spectrum = exp_spectrum.copy() + guess_spectrum.y[0].components[:] = 0 # Initialize guess spectrum with zeros + + for i in range(n_dist): + _, _, amp = CzjzekDistribution( + sigma=params[f"czjzek_{i}_sigma"], polar=polar + ).pdf(pos) + + # Dot amplitude with kernel, then package as CSDM object + spec_tmp = guess_spectrum.copy() + spec_tmp.y[0].components[0] = np.dot(amp.flatten(), kernel) + + # Apply isotropic chemical shift to distribution using FFT shift theorem + spec_tmp = _apply_iso_shift( + csdm_obj=spec_tmp, + iso_shift_ppm=params[f"czjzek_{i}_iso_shift"].value, + larmor_freq_Hz=larmor_freq_Hz, + ) + + guess_spectrum += spec_tmp.real * params[f"czjzek_{i}_weight"].value + + if processor is not None: + _update_processors_from_LMFIT_params(params, [processor]) + guess_spectrum = processor.apply_operations(dataset=guess_spectrum).real + + return guess_spectrum + + +def LMFIT_min_function_Czjzek( + params: Parameters, + exp_spectrum: cp.CSDM, + pos: tuple, + polar: bool, + kernel: np.ndarray, + n_dist: int, + larmor_freq_Hz: float, + processor: sp.SignalProcessor = None, +) -> np.ndarray: + """The minimization routine for fitting a set of Czjzek models to an experimental + spectrum. + + Arguments: + (Parameters) params: The LMfit parameters object holding parameters used during + the least-squares minimization. + (cp.CSDM) exp_spectrum: The experimental spectrum to fit to. + (tuple) pos: A tuple of two np.ndarray objects defining the grid on which to + sample the distribution. + (bool) polar: True if the sample grid is in polar coordinates. False if the grid + is defined using the Haberlen components. + (np.ndarray) The pre-computed lineshape kernel. The kernel needs to be defined + on the same grid defined by pos. + (int) n_dist: The number of Czjzek distributions to fit for. + (float) larmor_freq_Hz: This value is used in conjunction with the FFT shift + theorem to apply an isotropic chemical shift to each distribution. + (sp.SignalProcessor) processor: A + :py:class:~`mrsimulator.signal_processor.Processor` object used to apply + post-simulation signal processing to the resulting spectrum. + TODO: expand processor to apply to individual distributions (as a list) + + Returns: + A residual array as a numpy array. + """ + guess_spectrum = _make_spectrum_from_Czjzek_distribution_parameters( + params, + exp_spectrum, + pos, + polar, + kernel, + n_dist, + larmor_freq_Hz, + processor, + ) + + return (exp_spectrum - guess_spectrum).y[0].components[0] + + +def bestfit_Czjzek( + params: Parameters, + exp_spectrum: cp.CSDM, + pos: tuple, + polar: bool, + kernel: np.ndarray, + n_dist: int, + larmor_freq_Hz: float, + processor: sp.SignalProcessor = None, +) -> cp.CSDM: + """Returns the best-fit spectrum as a CSDM object""" + return _make_spectrum_from_Czjzek_distribution_parameters( + params, + exp_spectrum, + pos, + polar, + kernel, + n_dist, + larmor_freq_Hz, + processor, + ) + + +def residuals_Czjzek( + params: Parameters, + exp_spectrum: cp.CSDM, + pos: tuple, + polar: bool, + kernel: np.ndarray, + n_dist: int, + larmor_freq_Hz: float, + processor: sp.SignalProcessor = None, +) -> cp.CSDM: + """Returns the residuals spectrum as a CSDM object""" + bestfit = _make_spectrum_from_Czjzek_distribution_parameters( + params, + exp_spectrum, + pos, + polar, + kernel, + n_dist, + larmor_freq_Hz, + processor, + ) + + return exp_spectrum - bestfit + + +def make_LMFIT_params_extended_Czjzek( + ext_cz_models: list, + weights: list = None, + iso_shifts: list = None, + processor: sp.SignalProcessor = None, + tensor_type: str = "shielding", +) -> Parameters: + """Generates the LMfit Parameters object with the proper keywords for an Extended + Czjzek distribution. Each distribution has the following parameters: + - ext_czjzek_i_zeta0 *or* ext_czjzek_i_Cq0 + - ext_czjzek_i_eta0 + - ext_czjzek_i_epsilon + - ext_czjzek_i_iso_shift + - ext_czjzek_i_weight + + where *i* refers to the :math:i^\text{th} Extended Czjzek distribution. Note that + Cq values are assumed to be in units of MHz. + + Arguments: + (list) ext_cz_models: A list of + :class:`~mrsimulator.models.ExtCzjzekDistribution` objects used to set + initial values. + (list) weights: An optional list of float values defining the relative weight of + each Czjzek distribution. The default is (1 / n_dist), that is, all + distributions have the same weight + (list) iso_shifts: An optional list of float values defining the central + isotropic chemical shift for each distribution. + (sp.SignalProcessor) processor: An optional signal processor to apply + post-simulation signal processing to each guess spectrum. + (str) tensor_type: A string literal describing if the Extended Czjzek models + define a distribution of symmetric shielding or quadrupolar tensors. The + allowed values are `shielding` and `quadrupolar`. + + Returns: + A LMfit Parameters object + """ + if tensor_type not in {"shielding", "quadrupolar"}: + raise ValueError(f"Unrecognized value of {tensor_type} for `tensor_type.") + + n_dist = len(ext_cz_models) + iso_shifts = [0] * n_dist if iso_shifts is None else iso_shifts + weights = np.ones(n_dist) if weights is None else np.asarray(weights) + weights /= weights.sum() # Normalize weights array to total of 1 + + params = Parameters() + for i in range(n_dist): + dominant_tensor = ext_cz_models[i].symmetric_tensor + if isinstance(dominant_tensor, dict): # Convert to a SymmetricTensor object + dominant_tensor = SymmetricTensor(**dominant_tensor) + + if tensor_type == "shielding": + params.add(f"ext_czjzek_{i}_zeta0", value=dominant_tensor.zeta) + elif tensor_type == "quadrupolar": + params.add(f"ext_czjzek_{i}_Cq0", value=dominant_tensor.Cq) + + params.add(f"ext_czjzek_{i}_eta0", value=dominant_tensor.eta, min=0, max=1) + params.add(f"ext_czjzek_{i}_epsilon", value=ext_cz_models[i].eps, min=0) + params.add(f"ext_czjzek_{i}_iso_shift", value=iso_shifts[i]) + params.add(f"ext_czjzek_{i}_weight", value=weights[i], min=0, max=1) + + # Set last weight parameter as a non-free variable + expr = "-".join([f"ext_czjzek_{i}_weight" for i in range(n_dist - 1)]) + expr = "1" if expr == "" else f"1-{expr}" + params[f"ext_czjzek_{n_dist-1}_weight"].vary = False + params[f"ext_czjzek_{n_dist-1}_weight"].expr = expr + + # Add SignalProcessor parameters, if requested + if processor is not None: + _make_params_single_processor(params, processor, 0) + + return params + + +def _make_spectrum_from_extended_Czjzek_distribution_parameters( + params: Parameters, + exp_spectrum: cp.CSDM, + pos: tuple, + polar: bool, + kernel: np.ndarray, + tensors: np.ndarray, + n_dist: int, + larmor_freq_Hz: float, + processor: sp.SignalProcessor = None, + tensor_type: str = "shielding", +) -> cp.CSDM: + """Helper function for generating a spectrum from a set of LMfit Parameters and + and arguments for defining the grid, kernel, etc. The other functions used in least- + squares minimization use this function to reduce code overlap. + + Arguments: + (Parameters) params: The LMfit parameters object holding parameters used during + the least-squares minimization. + (cp.CSDM) exp_spectrum: The experimental spectrum to fit to. + (tuple) pos: A tuple of two np.ndarray objects defining the grid on which to + sample the distribution. + (bool) polar: True if the sample grid is in polar coordinates. False if the grid + is defined using the Haberlen components. + (np.ndarray) The pre-computed lineshape kernel. The kernel needs to be defined + on the same grid defined by pos. + (int) n_dist: The number of Czjzek distributions to fit for. + (float) larmor_freq_Hz: This value is used in conjunction with the FFT shift + theorem to apply an isotropic chemical shift to each distribution. + (sp.SignalProcessor) processor: A + :py:class:~`mrsimulator.signal_processor.Processor` object used to apply + post-simulation signal processing to the resulting spectrum. + TODO: expand processor to apply to individual distributions (as a list) + + Returns: + Guess spectrum as a cp.CSDM object + + """ + if tensor_type not in {"shielding", "quadrupolar"}: + raise ValueError(f"Unknown `tensor_type` value of {tensor_type}.") + + guess_spectrum = exp_spectrum.copy() + guess_spectrum.y[0].components[:] = 0 # Initialize guess spectrum with zeros + + for i in range(n_dist): + zeta0 = ( + params[f"ext_czjzek_{i}_zeta0"].value + if tensor_type == "shielding" + else params[f"ext_czjzek_{i}_Cq0"].value + ) + eta0 = params[f"ext_czjzek_{i}_eta0"].value + epsilon = params[f"ext_czjzek_{i}_epsilon"].value + + # Draw the pdf from the set of static random tensors + _, _, amp = _extended_czjzek_pdf_from_tensors( + pos=pos, + tensors=tensors, + zeta0=zeta0, + eta0=eta0, + epsilon=epsilon, + polar=polar, + ) + + # Dot amplitude with kernel, then package as CSDM object + spec_tmp = guess_spectrum.copy() + spec_tmp.y[0].components[0] = np.dot(amp.flatten(), kernel) + + # Apply isotropic chemical shift to distribution using FFT shift theorem + spec_tmp = _apply_iso_shift( + csdm_obj=spec_tmp, + iso_shift_ppm=params[f"ext_czjzek_{i}_iso_shift"].value, + larmor_freq_Hz=larmor_freq_Hz, + ) + + guess_spectrum += spec_tmp.real * params[f"ext_czjzek_{i}_weight"].value + + if processor is not None: + _update_processors_from_LMFIT_params(params, [processor]) + guess_spectrum = processor.apply_operations(dataset=guess_spectrum).real + + return guess_spectrum + + +def LMFIT_min_function_extended_Czjzek( + params: Parameters, + exp_spectrum: cp.CSDM, + pos: tuple, + polar: bool, + kernel: np.ndarray, + tensors: np.ndarray, + n_dist: int, + larmor_freq_Hz: float, + processor: sp.SignalProcessor = None, + tensor_type: str = "shielding", +) -> np.ndarray: + """The minimization routine for fitting a set of Czjzek models to an experimental + spectrum. + + Arguments: + (Parameters) params: The LMfit parameters object holding parameters used during + the least-squares minimization. + (cp.CSDM) exp_spectrum: The experimental spectrum to fit to. + (tuple) pos: A tuple of two np.ndarray objects defining the grid on which to + sample the distribution. + (bool) polar: True if the sample grid is in polar coordinates. False if the grid + is defined using the Haberlen components. + (np.ndarray) The pre-computed lineshape kernel. The kernel needs to be defined + on the same grid defined by pos. + (np.ndarray) tensors: A set of random Czjzek tensors in 5-dimensional space to + calculate the Extended Czjzek tensors from. + (int) n_dist: The number of Czjzek distributions to fit for. + (float) larmor_freq_Hz: This value is used in conjunction with the FFT shift + theorem to apply an isotropic chemical shift to each distribution. + (sp.SignalProcessor) processor: A + :py:class:~`mrsimulator.signal_processor.Processor` object used to apply + post-simulation signal processing to the resulting spectrum. + TODO: expand processor to apply to individual distributions (as a list) + (str) tensor_type: A string literal describing if the Extended Czjzek models + define a distribution of symmetric shielding or quadrupolar tensors. The + allowed values are `shielding` and `quadrupolar`. + + Returns: + A residual array as a numpy array. + """ + guess_spectrum = _make_spectrum_from_extended_Czjzek_distribution_parameters( + params, + exp_spectrum, + pos, + polar, + kernel, + tensors, + n_dist, + larmor_freq_Hz, + processor, + tensor_type, + ) + + return (exp_spectrum - guess_spectrum).y[0].components[0] + + +def bestfit_extended_Czjzek( + params: Parameters, + exp_spectrum: cp.CSDM, + pos: tuple, + polar: bool, + kernel: np.ndarray, + tensors: np.ndarray, + n_dist: int, + larmor_freq_Hz: float, + processor: sp.SignalProcessor = None, + tensor_type: str = "shielding", +) -> cp.CSDM: + """Returns the bestfit spectrum as a CSDM object""" + return _make_spectrum_from_extended_Czjzek_distribution_parameters( + params, + exp_spectrum, + pos, + polar, + kernel, + tensors, + n_dist, + larmor_freq_Hz, + processor, + tensor_type, + ) + + +def residuals_extended_Czjzek( + params: Parameters, + exp_spectrum: cp.CSDM, + pos: tuple, + polar: bool, + kernel: np.ndarray, + tensors: np.ndarray, + n_dist: int, + larmor_freq_Hz: float, + processor: sp.SignalProcessor = None, + tensor_type: str = "shielding", +) -> cp.CSDM: + """Returns the residuals spectrum as a CSDM object""" + bestfit = _make_spectrum_from_extended_Czjzek_distribution_parameters( + params, + exp_spectrum, + pos, + polar, + kernel, + tensors, + n_dist, + larmor_freq_Hz, + processor, + tensor_type, + ) + + return exp_spectrum - bestfit From de280b10392a84badd90627ab94c1f862241ae49 Mon Sep 17 00:00:00 2001 From: Matthew Giammar <43814525+mgiammar@users.noreply.github.com> Date: Tue, 11 Jul 2023 09:42:56 -0400 Subject: [PATCH 03/21] Add kernel generation method --- src/mrsimulator/models/czjzek.py | 3 +- src/mrsimulator/models/utils.py | 61 ++++++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 1 deletion(-) diff --git a/src/mrsimulator/models/czjzek.py b/src/mrsimulator/models/czjzek.py index c29af0a0a..cd798e1d8 100644 --- a/src/mrsimulator/models/czjzek.py +++ b/src/mrsimulator/models/czjzek.py @@ -257,7 +257,8 @@ def pdf(self, pos, **kwargs): amp = amp.reshape((pos[1].size, pos[0].size)) # Reshape into 2D grid array # Meshgrid called again to handle the polar and cartesian case - return *np.meshgrid(pos[0], pos[1]), amp + dim0, dim1 = np.meshgrid(pos[0], pos[1]) + return dim0, dim1, amp class ExtCzjzekDistribution(AbstractDistribution): diff --git a/src/mrsimulator/models/utils.py b/src/mrsimulator/models/utils.py index fc743092d..4b849fa8d 100644 --- a/src/mrsimulator/models/utils.py +++ b/src/mrsimulator/models/utils.py @@ -1,4 +1,9 @@ import numpy as np +from mrsimulator import Method +from mrsimulator import Simulator +from mrsimulator import Site +from mrsimulator import SpinSystem +from mrsimulator.spin_system.isotope import Isotope __author__ = "Deepansh J. Srivastava" __email__ = "srivastava.89@osu.edu" @@ -80,3 +85,59 @@ def x_y_to_zeta_eta(x, y): eta[index] = (4.0 / np.pi) * np.arctan(x[index] / y[index]) return zeta.ravel(), eta.ravel() + + +def _simulate_spectra_over_zeta_and_eta(ZZ, ee, mth, tensor_type): + """Helper function to generate the kernel""" + isotope = Isotope.parse(mth.channels[0]).symbol # Grab isotope from Method object + + spin_systems = [ + SpinSystem(sites=[Site(isotope=isotope, quadrupolar=dict(Cq=Z * 1e6, eta=e))]) + if tensor_type == "quadrupolar" + else SpinSystem( + sites=[Site(isotope=isotope, shielding_symmetric=dict(zeta=Z, eta=e))] + ) + for Z, e in zip(ZZ.flatten(), ee.flatten()) + ] + sim = Simulator(spin_systems=spin_systems, methods=[mth]) + sim.config.decompose_spectrum = "spin_system" + sim.run(pack_as_csdm=False) + + amp = sim.methods[0].simulation.real + return amp + + +def generate_lineshape_kernel( + pos: tuple, mth: Method, polar: bool, tensor_type: str = "shielding" +) -> np.ndarray: + """Pre-compute a lineshape kernel used during least-squares fitting of an + experimental spectrum. The isotope for the spin system is grabbed from the isotope + at index zero of the channels method + + Note: The lineshape kernel is currently limited to simulating the spectrum of an + uncoupled spin system for a single Method over a range of either chemical shielding + or quadrupolar tensors. Functionality may expand in the future + + Arguments: + (tuple) pos: A set of numpy arrays defining the grid space + (mrsimulator.Method) mth: The :py:class:`~mrsimulator.method.Method` used to + simulate the spectra. + (bool) polar: Weather the grid is defined in polar coordinates + (str) tensor_type: A string enumeration literal describing which type of tensor + to use in the simulation. The allowed values are `shielding` and + `quadrupolar`. + + Returns: + (np.ndarray) kernel: The simulated lineshape kernel + """ + if tensor_type not in {"shielding", "quadrupolar"}: + raise ValueError(f"Unrecognized value of {tensor_type} for `tensor_type.") + + # If in polar coordinates, then (ZZ, ee) right now is really (xx, yy) + ZZ, ee = np.meshgrid(pos[0], pos[1], indexing="xy") + + # Convert polar coords to cartesian coords + if polar: + ZZ, ee = x_y_to_zeta_eta(ZZ, ee) + + return _simulate_spectra_over_zeta_and_eta(ZZ, ee, mth, tensor_type) From 70e8c5aa040ce9118f4dcd0b232c48baa11622fe Mon Sep 17 00:00:00 2001 From: Matthew Giammar <43814525+mgiammar@users.noreply.github.com> Date: Tue, 11 Jul 2023 18:12:16 -0400 Subject: [PATCH 04/21] Add extended Czjzek fitting example --- .../1D_fitting/plot_7_27Al_Extended_Czjzek.py | 234 ++++++++++++++++++ 1 file changed, 234 insertions(+) create mode 100644 fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py diff --git a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py new file mode 100644 index 000000000..045eb22a7 --- /dev/null +++ b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py @@ -0,0 +1,234 @@ +#!/usr/bin/env python +""" +Fitting an Extended Czjzek Model +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +""" +# %% +# In this example, we display the capability of mrsimulator to fit an Extended Czjzek +# distribution of quadrupolar tensors to an experimental $^{27}\text{Al}$ MAS spectrum +# of a phospho-aluminosilicate glass. Setting up a least-squares minimization for an +# Extended Czjzek distribution is slightly more complicated than that of a crystalline +# solid. There are 4 major steps involved in the process: +# +# 1. Importing the experimental dataset +# 2. Generating a pre-computed linshape kernel for the experiment +# 3. Creating parameters for the Extended Czjzek model from an initial guess +# 4. Minimizing and visualizing. +# +# Below, we first import the requisite libraries, functions, and classes. +import numpy as np +import csdmpy as cp +import matplotlib.pyplot as plt +import mrsimulator.signal_processor as sp +import mrsimulator.utils.spectral_fitting as sf +from lmfit import Minimizer + +from mrsimulator.method.lib import BlochDecayCTSpectrum +from mrsimulator.utils import get_spectral_dimensions + +from mrsimulator.models.czjzek import ExtCzjzekDistribution +from mrsimulator.models.czjzek import _czjzek_random_distribution_tensors +from mrsimulator.models.utils import generate_lineshape_kernel + +# sphinx_gallery_thumbnail_number = 3 + +# %% +# Import the experimental dataset +# ------------------------------- +# +# Below we import and visualize the experimental dataset. Line 3 is used to set the +# maximum intensity of the spectrum to 1, but has no other bering on the minimization. +host = "http://ssnmr.org/sites/default/files/mrsimulator/" +filename = "20K_20Al_10P_50Si_HahnEcho_27Al.csdf" +exp_data = cp.load(host + filename).real +exp_data.x[0].to("ppm", "nmr_frequency_ratio") +exp_data /= exp_data.max() + +exp_data.plot() +plt.show() + +# %% +# Generating a lineshapke kernel +# ------------------------------ +# +# A spectrum arising from an Extended Czjzek distribution can be modeled by drawing a +# random set of tensors, binning those tensors into a distribution on some grid of +# tensors parameters, then simulating the spectra of spin systems whose abundances +# correspond to the amplitude of that distribution. +# +# For least-squares fitting, however, simulating the spectra of many spin systems during +# each minimization step would be computationally costly, especially since a +# least-squares minimization could take hundreds of steps to complete. +# A more efficient way would be to define a grid for the tensor parameters, simulate a +# spectrum for each point on the grid before minimization, and only update the +# probability distribution on that grid during each minimization step. +# +# Below, we create a `Method` object for simulating the spectra of the given experiment, +# define a polar grid for generating the lineshape kernel, and finally call the kernel +# generation method -- `generate_lineshape_kernel`. +# The grid we defined is in polar coordinates, so we pass `polar=True` to the kernel +# generation method; if the grid were defined for the Haeberlen components of the +# tensor, then pass `polar=False`. +# We are also interested in the distribution of the EFG (quadrupolar) tensor, hence +# `tensor_type="quadrupolar"`, but a symmetric shielding lineshape kernel can also be +# generated by passing `tensor_type="shielding"`. +# Create a Method object to simulate the spectrum +spectral_dimensions = get_spectral_dimensions(exp_data) +mth = BlochDecayCTSpectrum( + channels=["27Al"], rotor_frequency=14.2e3, spectral_dimensions=spectral_dimensions +) + +# Define a polar grid for the lineshape kernel +x = np.linspace(0, 12, num=36) +y = np.linspace(0, 12, num=36) +pos = (x, y) + +# Call the kernel generation function with the correct arguments +kernel = generate_lineshape_kernel( + pos=pos, mth=mth, polar=True, tensor_type="quadrupolar" +) + +print("Desired Kernel shape: ", (x.size * y.size, spectral_dimensions[0]["count"])) +print("Actual Kernel shape: ", kernel.shape) + +# %% +# Create a set of random noise tensors +# ------------------------------------ +# Below we sample a set of 200,000 random noise tensors to use during the fit; this set +# of tensors remain the same throughout the least-squares minimization making the fit +# more stable. +tensors = _czjzek_random_distribution_tensors(sigma=1, n=200_000) + +# %% +# Create a Parameters object +# -------------------------- +# Next, we create an instance of the `ExtCzjzekDistribution` class with initial guess +# values along with a `SignalProcessor` object. A set of minimization parameters are +# then created from this initial guess. +# +# The initial guess and residual spectrum is also plotted to judge the quality of the +# initial guess. +# +# Note that the variable `addtl_sf_kwargs` holds some additional keyword arguments that +# many of the spectral fitting function takes in. This dictionary needs to be updated to +# reflect any changes made in the minimization. + +# Create initial guess ExtCzjzekDistribution +tensor = {"Cq": 3, "eta": 0.3} +ext_cz_model = ExtCzjzekDistribution(symmetric_tensor=tensor, eps=1.6) +dist_iso_shift = 60 # in ppm + +processor = sp.SignalProcessor( + operations=[ + sp.IFFT(), + sp.apodization.Gaussian(FWHM="420 Hz"), + sp.FFT(), + sp.Scale(factor=2.5), + ] +) + +# Make the Parameters object +params = sf.make_LMFIT_params_extended_Czjzek( + ext_cz_models=[ext_cz_model], + iso_shifts=[dist_iso_shift], + processor=processor, + tensor_type="quadrupolar", +) + + +# Additional keyword arguments passed to best-fit and residual functions. +addtl_sf_kwargs = dict( + exp_spectrum=exp_data, + pos=pos, + polar=True, + kernel=kernel, + tensors=tensors, + n_dist=1, + larmor_freq_Hz=mth.channels[0].larmor_freq(B0=mth.magnetic_flux_density), + processor=processor, + tensor_type="quadrupolar", +) + +# Make a guess and residuals spectrum from the initial guess +guess = sf.bestfit_extended_Czjzek(params=params, **addtl_sf_kwargs) +residuals = sf.residuals_extended_Czjzek(params=params, **addtl_sf_kwargs) + +plt.figure(figsize=(5, 4)) +ax = plt.subplot(projection="csdm") +ax.plot(exp_data, "k", alpha=0.5, label="Experiment") +ax.plot(guess, "r", alpha=0.3, label="Guess") +ax.plot(residuals, "b", alpha=0.3, label="Residuals") +plt.legend() +plt.grid() +plt.tight_layout() +plt.title("Initial Guess") +plt.show() + +# Print the Parameters object +params + +# %% +# Create and run a minimization +# ----------------------------- +# Finally, a `Minimizer` object is created and a minimization run using least-squares. +# The same arguments defined in the `addtl_sf_kwargs` variable are also passed to the +# minimizer. +scipy_minimization_kwargs = dict( + diff_step=1e-3, # Increase step size + gtol=1e-15, # Increase global convergence requirement (default 1e-8) + xtol=1e-15, # Increase variable convergence requirement (default 1e-8) + verbose=2, # Print minimization info during each step + loss="soft_l1", +) + + +minner = Minimizer( + sf.LMFIT_min_function_extended_Czjzek, + params, + fcn_kws=addtl_sf_kwargs, + **scipy_minimization_kwargs, +) +result = minner.minimize(method="least_squares") +result + +# %% +# Plot the best-fit spectrum +# -------------------------- +bestfit = sf.bestfit_extended_Czjzek(params=result.params, **addtl_sf_kwargs) +residuals = sf.residuals_extended_Czjzek(params=result.params, **addtl_sf_kwargs) + +plt.figure(figsize=(5, 4)) +ax = plt.subplot(projection="csdm") +ax.plot(exp_data, "k", alpha=0.5, label="Experiment") +ax.plot(bestfit, "r", alpha=0.3, label="Guess") +ax.plot(residuals, "b", alpha=0.3, label="Residuals") +plt.legend() +plt.grid() +plt.tight_layout() +plt.title("Best Fit") +plt.show() + +# %% +# Plot the best-fit distribution +# ------------------------------ +# +# Note that this cell is hardcoded for a single distribution. +epsilon = result.params["ext_czjzek_0_epsilon"].value +bf_tensor = {"eta": result.params["ext_czjzek_0_eta0"].value} +if addtl_sf_kwargs["tensor_type"] == "shielding": + bf_tensor["zeta"] = result.params["ext_czjzek_0_zeta0"].value +else: + bf_tensor["Cq"] = result.params["ext_czjzek_0_Cq0"].value + +bestfit_model = ExtCzjzekDistribution( + symmetric_tensor=tensor, eps=epsilon, polar=addtl_sf_kwargs["polar"] +) + +xx, yy, amp = bestfit_model.pdf(pos=pos, size=10000000) + +plt.figure(figsize=(5, 5)) +plt.contourf(xx, yy, amp) +plt.xlabel("x / (ppm)") +plt.ylabel("y / (ppm)") +plt.tight_layout() +plt.show() From 712abae7f1c6103caaeafc24021a0131f8a2f679 Mon Sep 17 00:00:00 2001 From: Matthew Giammar <43814525+mgiammar@users.noreply.github.com> Date: Fri, 14 Jul 2023 08:38:49 -0400 Subject: [PATCH 05/21] Update distribution documentation pages --- .../spin_system_distributions/czjzek.rst | 238 +++++++++++++++--- .../extended_czjzek.rst | 65 ++++- 2 files changed, 264 insertions(+), 39 deletions(-) diff --git a/docs/user_guide/spin_system_distributions/czjzek.rst b/docs/user_guide/spin_system_distributions/czjzek.rst index d2250cf3d..0e25ec426 100644 --- a/docs/user_guide/spin_system_distributions/czjzek.rst +++ b/docs/user_guide/spin_system_distributions/czjzek.rst @@ -4,12 +4,22 @@ 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. +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 +noise 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 +31,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`. @@ -39,13 +49,15 @@ function. Let's first draw points from this distribution, using the 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 +69,216 @@ 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 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. +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``. + +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 exactly 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. + +.. 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. + +.. 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, amp = cz_model.pdf(pos=[Cq_range, eta_range]) + + sys = single_site_system_generator( + isotope="13C", + 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) + + sys = single_site_system_generator( + isotope="13C", + shielding_symmetric={"zeta": zeta_grid, "eta": eta_grid}, + abundance=amp, + ) + + +--- .. minigallery:: mrsimulator.models.CzjzekDistribution :add-heading: Mini-gallery using the Czjzek distributions diff --git a/docs/user_guide/spin_system_distributions/extended_czjzek.rst b/docs/user_guide/spin_system_distributions/extended_czjzek.rst index 6b5b81c13..7b094ed8e 100644 --- a/docs/user_guide/spin_system_distributions/extended_czjzek.rst +++ b/docs/user_guide/spin_system_distributions/extended_czjzek.rst @@ -4,8 +4,10 @@ Extended Czjzek distribution ---------------------------- The Extended Czjzek distribution models random variations of a second-rank traceless -symmetric tensors about a non-zero tensor. See :ref:`ext_czjzek_model` and -references within for a brief description of the model. +symmetric tensors about a non-zero tensor. Unlike the Czjzek distribution, the Extended +Czjzek model has no known analytical function for the probability distribution. Therefore, +mrsimulator relies on random sampling to approximate the probability distribution function. +See :ref:`ext_czjzek_model` and references within for a further description of the model. Extended Czjzek distribution of symmetric shielding tensors ''''''''''''''''''''''''''''''''''''''''''''''''''''''''''' @@ -76,20 +78,24 @@ function using the :py:meth:`~mrsimulator.models.ExtCzjzekDistribution.pdf` meth import numpy as np - Cq_range = np.arange(100) * 0.04 + 2 # pre-defined Cq range in MHz. - eta_range = np.arange(21) / 20 # pre-defined eta range. + Cq_range = np.linspace(2, 6, num=100) # pre-defined Cq range in MHz. + eta_range = np.linspace(0, 1, num=20) # pre-defined eta range. quad_tensor = {"Cq": 3.5, "eta": 0.23} # Cq assumed in MHz model_quad = ExtCzjzekDistribution(quad_tensor, eps=0.2) - Cq, eta, amp = model_quad.pdf(pos=[Cq_range, eta_range]) + Cq_grid, eta_grid, amp = model_quad.pdf(pos=[Cq_range, eta_range], size=400000) As with the case with Czjzek distribution, to generate a probability distribution of the extended Czjzek 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.ExtCzjzekDistribution.pdf` method. The output -``Cq``, ``eta``, and ``amp`` hold the two coordinates and amplitude, respectively. +argument to the :py:meth:`~mrsimulator.models.ExtCzjzekDistribution.pdf` method. The pdf +method also accepts the keyword argument ``size`` which defines the number of random samples +used to approximate the probability distribution. A larger number will create better +approximations, although this increased quality comes at the expense of computation time. +The output ``Cq_grid``, ``eta_grid``, and ``amp`` hold the two coordinates and +amplitude, respectively. The plot of the extended Czjzek probability distribution is shown below. @@ -101,12 +107,55 @@ The plot of the extended Czjzek probability distribution is shown below. import matplotlib.pyplot as plt - plt.contourf(Cq, eta, amp, levels=10) + plt.contourf(Cq_grid, eta_grid, amp, levels=10) plt.xlabel("$C_q$ / MHz") plt.ylabel("$\eta$") plt.tight_layout() plt.show() +Extended Czjzek distribution in polar coordinates +''''''''''''''''''''''''''''''''''''''''''''''''' + +As with the Czjzek distribution, we can sample an Extended Czjzek distribution on a polar +(x, y) grid. Below, we construct two equivalent +:py:class:`~mrsimulator.models.ExtCzjzekDistribution` objects, except one is defined in polar +coordinates. + +.. skip: next + +.. plot:: + :context: close-figs + :caption: Two equivalent Extended Czjzek distributions in Cartesian :math:`\left(\zeta, \eta\right)` coordinates (left) and in polar :math:`\left(x, y\right)` coordinates (right). + + quad_tensor = {"Cq": 4.2, "eta": 0.15} # Cq assumed in MHz + ext_cz_model = ExtCzjzekDistribution(quad_tensor, eps=0.4) + ext_cz_model_polar = ExtCzjzekDistribution(quad_tensor, eps=0.4, polar=True) + + # Distribution in cartesian (zeta, eta) coordinates + Cq_range = np.linspace(2, 8, num=50) + eta_range = np.linspace(0, 1, num=20) + Cq_grid, eta_grid, amp = ext_cz_model.pdf(pos=[Cq_range, eta_range], size=2000000) + + # Distribution in polar coordinates + x_range = np.linspace(0, 6, num=36) + y_range = np.linspace(0, 6, num=36) + x_grid, y_grid, amp_polar = ext_cz_model_polar.pdf(pos=[x_range, y_range], size=2000000) + + # Plot the distributions + fig, ax = plt.subplots(1, 2, figsize=(9, 4), gridspec_kw={"width_ratios": (5, 4)}) + ax[0].contourf(Cq_grid, eta_grid, amp, levels=10) + ax[0].set_xlabel("$C_q$ / MHz") + ax[0].set_ylabel("$\eta$") + ax[0].set_title("Cartesian coordinates") + ax[1].contourf(x_grid, y_grid, amp_polar, levels=10) + ax[1].set_xlabel("x / MHz") + ax[1].set_ylabel("y / MHz") + ax[1].set_title("Polar coordinates") + + 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 From 99fb5d3d67268565a5c984b39652e997ca278de4 Mon Sep 17 00:00:00 2001 From: deepanshs <21365911+deepanshs@users.noreply.github.com> Date: Sun, 17 Mar 2024 15:49:59 -0400 Subject: [PATCH 06/21] lint --- src/mrsimulator/models/czjzek.py | 1 - src/mrsimulator/models/tests/test_czjzek_and_extended_czjzek.py | 2 +- src/mrsimulator/models/tests/test_model_utils.py | 2 +- 3 files changed, 2 insertions(+), 3 deletions(-) diff --git a/src/mrsimulator/models/czjzek.py b/src/mrsimulator/models/czjzek.py index 59fb87261..3e68652bb 100644 --- a/src/mrsimulator/models/czjzek.py +++ b/src/mrsimulator/models/czjzek.py @@ -14,7 +14,6 @@ from .utils import zeta_eta_to_x_y - __author__ = "Deepansh J. Srivastava" __email__ = "srivastava.89@osu.edu" diff --git a/src/mrsimulator/models/tests/test_czjzek_and_extended_czjzek.py b/src/mrsimulator/models/tests/test_czjzek_and_extended_czjzek.py index 2cbd78bb2..5d58b7900 100644 --- a/src/mrsimulator/models/tests/test_czjzek_and_extended_czjzek.py +++ b/src/mrsimulator/models/tests/test_czjzek_and_extended_czjzek.py @@ -3,8 +3,8 @@ import numpy as np from mrsimulator.models import CzjzekDistribution from mrsimulator.models import ExtCzjzekDistribution -from mrsimulator.models.utils import zeta_eta_to_x_y from mrsimulator.models.utils import x_y_to_zeta_eta +from mrsimulator.models.utils import zeta_eta_to_x_y __author__ = "Deepansh J. Srivastava" __email__ = "srivastava.89@osu.edu" diff --git a/src/mrsimulator/models/tests/test_model_utils.py b/src/mrsimulator/models/tests/test_model_utils.py index 1857b911b..abb72b389 100644 --- a/src/mrsimulator/models/tests/test_model_utils.py +++ b/src/mrsimulator/models/tests/test_model_utils.py @@ -1,8 +1,8 @@ import numpy as np from mrsimulator.models.utils import get_Haeberlen_components from mrsimulator.models.utils import get_principal_components -from mrsimulator.models.utils import zeta_eta_to_x_y from mrsimulator.models.utils import x_y_to_zeta_eta +from mrsimulator.models.utils import zeta_eta_to_x_y __author__ = "Deepansh J. Srivastava" From 7081a964c3079ff9dd815e046f82a4821bc3ba7f Mon Sep 17 00:00:00 2001 From: deepanshs <21365911+deepanshs@users.noreply.github.com> Date: Sun, 17 Mar 2024 19:13:48 -0400 Subject: [PATCH 07/21] fix crashes --- .../spin_system_distributions/czjzek.rst | 39 +++++---------- .../extended_czjzek.rst | 8 +--- .../1D_fitting/plot_7_27Al_Extended_Czjzek.py | 2 +- .../models/analytical_distributions.py | 4 +- src/mrsimulator/models/czjzek.py | 48 ------------------- 5 files changed, 18 insertions(+), 83 deletions(-) diff --git a/docs/user_guide/spin_system_distributions/czjzek.rst b/docs/user_guide/spin_system_distributions/czjzek.rst index 0e25ec426..96622f1d9 100644 --- a/docs/user_guide/spin_system_distributions/czjzek.rst +++ b/docs/user_guide/spin_system_distributions/czjzek.rst @@ -3,19 +3,17 @@ Czjzek distribution ------------------- -The Czjzek distribution models random variations of a second-rank traceless +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 .. 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) + 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 -noise parameter. -See :ref:`czjzek_model` for a further mathematical description of the model. +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 +: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 @@ -46,12 +44,7 @@ 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. -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. +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 @@ -72,9 +65,7 @@ The scatter plot of these coordinates is shown below. 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 +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:: @@ -85,8 +76,7 @@ Cartesian grid is further described in mrinversion (cite), and the following def \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, +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:: @@ -164,7 +154,7 @@ on the grid. Below, the distribution is plotted --- The probability distribution function can also be generated in polar coordinates. The workflow -is exactly the same, except we now define an (x, y) grid system using the variables ``x_range`` +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 @@ -202,6 +192,8 @@ can also be drawn from the Czjzek distribution in the same manner; however, the are assumed to be in units of MHz. The following code draws a distribution of quadrupolar tensor parameters. +.. skip: next + .. plot:: :context: close-figs @@ -213,6 +205,8 @@ tensor parameters. 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 @@ -242,7 +236,7 @@ tensor parameters follow the Czjzek distribution. Cq_grid, eta_grid, amp = cz_model.pdf(pos=[Cq_range, eta_range]) sys = single_site_system_generator( - isotope="13C", + isotope="27Al", quadrupolar={"Cq": Cq_grid * 1e6, "eta": eta_grid}, # Cq argument in units of Hz abundance=amp, ) @@ -271,13 +265,6 @@ below. # To transformation (x, y) -> (zeta, eta) zeta_grid, eta_grid = x_y_to_zeta_eta(x_grid, y_grid) - sys = single_site_system_generator( - isotope="13C", - shielding_symmetric={"zeta": zeta_grid, "eta": eta_grid}, - abundance=amp, - ) - - --- .. minigallery:: mrsimulator.models.CzjzekDistribution diff --git a/docs/user_guide/spin_system_distributions/extended_czjzek.rst b/docs/user_guide/spin_system_distributions/extended_czjzek.rst index 7b094ed8e..5af596864 100644 --- a/docs/user_guide/spin_system_distributions/extended_czjzek.rst +++ b/docs/user_guide/spin_system_distributions/extended_czjzek.rst @@ -3,11 +3,7 @@ Extended Czjzek distribution ---------------------------- -The Extended Czjzek distribution models random variations of a second-rank traceless -symmetric tensors about a non-zero tensor. Unlike the Czjzek distribution, the Extended -Czjzek model has no known analytical function for the probability distribution. Therefore, -mrsimulator relies on random sampling to approximate the probability distribution function. -See :ref:`ext_czjzek_model` and references within for a further description of the model. +The Extended Czjzek distribution models random variations of second-rank traceless symmetric tensors about a non-zero tensor. Unlike the Czjzek distribution, the Extended Czjzek model has no known analytical function for the probability distribution. Therefore, mrsimulator relies on random sampling to approximate the probability distribution function. See :ref:`ext_czjzek_model` and references within for a further description of the model. Extended Czjzek distribution of symmetric shielding tensors ''''''''''''''''''''''''''''''''''''''''''''''''''''''''''' @@ -85,7 +81,7 @@ function using the :py:meth:`~mrsimulator.models.ExtCzjzekDistribution.pdf` meth model_quad = ExtCzjzekDistribution(quad_tensor, eps=0.2) Cq_grid, eta_grid, amp = model_quad.pdf(pos=[Cq_range, eta_range], size=400000) -As with the case with Czjzek distribution, to generate a probability distribution of the +As with the case of the Czjzek distribution, to generate a probability distribution of the extended Czjzek 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 diff --git a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py index 045eb22a7..991b91064 100644 --- a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py +++ b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py @@ -123,7 +123,7 @@ sp.IFFT(), sp.apodization.Gaussian(FWHM="420 Hz"), sp.FFT(), - sp.Scale(factor=2.5), + sp.Scale(factor=25), ] ) diff --git a/src/mrsimulator/models/analytical_distributions.py b/src/mrsimulator/models/analytical_distributions.py index f36e42e7d..43ce106a5 100644 --- a/src/mrsimulator/models/analytical_distributions.py +++ b/src/mrsimulator/models/analytical_distributions.py @@ -1,5 +1,5 @@ import numpy as np -from mrsimulator.models.utils import x_y_from_zeta_eta +from mrsimulator.models.utils import zeta_eta_to_x_y __author__ = "Deepansh J. Srivastava" @@ -73,7 +73,7 @@ def czjzek_polar(sigma: float, pos: list): pdf_model = czjzek_distribution(sigma, zeta, eta) eta = eta.ravel() zeta = zeta.ravel() - x, y = x_y_from_zeta_eta(zeta, eta) + x, y = zeta_eta_to_x_y(zeta, eta) eta_idx = np.where(eta == 1) pdf_model = pdf_model.ravel() diff --git a/src/mrsimulator/models/czjzek.py b/src/mrsimulator/models/czjzek.py index 3e68652bb..137a57922 100644 --- a/src/mrsimulator/models/czjzek.py +++ b/src/mrsimulator/models/czjzek.py @@ -10,7 +10,6 @@ from .utils import get_Haeberlen_components from .utils import get_principal_components -from .utils import x_y_to_zeta_eta from .utils import zeta_eta_to_x_y @@ -20,26 +19,6 @@ ANALYTICAL_AVAILABLE = {"czjzek": analytical_dist.czjzek} -def _analytical_czjzek_pdf(zeta, eta, sigma): - """Computes the probability density on a (zeta, eta) grid point for a Czjzek - distribution for a given value of sigma. - - Arguments: - (float) zeta: Zeta value of probability to calculate - (float) eta: Eta value of probability to calculate - (float) sigma: The size of the noise for the Czjzek distribution - - Returns: - The normalized probability density function as a numpy array - """ - denom = (2 * np.pi) ** 0.5 * sigma**5 - res = (zeta**4 * eta) * (1 - eta**2 / 9) / denom - res *= np.exp(-(zeta**2 * (1 + (eta**2 / 3))) / (2 * sigma**2)) - res /= res.sum() # Normalize total probability to 1 - - return res - - def _extended_czjzek_pdf_from_tensors(pos, tensors, zeta0, eta0, epsilon, polar): """Takes in a list of random noise tensors along with the parameters for a given Extended Czjzek distribution, applies the requisite math to @@ -251,33 +230,6 @@ def rvs(self, size: int): return get_Haeberlen_components(tensors) return zeta_eta_to_x_y(*get_Haeberlen_components(tensors)) - def pdf(self, pos, **kwargs): - """Overloaded function which uses the analytical expression for the Czjzek - distribution to calculate the probability density on the given grid. - Additional arguments such as size are ignored. - - Arguments: - (tuple) pos: Coordinates along the two dimensions given as numpy arrays to - calculate the probability density on - - Returns: - The me probability density on the grid given as a numpy array - """ - _sigma = 2 * self.sigma - VV, ee = np.meshgrid(pos[0], pos[1]) - - # If the polar attribute is true, pos is assumed to be (x, y) grid and must be - # converted to (zeta, eta) before passing to the analytical pdf function - if self.polar: - VV, ee = x_y_to_zeta_eta(x=VV, y=ee) - - amp = _analytical_czjzek_pdf(zeta=VV, eta=ee, sigma=_sigma) - amp = amp.reshape((pos[1].size, pos[0].size)) # Reshape into 2D grid array - - # Meshgrid called again to handle the polar and cartesian case - dim0, dim1 = np.meshgrid(pos[0], pos[1]) - return dim0, dim1, amp - class ExtCzjzekDistribution(AbstractDistribution): r"""An extended Czjzek distribution distribution model. From c855585c8f8d6f11a0040729dcafab59c8026767 Mon Sep 17 00:00:00 2001 From: deepanshs <21365911+deepanshs@users.noreply.github.com> Date: Mon, 18 Mar 2024 07:38:08 -0400 Subject: [PATCH 08/21] simplify code --- src/mrsimulator/models/czjzek.py | 36 ++++++++++++++++++----- src/mrsimulator/models/utils.py | 18 +++++++----- src/mrsimulator/utils/spectral_fitting.py | 2 +- 3 files changed, 40 insertions(+), 16 deletions(-) diff --git a/src/mrsimulator/models/czjzek.py b/src/mrsimulator/models/czjzek.py index 137a57922..8a538572b 100644 --- a/src/mrsimulator/models/czjzek.py +++ b/src/mrsimulator/models/czjzek.py @@ -19,6 +19,7 @@ ANALYTICAL_AVAILABLE = {"czjzek": analytical_dist.czjzek} +# this function does the same as self.pdf() but with cached tensors! def _extended_czjzek_pdf_from_tensors(pos, tensors, zeta0, eta0, epsilon, polar): """Takes in a list of random noise tensors along with the parameters for a given Extended Czjzek distribution, applies the requisite math to @@ -38,7 +39,7 @@ def _extended_czjzek_pdf_from_tensors(pos, tensors, zeta0, eta0, epsilon, polar) Returns: (np.ndarray, np.ndarray) of the zeta and eta distributions, respectively """ - T0 = np.asarray([zeta0 * (eta0 - 1) / 2, -zeta0 * (eta0 + 1) / 2, zeta0]) + T0 = get_principal_components(zeta0, eta0) norm_T0 = np.linalg.norm(T0) rho = epsilon * norm_T0 / np.sqrt(30) @@ -129,6 +130,10 @@ def _czjzek_random_distribution_tensors(sigma, n): class AbstractDistribution: + def __init__(self, cache_tensors=False): + self._cache_tensors = cache_tensors + self._tensors = None + def pdf(self, pos, size: int = 400000, analytical: bool = True): """Generates a probability distribution function by binning the random variates of length size onto the given grid system. @@ -207,15 +212,17 @@ class CzjzekDistribution(AbstractDistribution): """ model_name = "czjzek" - def __init__(self, sigma: float, polar=False): + def __init__(self, sigma: float, polar=False, cache=False): + super().__init__(cache_tensors=cache) self.sigma = sigma self.polar = polar - def rvs(self, size: int): + def rvs(self, size: int, tensors: np.ndarray): """Draw random variates of length `size` from the distribution. Args: size: The number of random points to draw. + tensors: Pre-computed array of tensors. Returns: A list of two NumPy array, where the first and the second array are the @@ -225,7 +232,13 @@ def rvs(self, size: int): Example: >>> Cq_dist, eta_dist = cz_model.rvs(size=1000000) """ - tensors = _czjzek_random_distribution_tensors(self.sigma, size) + if self._cache_tensors: + if self._tensors is None: + self._tensors = _czjzek_random_distribution_tensors(self.sigma, size) + tensors = self._tensors + else: + tensors = _czjzek_random_distribution_tensors(self.sigma, size) + if not self.polar: return get_Haeberlen_components(tensors) return zeta_eta_to_x_y(*get_Haeberlen_components(tensors)) @@ -273,16 +286,20 @@ class ExtCzjzekDistribution(AbstractDistribution): """ model_name = "extended czjzek" - def __init__(self, symmetric_tensor: SymmetricTensor, eps: float, polar=False): + def __init__( + self, symmetric_tensor: SymmetricTensor, eps: float, polar=False, cache=False + ): + super().__init__(cache_tensors=cache) self.symmetric_tensor = symmetric_tensor self.eps = eps self.polar = polar - def rvs(self, size: int): + def rvs(self, size: int, tensors=None): """Draw random variates of length `size` from the distribution. Args: size: The number of random points to draw. + tensors: Pre-computed array of tensors. Returns: A list of two NumPy array, where the first and the second array are the @@ -294,7 +311,12 @@ def rvs(self, size: int): """ # czjzek_random_distribution model - tensors = _czjzek_random_distribution_tensors(1, size) + if self._cache_tensors: + if self._tensors is None: + self._tensors = _czjzek_random_distribution_tensors(1, size) + tensors = self._tensors + else: + tensors = _czjzek_random_distribution_tensors(1, size) symmetric_tensor = self.symmetric_tensor diff --git a/src/mrsimulator/models/utils.py b/src/mrsimulator/models/utils.py index 25290322a..cc96c6b71 100644 --- a/src/mrsimulator/models/utils.py +++ b/src/mrsimulator/models/utils.py @@ -87,19 +87,21 @@ def x_y_to_zeta_eta(x, y): return zeta, eta -def _simulate_spectra_over_zeta_and_eta(ZZ, ee, mth, tensor_type): +def _simulate_spectra_over_zeta_and_eta(ZZ, ee, method, tensor_type): """Helper function to generate the kernel""" - isotope = Isotope.parse(mth.channels[0]).symbol # Grab isotope from Method object + isotope = Isotope.parse( + method.channels[0] + ).symbol # Grab isotope from Method object spin_systems = [ - SpinSystem(sites=[Site(isotope=isotope, quadrupolar=dict(Cq=Z * 1e6, eta=e))]) + SpinSystem(sites=[Site(isotope=isotope, quadrupolar=dict(Cq=Z, eta=e))]) if tensor_type == "quadrupolar" else SpinSystem( sites=[Site(isotope=isotope, shielding_symmetric=dict(zeta=Z, eta=e))] ) - for Z, e in zip(ZZ.flatten(), ee.flatten()) + for Z, e in zip(ZZ.ravel(), ee.ravel()) ] - sim = Simulator(spin_systems=spin_systems, methods=[mth]) + sim = Simulator(spin_systems=spin_systems, methods=[method]) sim.config.decompose_spectrum = "spin_system" sim.run(pack_as_csdm=False) @@ -108,7 +110,7 @@ def _simulate_spectra_over_zeta_and_eta(ZZ, ee, mth, tensor_type): def generate_lineshape_kernel( - pos: tuple, mth: Method, polar: bool, tensor_type: str = "shielding" + pos: tuple, method: Method, polar: bool, tensor_type: str = "shielding" ) -> np.ndarray: """Pre-compute a lineshape kernel used during least-squares fitting of an experimental spectrum. The isotope for the spin system is grabbed from the isotope @@ -120,7 +122,7 @@ def generate_lineshape_kernel( Arguments: (tuple) pos: A set of numpy arrays defining the grid space - (mrsimulator.Method) mth: The :py:class:`~mrsimulator.method.Method` used to + (mrsimulator.Method) method: The :py:class:`~mrsimulator.method.Method` used to simulate the spectra. (bool) polar: Weather the grid is defined in polar coordinates (str) tensor_type: A string enumeration literal describing which type of tensor @@ -140,4 +142,4 @@ def generate_lineshape_kernel( if polar: ZZ, ee = x_y_to_zeta_eta(ZZ.ravel(), ee.ravel()) - return _simulate_spectra_over_zeta_and_eta(ZZ, ee, mth, tensor_type) + return _simulate_spectra_over_zeta_and_eta(ZZ, ee, method, tensor_type) diff --git a/src/mrsimulator/utils/spectral_fitting.py b/src/mrsimulator/utils/spectral_fitting.py index 2fea1315f..89ccf1179 100644 --- a/src/mrsimulator/utils/spectral_fitting.py +++ b/src/mrsimulator/utils/spectral_fitting.py @@ -890,7 +890,7 @@ def _make_spectrum_from_extended_Czjzek_distribution_parameters( # Dot amplitude with kernel, then package as CSDM object spec_tmp = guess_spectrum.copy() - spec_tmp.y[0].components[0] = np.dot(amp.flatten(), kernel) + spec_tmp.y[0].components[0] = np.dot(amp.ravel(), kernel) # Apply isotropic chemical shift to distribution using FFT shift theorem spec_tmp = _apply_iso_shift( From 62761894899c5b04e8e2ed2020c5bba7cb3e118b Mon Sep 17 00:00:00 2001 From: deepanshs <21365911+deepanshs@users.noreply.github.com> Date: Tue, 19 Mar 2024 07:24:43 -0400 Subject: [PATCH 09/21] code simplify --- .../1D_fitting/plot_7_27Al_Extended_Czjzek.py | 47 +++++++-------- src/mrsimulator/models/czjzek.py | 58 +----------------- src/mrsimulator/utils/spectral_fitting.py | 60 ++++++------------- 3 files changed, 41 insertions(+), 124 deletions(-) diff --git a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py index 991b91064..ac58d45fe 100644 --- a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py +++ b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py @@ -27,7 +27,6 @@ from mrsimulator.utils import get_spectral_dimensions from mrsimulator.models.czjzek import ExtCzjzekDistribution -from mrsimulator.models.czjzek import _czjzek_random_distribution_tensors from mrsimulator.models.utils import generate_lineshape_kernel # sphinx_gallery_thumbnail_number = 3 @@ -41,10 +40,14 @@ host = "http://ssnmr.org/sites/default/files/mrsimulator/" filename = "20K_20Al_10P_50Si_HahnEcho_27Al.csdf" exp_data = cp.load(host + filename).real + exp_data.x[0].to("ppm", "nmr_frequency_ratio") exp_data /= exp_data.max() -exp_data.plot() +plt.figure(figsize=(4, 3)) +ax = plt.subplot(projection="csdm") +ax.plot(exp_data) +plt.tight_layout() plt.show() # %% @@ -74,31 +77,23 @@ # generated by passing `tensor_type="shielding"`. # Create a Method object to simulate the spectrum spectral_dimensions = get_spectral_dimensions(exp_data) -mth = BlochDecayCTSpectrum( +method = BlochDecayCTSpectrum( channels=["27Al"], rotor_frequency=14.2e3, spectral_dimensions=spectral_dimensions ) # Define a polar grid for the lineshape kernel -x = np.linspace(0, 12, num=36) -y = np.linspace(0, 12, num=36) +x = np.linspace(0, 2e7, num=36) +y = np.linspace(0, 2e7, num=36) pos = (x, y) # Call the kernel generation function with the correct arguments kernel = generate_lineshape_kernel( - pos=pos, mth=mth, polar=True, tensor_type="quadrupolar" + pos=pos, method=method, polar=True, tensor_type="quadrupolar" ) print("Desired Kernel shape: ", (x.size * y.size, spectral_dimensions[0]["count"])) print("Actual Kernel shape: ", kernel.shape) -# %% -# Create a set of random noise tensors -# ------------------------------------ -# Below we sample a set of 200,000 random noise tensors to use during the fit; this set -# of tensors remain the same throughout the least-squares minimization making the fit -# more stable. -tensors = _czjzek_random_distribution_tensors(sigma=1, n=200_000) - # %% # Create a Parameters object # -------------------------- @@ -114,8 +109,10 @@ # reflect any changes made in the minimization. # Create initial guess ExtCzjzekDistribution -tensor = {"Cq": 3, "eta": 0.3} -ext_cz_model = ExtCzjzekDistribution(symmetric_tensor=tensor, eps=1.6) +tensor_guess = {"Cq": 1e6, "eta": 0.3} +ext_cz_model = ExtCzjzekDistribution( + symmetric_tensor=tensor_guess, eps=6, polar=True, cache=True +) dist_iso_shift = 60 # in ppm processor = sp.SignalProcessor( @@ -140,11 +137,9 @@ addtl_sf_kwargs = dict( exp_spectrum=exp_data, pos=pos, - polar=True, kernel=kernel, - tensors=tensors, - n_dist=1, - larmor_freq_Hz=mth.channels[0].larmor_freq(B0=mth.magnetic_flux_density), + models=[ext_cz_model], + larmor_freq_Hz=method.channels[0].larmor_freq(B0=method.magnetic_flux_density), processor=processor, tensor_type="quadrupolar", ) @@ -160,8 +155,8 @@ ax.plot(residuals, "b", alpha=0.3, label="Residuals") plt.legend() plt.grid() -plt.tight_layout() plt.title("Initial Guess") +plt.tight_layout() plt.show() # Print the Parameters object @@ -204,8 +199,8 @@ ax.plot(residuals, "b", alpha=0.3, label="Residuals") plt.legend() plt.grid() -plt.tight_layout() plt.title("Best Fit") +plt.tight_layout() plt.show() # %% @@ -221,14 +216,14 @@ bf_tensor["Cq"] = result.params["ext_czjzek_0_Cq0"].value bestfit_model = ExtCzjzekDistribution( - symmetric_tensor=tensor, eps=epsilon, polar=addtl_sf_kwargs["polar"] + symmetric_tensor=bf_tensor, eps=epsilon, polar=True ) xx, yy, amp = bestfit_model.pdf(pos=pos, size=10000000) plt.figure(figsize=(5, 5)) -plt.contourf(xx, yy, amp) -plt.xlabel("x / (ppm)") -plt.ylabel("y / (ppm)") +plt.contourf(xx / 1e6, yy / 1e6, amp) +plt.xlabel("x / MHz ") +plt.ylabel("y / MHz") plt.tight_layout() plt.show() diff --git a/src/mrsimulator/models/czjzek.py b/src/mrsimulator/models/czjzek.py index 8a538572b..7ab607bd4 100644 --- a/src/mrsimulator/models/czjzek.py +++ b/src/mrsimulator/models/czjzek.py @@ -19,58 +19,6 @@ ANALYTICAL_AVAILABLE = {"czjzek": analytical_dist.czjzek} -# this function does the same as self.pdf() but with cached tensors! -def _extended_czjzek_pdf_from_tensors(pos, tensors, zeta0, eta0, epsilon, polar): - """Takes in a list of random noise tensors along with the parameters for - a given Extended Czjzek distribution, applies the requisite math to - the tensors, then diagonalizes the tensors to get the (zeta, eta) - distribution. This allows the same sampling points to be drawn between different - minimization steps - - Arguments: - (tuple) pos: Two np.ndarrays defining the grid on which to calculate the pdf - (np.ndarray) tensors: A np.ndarray with shape (n, 3, 3) representing - the random tensors in the Extended Czjzek model - (float) zeta0: The zeta value of the central tensor - (float) eta0: The eta value of the central tensor - (float) epsilon: The noise parameter for the Extended Czjzek distribution - (bool) polar: Weather the distribution should be sampled in polar coordinates. - - Returns: - (np.ndarray, np.ndarray) of the zeta and eta distributions, respectively - """ - T0 = get_principal_components(zeta0, eta0) - - norm_T0 = np.linalg.norm(T0) - rho = epsilon * norm_T0 / np.sqrt(30) - - if rho != 0: - ext_tensors = (rho * tensors) + np.diag(T0) - else: - ext_tensors = tensors.copy() - - zeta_dist, eta_dist = get_Haeberlen_components(ext_tensors) - if polar: - zeta_dist, eta_dist = zeta_eta_to_x_y(zeta_dist, eta_dist) - - delta_0 = (pos[0][1] - pos[0][0]) / 2 - delta_1 = (pos[1][1] - pos[1][0]) / 2 - pts_0 = [pos[0][0] - delta_0, pos[0][-1] + delta_0] - pts_1 = [pos[1][0] - delta_1, pos[1][-1] + delta_1] - - size_0 = pos[0].size - size_1 = pos[1].size - - hist, _, _ = np.histogram2d( - zeta_dist, eta_dist, bins=[size_0, size_1], range=[pts_0, pts_1] - ) - - hist /= hist.sum() - xx, yy = np.meshgrid(pos[0], pos[1]) - - return xx, yy, hist.T - - def _czjzek_random_distribution_tensors(sigma, n): r"""Czjzek random distribution model. @@ -217,12 +165,11 @@ def __init__(self, sigma: float, polar=False, cache=False): self.sigma = sigma self.polar = polar - def rvs(self, size: int, tensors: np.ndarray): + def rvs(self, size: int): """Draw random variates of length `size` from the distribution. Args: size: The number of random points to draw. - tensors: Pre-computed array of tensors. Returns: A list of two NumPy array, where the first and the second array are the @@ -294,12 +241,11 @@ def __init__( self.eps = eps self.polar = polar - def rvs(self, size: int, tensors=None): + def rvs(self, size: int): """Draw random variates of length `size` from the distribution. Args: size: The number of random points to draw. - tensors: Pre-computed array of tensors. Returns: A list of two NumPy array, where the first and the second array are the diff --git a/src/mrsimulator/utils/spectral_fitting.py b/src/mrsimulator/utils/spectral_fitting.py index 89ccf1179..e2f61766e 100644 --- a/src/mrsimulator/utils/spectral_fitting.py +++ b/src/mrsimulator/utils/spectral_fitting.py @@ -3,7 +3,6 @@ import numpy as np from lmfit import Parameters from mrsimulator import Simulator -from mrsimulator.models.czjzek import _extended_czjzek_pdf_from_tensors from mrsimulator.models.czjzek import CzjzekDistribution from mrsimulator.spin_system.tensors import SymmetricTensor @@ -829,10 +828,8 @@ def _make_spectrum_from_extended_Czjzek_distribution_parameters( params: Parameters, exp_spectrum: cp.CSDM, pos: tuple, - polar: bool, kernel: np.ndarray, - tensors: np.ndarray, - n_dist: int, + models: list, larmor_freq_Hz: float, processor: sp.SignalProcessor = None, tensor_type: str = "shielding", @@ -847,11 +844,9 @@ def _make_spectrum_from_extended_Czjzek_distribution_parameters( (cp.CSDM) exp_spectrum: The experimental spectrum to fit to. (tuple) pos: A tuple of two np.ndarray objects defining the grid on which to sample the distribution. - (bool) polar: True if the sample grid is in polar coordinates. False if the grid - is defined using the Haberlen components. - (np.ndarray) The pre-computed lineshape kernel. The kernel needs to be defined - on the same grid defined by pos. - (int) n_dist: The number of Czjzek distributions to fit for. + (np.ndarray) kernel: The pre-computed lineshape kernel. The kernel needs to be + defined on the same grid defined by pos. + (list) models: A list of distribution models. (float) larmor_freq_Hz: This value is used in conjunction with the FFT shift theorem to apply an isotropic chemical shift to each distribution. (sp.SignalProcessor) processor: A @@ -863,6 +858,7 @@ def _make_spectrum_from_extended_Czjzek_distribution_parameters( Guess spectrum as a cp.CSDM object """ + n_dist = len(models) if tensor_type not in {"shielding", "quadrupolar"}: raise ValueError(f"Unknown `tensor_type` value of {tensor_type}.") @@ -879,14 +875,10 @@ def _make_spectrum_from_extended_Czjzek_distribution_parameters( epsilon = params[f"ext_czjzek_{i}_epsilon"].value # Draw the pdf from the set of static random tensors - _, _, amp = _extended_czjzek_pdf_from_tensors( - pos=pos, - tensors=tensors, - zeta0=zeta0, - eta0=eta0, - epsilon=epsilon, - polar=polar, - ) + tensor = {"Cq": zeta0, "eta": eta0} + models[i].symmetric_tensor = tensor + models[i].eps = epsilon + _, _, amp = models[i].pdf(pos, size=200_000) # Dot amplitude with kernel, then package as CSDM object spec_tmp = guess_spectrum.copy() @@ -912,10 +904,8 @@ def LMFIT_min_function_extended_Czjzek( params: Parameters, exp_spectrum: cp.CSDM, pos: tuple, - polar: bool, kernel: np.ndarray, - tensors: np.ndarray, - n_dist: int, + models: list, larmor_freq_Hz: float, processor: sp.SignalProcessor = None, tensor_type: str = "shielding", @@ -929,13 +919,9 @@ def LMFIT_min_function_extended_Czjzek( (cp.CSDM) exp_spectrum: The experimental spectrum to fit to. (tuple) pos: A tuple of two np.ndarray objects defining the grid on which to sample the distribution. - (bool) polar: True if the sample grid is in polar coordinates. False if the grid - is defined using the Haberlen components. - (np.ndarray) The pre-computed lineshape kernel. The kernel needs to be defined - on the same grid defined by pos. - (np.ndarray) tensors: A set of random Czjzek tensors in 5-dimensional space to - calculate the Extended Czjzek tensors from. - (int) n_dist: The number of Czjzek distributions to fit for. + (np.ndarray) kernel: The pre-computed lineshape kernel. The kernel needs to be + defined on the same grid defined by pos. + (list) models: A list of distribution models. (float) larmor_freq_Hz: This value is used in conjunction with the FFT shift theorem to apply an isotropic chemical shift to each distribution. (sp.SignalProcessor) processor: A @@ -953,10 +939,8 @@ def LMFIT_min_function_extended_Czjzek( params, exp_spectrum, pos, - polar, kernel, - tensors, - n_dist, + models, larmor_freq_Hz, processor, tensor_type, @@ -969,10 +953,8 @@ def bestfit_extended_Czjzek( params: Parameters, exp_spectrum: cp.CSDM, pos: tuple, - polar: bool, kernel: np.ndarray, - tensors: np.ndarray, - n_dist: int, + models: list, larmor_freq_Hz: float, processor: sp.SignalProcessor = None, tensor_type: str = "shielding", @@ -982,10 +964,8 @@ def bestfit_extended_Czjzek( params, exp_spectrum, pos, - polar, kernel, - tensors, - n_dist, + models, larmor_freq_Hz, processor, tensor_type, @@ -996,10 +976,8 @@ def residuals_extended_Czjzek( params: Parameters, exp_spectrum: cp.CSDM, pos: tuple, - polar: bool, kernel: np.ndarray, - tensors: np.ndarray, - n_dist: int, + models: list, larmor_freq_Hz: float, processor: sp.SignalProcessor = None, tensor_type: str = "shielding", @@ -1009,10 +987,8 @@ def residuals_extended_Czjzek( params, exp_spectrum, pos, - polar, kernel, - tensors, - n_dist, + models, larmor_freq_Hz, processor, tensor_type, From 3ec2457391fe18d59ca983b3725c6e7fe72ef4de Mon Sep 17 00:00:00 2001 From: Deepansh Srivastava Date: Tue, 19 Mar 2024 12:15:26 -0400 Subject: [PATCH 10/21] simplify code --- .../1D_fitting/plot_7_27Al_Extended_Czjzek.py | 45 +-- src/mrsimulator/models/czjzek.py | 104 +++++- src/mrsimulator/utils/spectral_fitting.py | 353 ++---------------- 3 files changed, 150 insertions(+), 352 deletions(-) diff --git a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py index ac58d45fe..a80b31ed0 100644 --- a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py +++ b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py @@ -11,7 +11,7 @@ # solid. There are 4 major steps involved in the process: # # 1. Importing the experimental dataset -# 2. Generating a pre-computed linshape kernel for the experiment +# 2. Generating a pre-computed line shape kernel for the experiment # 3. Creating parameters for the Extended Czjzek model from an initial guess # 4. Minimizing and visualizing. # @@ -51,7 +51,7 @@ plt.show() # %% -# Generating a lineshapke kernel +# Generating a line shape kernel # ------------------------------ # # A spectrum arising from an Extended Czjzek distribution can be modeled by drawing a @@ -67,7 +67,7 @@ # probability distribution on that grid during each minimization step. # # Below, we create a `Method` object for simulating the spectra of the given experiment, -# define a polar grid for generating the lineshape kernel, and finally call the kernel +# define a polar grid for generating the line shape kernel, and finally call the kernel # generation method -- `generate_lineshape_kernel`. # The grid we defined is in polar coordinates, so we pass `polar=True` to the kernel # generation method; if the grid were defined for the Haeberlen components of the @@ -111,9 +111,12 @@ # Create initial guess ExtCzjzekDistribution tensor_guess = {"Cq": 1e6, "eta": 0.3} ext_cz_model = ExtCzjzekDistribution( - symmetric_tensor=tensor_guess, eps=6, polar=True, cache=True + mean_isotropic_chemical_shift=60, + symmetric_tensor=tensor_guess, + eps=6, + polar=True, + cache=True, ) -dist_iso_shift = 60 # in ppm processor = sp.SignalProcessor( operations=[ @@ -125,11 +128,9 @@ ) # Make the Parameters object -params = sf.make_LMFIT_params_extended_Czjzek( - ext_cz_models=[ext_cz_model], - iso_shifts=[dist_iso_shift], +params = sf.make_LMFIT_distribution_params( + distribution_models=[ext_cz_model], processor=processor, - tensor_type="quadrupolar", ) @@ -141,12 +142,11 @@ models=[ext_cz_model], larmor_freq_Hz=method.channels[0].larmor_freq(B0=method.magnetic_flux_density), processor=processor, - tensor_type="quadrupolar", ) # Make a guess and residuals spectrum from the initial guess -guess = sf.bestfit_extended_Czjzek(params=params, **addtl_sf_kwargs) -residuals = sf.residuals_extended_Czjzek(params=params, **addtl_sf_kwargs) +guess = sf.bestfit_dist(params=params, **addtl_sf_kwargs) +residuals = sf.residuals_dist(params=params, **addtl_sf_kwargs) plt.figure(figsize=(5, 4)) ax = plt.subplot(projection="csdm") @@ -178,7 +178,7 @@ minner = Minimizer( - sf.LMFIT_min_function_extended_Czjzek, + sf.LMFIT_min_function_dist, params, fcn_kws=addtl_sf_kwargs, **scipy_minimization_kwargs, @@ -189,8 +189,8 @@ # %% # Plot the best-fit spectrum # -------------------------- -bestfit = sf.bestfit_extended_Czjzek(params=result.params, **addtl_sf_kwargs) -residuals = sf.residuals_extended_Czjzek(params=result.params, **addtl_sf_kwargs) +bestfit = sf.bestfit_dist(params=result.params, **addtl_sf_kwargs) +residuals = sf.residuals_dist(params=result.params, **addtl_sf_kwargs) plt.figure(figsize=(5, 4)) ax = plt.subplot(projection="csdm") @@ -207,19 +207,10 @@ # Plot the best-fit distribution # ------------------------------ # -# Note that this cell is hardcoded for a single distribution. -epsilon = result.params["ext_czjzek_0_epsilon"].value -bf_tensor = {"eta": result.params["ext_czjzek_0_eta0"].value} -if addtl_sf_kwargs["tensor_type"] == "shielding": - bf_tensor["zeta"] = result.params["ext_czjzek_0_zeta0"].value -else: - bf_tensor["Cq"] = result.params["ext_czjzek_0_Cq0"].value - -bestfit_model = ExtCzjzekDistribution( - symmetric_tensor=bf_tensor, eps=epsilon, polar=True -) +for i, model in enumerate([ext_cz_model]): + model.update_lmfit_params(result.params, i) -xx, yy, amp = bestfit_model.pdf(pos=pos, size=10000000) +xx, yy, amp = ext_cz_model.pdf(pos=pos, size=10000000) plt.figure(figsize=(5, 5)) plt.contourf(xx / 1e6, yy / 1e6, amp) diff --git a/src/mrsimulator/models/czjzek.py b/src/mrsimulator/models/czjzek.py index 7ab607bd4..9cd2f8b94 100644 --- a/src/mrsimulator/models/czjzek.py +++ b/src/mrsimulator/models/czjzek.py @@ -78,9 +78,19 @@ def _czjzek_random_distribution_tensors(sigma, n): class AbstractDistribution: - def __init__(self, cache_tensors=False): + def __init__( + self, + mean_isotropic_chemical_shift=0.0, + abundance=1.0, + polar=False, + cache_tensors=False, + ): + """Basic class attributes for distributions""" self._cache_tensors = cache_tensors self._tensors = None + self.mean_isotropic_chemical_shift = mean_isotropic_chemical_shift + self.abundance = abundance + self.polar = polar def pdf(self, pos, size: int = 400000, analytical: bool = True): """Generates a probability distribution function by binning the random @@ -160,10 +170,21 @@ class CzjzekDistribution(AbstractDistribution): """ model_name = "czjzek" - def __init__(self, sigma: float, polar=False, cache=False): - super().__init__(cache_tensors=cache) + def __init__( + self, + sigma: float, + mean_isotropic_chemical_shift: float = 0.0, + abundance: float = 1.0, + polar=False, + cache=False, + ): + super().__init__( + cache_tensors=cache, + polar=polar, + mean_isotropic_chemical_shift=mean_isotropic_chemical_shift, + abundance=abundance, + ) self.sigma = sigma - self.polar = polar def rvs(self, size: int): """Draw random variates of length `size` from the distribution. @@ -190,6 +211,25 @@ def rvs(self, size: int): return get_Haeberlen_components(tensors) return zeta_eta_to_x_y(*get_Haeberlen_components(tensors)) + def param_prefix(self): + return "czjzek" + + def get_lmfit_params(self, params, i): + """Create lmfit params for index i""" + params.add(f"czjzek_{i}_sigma", value=self.sigma, min=0) + params.add(f"czjzek_{i}_iso_shift", value=self.mean_isotropic_chemical_shift) + params.add(f"czjzek_{i}_weight", value=self.abundance, min=0, max=1) + return params + + def update_lmfit_params(self, params, i): + """Create lmfit params for index i""" + prefix = self.param_prefix() + + self.sigma = params[f"{prefix}_{i}_sigma"].value + self.eps = params[f"{prefix}_{i}_epsilon"].value + self.mean_isotropic_chemical_shift = params[f"{prefix}_{i}_iso_shift"].value + self.abundance = params[f"{prefix}_{i}_weight"].value + class ExtCzjzekDistribution(AbstractDistribution): r"""An extended Czjzek distribution distribution model. @@ -234,12 +274,25 @@ class ExtCzjzekDistribution(AbstractDistribution): model_name = "extended czjzek" def __init__( - self, symmetric_tensor: SymmetricTensor, eps: float, polar=False, cache=False + self, + symmetric_tensor: SymmetricTensor, + eps: float, + mean_isotropic_chemical_shift: float = 0.0, + abundance: float = 1.0, + polar=False, + cache=False, ): - super().__init__(cache_tensors=cache) - self.symmetric_tensor = symmetric_tensor + super().__init__( + cache_tensors=cache, + polar=polar, + mean_isotropic_chemical_shift=mean_isotropic_chemical_shift, + abundance=abundance, + ) + if isinstance(symmetric_tensor, dict): + self.symmetric_tensor = SymmetricTensor(**symmetric_tensor) + else: + self.symmetric_tensor = symmetric_tensor self.eps = eps - self.polar = polar def rvs(self, size: int): """Draw random variates of length `size` from the distribution. @@ -266,9 +319,6 @@ def rvs(self, size: int): symmetric_tensor = self.symmetric_tensor - if isinstance(symmetric_tensor, dict): - symmetric_tensor = SymmetricTensor(**symmetric_tensor) - zeta = symmetric_tensor.zeta or symmetric_tensor.Cq eta = symmetric_tensor.eta @@ -289,3 +339,35 @@ def rvs(self, size: int): if not self.polar: return get_Haeberlen_components(total_tensors) return zeta_eta_to_x_y(*get_Haeberlen_components(total_tensors)) + + def param_prefix(self): + return "ext_czjzek" + + def get_lmfit_params(self, params, i): + """Create lmfit params for index i""" + prefix = self.param_prefix() + if self.symmetric_tensor.zeta is not None: + zeta = self.symmetric_tensor.zeta + else: + zeta = self.symmetric_tensor.Cq + params.add(f"{prefix}_{i}_zeta0", value=zeta) + params.add(f"{prefix}_{i}_eta0", value=self.symmetric_tensor.eta, min=0, max=1) + params.add(f"{prefix}_{i}_epsilon", value=self.eps, min=0) + params.add(f"{prefix}_{i}_iso_shift", value=self.mean_isotropic_chemical_shift) + params.add(f"{prefix}_{i}_weight", value=self.abundance, min=0, max=1) + return params + + def update_lmfit_params(self, params, i): + """Create lmfit params for index i""" + prefix = self.param_prefix() + + zeta = params[f"{prefix}_{i}_zeta0"].value + if self.symmetric_tensor.zeta is not None: + self.symmetric_tensor.zeta = zeta + else: + self.symmetric_tensor.Cq = zeta + + self.symmetric_tensor.eta = params[f"{prefix}_{i}_eta0"].value + self.eps = params[f"{prefix}_{i}_epsilon"].value + self.mean_isotropic_chemical_shift = params[f"{prefix}_{i}_iso_shift"].value + self.abundance = params[f"{prefix}_{i}_weight"].value diff --git a/src/mrsimulator/utils/spectral_fitting.py b/src/mrsimulator/utils/spectral_fitting.py index e2f61766e..6d66a127a 100644 --- a/src/mrsimulator/utils/spectral_fitting.py +++ b/src/mrsimulator/utils/spectral_fitting.py @@ -3,8 +3,6 @@ import numpy as np from lmfit import Parameters from mrsimulator import Simulator -from mrsimulator.models.czjzek import CzjzekDistribution -from mrsimulator.spin_system.tensors import SymmetricTensor __author__ = ["Maxwell C Venetos", "Deepansh Srivastava"] __email__ = ["maxvenetos@gmail.com", "srivastava.89@osu.edu"] @@ -539,50 +537,34 @@ def _apply_iso_shift(csdm_obj, iso_shift_ppm, larmor_freq_Hz): return csdm_obj -def make_LMFIT_params_Czjzek( - cz_models: list, - weights: list = None, - iso_shifts: list = None, - processor: sp.SignalProcessor = None, +def make_LMFIT_distribution_params( + distribution_models: list, processor: sp.SignalProcessor = None ) -> Parameters: - """Generates the LMfit Parameters object with the proper keywords for a Czjzek - distribution. Each distribution has the following parameters: - - czjzek_i_sigma - - czjzek_i_iso_shift - - czjzek_i_weight - - where *i* refers to the :math:i^\text{th} Czjzek distribution - - Arguments: - (list) cz_models: A list of :class:`~mrsimulator.models.CzjzekDistribution` - objects used to set initial values. - (list) weights: An optional list of float values defining the relative weight of - each Czjzek distribution. The default is (1 / n_dist), that is, all - distributions have the same weight - (list) iso_shifts: An optional list of float values defining the central - isotropic chemical shift for each distribution. - (sp.SignalProcessor) processor: An optional signal processor to apply - post-simulation signal processing to each guess spectrum. - - Returns: - A LMfit Parameters object + """Generate LMfit Parameters object for spin system distribution model. + The distribution has the following parameters: """ - n_dist = len(cz_models) - iso_shifts = [0] * n_dist if iso_shifts is None else iso_shifts - weights = np.ones(n_dist) if weights is None else np.asarray(weights) - weights /= weights.sum() # Normalize weights array to total of 1 + n_dist = len(distribution_models) - params = Parameters() - for i in range(n_dist): - params.add(f"czjzek_{i}_sigma", value=cz_models[i].sigma, min=0) - params.add(f"czjzek_{i}_iso_shift", value=iso_shifts[i]) - params.add(f"czjzek_{i}_weight", value=weights[i], min=0, max=1) + norm = 0.0 + for model in distribution_models: + norm += model.abundance - # Set last weight parameter as a non-free variable - expr = "-".join([f"czjzek_{i}_weight" for i in range(n_dist - 1)]) - expr = "1" if expr == "" else f"1-{expr}" - params[f"czjzek_{n_dist-1}_weight"].vary = False - params[f"czjzek_{n_dist-1}_weight"].expr = expr + expr_terms = [] + params = Parameters() + for i, model in enumerate(distribution_models): + # normalize the abundance + model.abundance /= norm + params = model.get_lmfit_params(params, i) + + # Set last weight parameter as a non-free variable + model_prefix = model.param_prefix() + if i < n_dist - 1: + expr_terms.append(f"{model_prefix}_{i}_weight") + else: + expr = "-".join(expr_terms) + expr = "1" if expr == "" else f"1-{expr}" + params[f"{model_prefix}_{i}_weight"].vary = False + params[f"{model_prefix}_{i}_weight"].expr = expr # Add SignalProcessor parameters, if requested if processor is not None: @@ -591,13 +573,12 @@ def make_LMFIT_params_Czjzek( return params -def _make_spectrum_from_Czjzek_distribution_parameters( +def _generate_distribution_spectrum( params: Parameters, exp_spectrum: cp.CSDM, pos: tuple, - polar: bool, kernel: np.ndarray, - n_dist: int, + distribution_models: list, larmor_freq_Hz: float, processor: sp.SignalProcessor = None, ) -> cp.CSDM: @@ -611,11 +592,8 @@ def _make_spectrum_from_Czjzek_distribution_parameters( (cp.CSDM) exp_spectrum: The experimental spectrum to fit to. (tuple) pos: A tuple of two np.ndarray objects defining the grid on which to sample the distribution. - (bool) polar: True if the sample grid is in polar coordinates. False if the grid - is defined using the Haberlen components. (np.ndarray) The pre-computed lineshape kernel. The kernel needs to be defined on the same grid defined by pos. - (int) n_dist: The number of Czjzek distributions to fit for. (float) larmor_freq_Hz: This value is used in conjunction with the FFT shift theorem to apply an isotropic chemical shift to each distribution. (sp.SignalProcessor) processor: A @@ -630,10 +608,9 @@ def _make_spectrum_from_Czjzek_distribution_parameters( guess_spectrum = exp_spectrum.copy() guess_spectrum.y[0].components[:] = 0 # Initialize guess spectrum with zeros - for i in range(n_dist): - _, _, amp = CzjzekDistribution( - sigma=params[f"czjzek_{i}_sigma"], polar=polar - ).pdf(pos) + for i, model in enumerate(distribution_models): + model.update_lmfit_params(params, i) + _, _, amp = model.pdf(pos) # Dot amplitude with kernel, then package as CSDM object spec_tmp = guess_spectrum.copy() @@ -642,11 +619,11 @@ def _make_spectrum_from_Czjzek_distribution_parameters( # Apply isotropic chemical shift to distribution using FFT shift theorem spec_tmp = _apply_iso_shift( csdm_obj=spec_tmp, - iso_shift_ppm=params[f"czjzek_{i}_iso_shift"].value, + iso_shift_ppm=model.mean_isotropic_chemical_shift, larmor_freq_Hz=larmor_freq_Hz, ) - guess_spectrum += spec_tmp.real * params[f"czjzek_{i}_weight"].value + guess_spectrum += spec_tmp.real * model.abundance if processor is not None: _update_processors_from_LMFIT_params(params, [processor]) @@ -655,13 +632,12 @@ def _make_spectrum_from_Czjzek_distribution_parameters( return guess_spectrum -def LMFIT_min_function_Czjzek( +def LMFIT_min_function_dist( params: Parameters, exp_spectrum: cp.CSDM, pos: tuple, - polar: bool, kernel: np.ndarray, - n_dist: int, + models: list, larmor_freq_Hz: float, processor: sp.SignalProcessor = None, ) -> np.ndarray: @@ -689,253 +665,7 @@ def LMFIT_min_function_Czjzek( Returns: A residual array as a numpy array. """ - guess_spectrum = _make_spectrum_from_Czjzek_distribution_parameters( - params, - exp_spectrum, - pos, - polar, - kernel, - n_dist, - larmor_freq_Hz, - processor, - ) - - return (exp_spectrum - guess_spectrum).y[0].components[0] - - -def bestfit_Czjzek( - params: Parameters, - exp_spectrum: cp.CSDM, - pos: tuple, - polar: bool, - kernel: np.ndarray, - n_dist: int, - larmor_freq_Hz: float, - processor: sp.SignalProcessor = None, -) -> cp.CSDM: - """Returns the best-fit spectrum as a CSDM object""" - return _make_spectrum_from_Czjzek_distribution_parameters( - params, - exp_spectrum, - pos, - polar, - kernel, - n_dist, - larmor_freq_Hz, - processor, - ) - - -def residuals_Czjzek( - params: Parameters, - exp_spectrum: cp.CSDM, - pos: tuple, - polar: bool, - kernel: np.ndarray, - n_dist: int, - larmor_freq_Hz: float, - processor: sp.SignalProcessor = None, -) -> cp.CSDM: - """Returns the residuals spectrum as a CSDM object""" - bestfit = _make_spectrum_from_Czjzek_distribution_parameters( - params, - exp_spectrum, - pos, - polar, - kernel, - n_dist, - larmor_freq_Hz, - processor, - ) - - return exp_spectrum - bestfit - - -def make_LMFIT_params_extended_Czjzek( - ext_cz_models: list, - weights: list = None, - iso_shifts: list = None, - processor: sp.SignalProcessor = None, - tensor_type: str = "shielding", -) -> Parameters: - """Generates the LMfit Parameters object with the proper keywords for an Extended - Czjzek distribution. Each distribution has the following parameters: - - ext_czjzek_i_zeta0 *or* ext_czjzek_i_Cq0 - - ext_czjzek_i_eta0 - - ext_czjzek_i_epsilon - - ext_czjzek_i_iso_shift - - ext_czjzek_i_weight - - where *i* refers to the :math:i^\text{th} Extended Czjzek distribution. Note that - Cq values are assumed to be in units of MHz. - - Arguments: - (list) ext_cz_models: A list of - :class:`~mrsimulator.models.ExtCzjzekDistribution` objects used to set - initial values. - (list) weights: An optional list of float values defining the relative weight of - each Czjzek distribution. The default is (1 / n_dist), that is, all - distributions have the same weight - (list) iso_shifts: An optional list of float values defining the central - isotropic chemical shift for each distribution. - (sp.SignalProcessor) processor: An optional signal processor to apply - post-simulation signal processing to each guess spectrum. - (str) tensor_type: A string literal describing if the Extended Czjzek models - define a distribution of symmetric shielding or quadrupolar tensors. The - allowed values are `shielding` and `quadrupolar`. - - Returns: - A LMfit Parameters object - """ - if tensor_type not in {"shielding", "quadrupolar"}: - raise ValueError(f"Unrecognized value of {tensor_type} for `tensor_type.") - - n_dist = len(ext_cz_models) - iso_shifts = [0] * n_dist if iso_shifts is None else iso_shifts - weights = np.ones(n_dist) if weights is None else np.asarray(weights) - weights /= weights.sum() # Normalize weights array to total of 1 - - params = Parameters() - for i in range(n_dist): - dominant_tensor = ext_cz_models[i].symmetric_tensor - if isinstance(dominant_tensor, dict): # Convert to a SymmetricTensor object - dominant_tensor = SymmetricTensor(**dominant_tensor) - - if tensor_type == "shielding": - params.add(f"ext_czjzek_{i}_zeta0", value=dominant_tensor.zeta) - elif tensor_type == "quadrupolar": - params.add(f"ext_czjzek_{i}_Cq0", value=dominant_tensor.Cq) - - params.add(f"ext_czjzek_{i}_eta0", value=dominant_tensor.eta, min=0, max=1) - params.add(f"ext_czjzek_{i}_epsilon", value=ext_cz_models[i].eps, min=0) - params.add(f"ext_czjzek_{i}_iso_shift", value=iso_shifts[i]) - params.add(f"ext_czjzek_{i}_weight", value=weights[i], min=0, max=1) - - # Set last weight parameter as a non-free variable - expr = "-".join([f"ext_czjzek_{i}_weight" for i in range(n_dist - 1)]) - expr = "1" if expr == "" else f"1-{expr}" - params[f"ext_czjzek_{n_dist-1}_weight"].vary = False - params[f"ext_czjzek_{n_dist-1}_weight"].expr = expr - - # Add SignalProcessor parameters, if requested - if processor is not None: - _make_params_single_processor(params, processor, 0) - - return params - - -def _make_spectrum_from_extended_Czjzek_distribution_parameters( - params: Parameters, - exp_spectrum: cp.CSDM, - pos: tuple, - kernel: np.ndarray, - models: list, - larmor_freq_Hz: float, - processor: sp.SignalProcessor = None, - tensor_type: str = "shielding", -) -> cp.CSDM: - """Helper function for generating a spectrum from a set of LMfit Parameters and - and arguments for defining the grid, kernel, etc. The other functions used in least- - squares minimization use this function to reduce code overlap. - - Arguments: - (Parameters) params: The LMfit parameters object holding parameters used during - the least-squares minimization. - (cp.CSDM) exp_spectrum: The experimental spectrum to fit to. - (tuple) pos: A tuple of two np.ndarray objects defining the grid on which to - sample the distribution. - (np.ndarray) kernel: The pre-computed lineshape kernel. The kernel needs to be - defined on the same grid defined by pos. - (list) models: A list of distribution models. - (float) larmor_freq_Hz: This value is used in conjunction with the FFT shift - theorem to apply an isotropic chemical shift to each distribution. - (sp.SignalProcessor) processor: A - :py:class:~`mrsimulator.signal_processor.Processor` object used to apply - post-simulation signal processing to the resulting spectrum. - TODO: expand processor to apply to individual distributions (as a list) - - Returns: - Guess spectrum as a cp.CSDM object - - """ - n_dist = len(models) - if tensor_type not in {"shielding", "quadrupolar"}: - raise ValueError(f"Unknown `tensor_type` value of {tensor_type}.") - - guess_spectrum = exp_spectrum.copy() - guess_spectrum.y[0].components[:] = 0 # Initialize guess spectrum with zeros - - for i in range(n_dist): - zeta0 = ( - params[f"ext_czjzek_{i}_zeta0"].value - if tensor_type == "shielding" - else params[f"ext_czjzek_{i}_Cq0"].value - ) - eta0 = params[f"ext_czjzek_{i}_eta0"].value - epsilon = params[f"ext_czjzek_{i}_epsilon"].value - - # Draw the pdf from the set of static random tensors - tensor = {"Cq": zeta0, "eta": eta0} - models[i].symmetric_tensor = tensor - models[i].eps = epsilon - _, _, amp = models[i].pdf(pos, size=200_000) - - # Dot amplitude with kernel, then package as CSDM object - spec_tmp = guess_spectrum.copy() - spec_tmp.y[0].components[0] = np.dot(amp.ravel(), kernel) - - # Apply isotropic chemical shift to distribution using FFT shift theorem - spec_tmp = _apply_iso_shift( - csdm_obj=spec_tmp, - iso_shift_ppm=params[f"ext_czjzek_{i}_iso_shift"].value, - larmor_freq_Hz=larmor_freq_Hz, - ) - - guess_spectrum += spec_tmp.real * params[f"ext_czjzek_{i}_weight"].value - - if processor is not None: - _update_processors_from_LMFIT_params(params, [processor]) - guess_spectrum = processor.apply_operations(dataset=guess_spectrum).real - - return guess_spectrum - - -def LMFIT_min_function_extended_Czjzek( - params: Parameters, - exp_spectrum: cp.CSDM, - pos: tuple, - kernel: np.ndarray, - models: list, - larmor_freq_Hz: float, - processor: sp.SignalProcessor = None, - tensor_type: str = "shielding", -) -> np.ndarray: - """The minimization routine for fitting a set of Czjzek models to an experimental - spectrum. - - Arguments: - (Parameters) params: The LMfit parameters object holding parameters used during - the least-squares minimization. - (cp.CSDM) exp_spectrum: The experimental spectrum to fit to. - (tuple) pos: A tuple of two np.ndarray objects defining the grid on which to - sample the distribution. - (np.ndarray) kernel: The pre-computed lineshape kernel. The kernel needs to be - defined on the same grid defined by pos. - (list) models: A list of distribution models. - (float) larmor_freq_Hz: This value is used in conjunction with the FFT shift - theorem to apply an isotropic chemical shift to each distribution. - (sp.SignalProcessor) processor: A - :py:class:~`mrsimulator.signal_processor.Processor` object used to apply - post-simulation signal processing to the resulting spectrum. - TODO: expand processor to apply to individual distributions (as a list) - (str) tensor_type: A string literal describing if the Extended Czjzek models - define a distribution of symmetric shielding or quadrupolar tensors. The - allowed values are `shielding` and `quadrupolar`. - - Returns: - A residual array as a numpy array. - """ - guess_spectrum = _make_spectrum_from_extended_Czjzek_distribution_parameters( + guess_spectrum = _generate_distribution_spectrum( params, exp_spectrum, pos, @@ -943,13 +673,12 @@ def LMFIT_min_function_extended_Czjzek( models, larmor_freq_Hz, processor, - tensor_type, ) return (exp_spectrum - guess_spectrum).y[0].components[0] -def bestfit_extended_Czjzek( +def bestfit_dist( params: Parameters, exp_spectrum: cp.CSDM, pos: tuple, @@ -957,10 +686,9 @@ def bestfit_extended_Czjzek( models: list, larmor_freq_Hz: float, processor: sp.SignalProcessor = None, - tensor_type: str = "shielding", ) -> cp.CSDM: - """Returns the bestfit spectrum as a CSDM object""" - return _make_spectrum_from_extended_Czjzek_distribution_parameters( + """Returns the best-fit spectrum as a CSDM object""" + return _generate_distribution_spectrum( params, exp_spectrum, pos, @@ -968,11 +696,10 @@ def bestfit_extended_Czjzek( models, larmor_freq_Hz, processor, - tensor_type, ) -def residuals_extended_Czjzek( +def residuals_dist( params: Parameters, exp_spectrum: cp.CSDM, pos: tuple, @@ -980,10 +707,9 @@ def residuals_extended_Czjzek( models: list, larmor_freq_Hz: float, processor: sp.SignalProcessor = None, - tensor_type: str = "shielding", ) -> cp.CSDM: """Returns the residuals spectrum as a CSDM object""" - bestfit = _make_spectrum_from_extended_Czjzek_distribution_parameters( + bestfit = _generate_distribution_spectrum( params, exp_spectrum, pos, @@ -991,7 +717,6 @@ def residuals_extended_Czjzek( models, larmor_freq_Hz, processor, - tensor_type, ) return exp_spectrum - bestfit From b19e26a703dc3640632c79b65e34f098a2f0eba3 Mon Sep 17 00:00:00 2001 From: Deepansh Srivastava Date: Tue, 19 Mar 2024 19:45:26 -0400 Subject: [PATCH 11/21] more simplification --- .../1D_fitting/plot_7_27Al_Extended_Czjzek.py | 23 ++++---- src/mrsimulator/simulator/__init__.py | 2 + src/mrsimulator/utils/spectral_fitting.py | 52 +++++-------------- 3 files changed, 30 insertions(+), 47 deletions(-) diff --git a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py index a80b31ed0..d617614f6 100644 --- a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py +++ b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py @@ -20,6 +20,7 @@ import csdmpy as cp import matplotlib.pyplot as plt import mrsimulator.signal_processor as sp +from mrsimulator.simulator import Simulator import mrsimulator.utils.spectral_fitting as sf from lmfit import Minimizer @@ -78,7 +79,10 @@ # Create a Method object to simulate the spectrum spectral_dimensions = get_spectral_dimensions(exp_data) method = BlochDecayCTSpectrum( - channels=["27Al"], rotor_frequency=14.2e3, spectral_dimensions=spectral_dimensions + channels=["27Al"], + rotor_frequency=14.2e3, + spectral_dimensions=spectral_dimensions, + experiment=exp_data, ) # Define a polar grid for the lineshape kernel @@ -109,14 +113,14 @@ # reflect any changes made in the minimization. # Create initial guess ExtCzjzekDistribution -tensor_guess = {"Cq": 1e6, "eta": 0.3} ext_cz_model = ExtCzjzekDistribution( - mean_isotropic_chemical_shift=60, - symmetric_tensor=tensor_guess, - eps=6, + mean_isotropic_chemical_shift=60.0, # in ppm + symmetric_tensor={"Cq": 1e6, "eta": 0.3}, # Cq in Hz + eps=6.0, polar=True, cache=True, ) +all_models = [ext_cz_model] processor = sp.SignalProcessor( operations=[ @@ -127,20 +131,21 @@ ] ) +sim = Simulator(spin_system_models=[ext_cz_model], methods=[method]) +sim.config.number_of_sidebands = 8 + # Make the Parameters object params = sf.make_LMFIT_distribution_params( - distribution_models=[ext_cz_model], + distribution_models=all_models, processor=processor, ) # Additional keyword arguments passed to best-fit and residual functions. addtl_sf_kwargs = dict( - exp_spectrum=exp_data, pos=pos, kernel=kernel, - models=[ext_cz_model], - larmor_freq_Hz=method.channels[0].larmor_freq(B0=method.magnetic_flux_density), + sim=sim, processor=processor, ) diff --git a/src/mrsimulator/simulator/__init__.py b/src/mrsimulator/simulator/__init__.py index c8d53fe22..b650f1452 100644 --- a/src/mrsimulator/simulator/__init__.py +++ b/src/mrsimulator/simulator/__init__.py @@ -1,6 +1,7 @@ """Base Simulator class.""" import json from copy import deepcopy +from typing import Any from typing import List import csdmpy as cp @@ -130,6 +131,7 @@ class Simulator(Parseable): """ spin_systems: List[SpinSystem] = [] + spin_system_models: List[Any] = [] methods: List[Method] = [] config: ConfigSimulator = ConfigSimulator() diff --git a/src/mrsimulator/utils/spectral_fitting.py b/src/mrsimulator/utils/spectral_fitting.py index 6d66a127a..603f8ee74 100644 --- a/src/mrsimulator/utils/spectral_fitting.py +++ b/src/mrsimulator/utils/spectral_fitting.py @@ -575,11 +575,9 @@ def make_LMFIT_distribution_params( def _generate_distribution_spectrum( params: Parameters, - exp_spectrum: cp.CSDM, pos: tuple, kernel: np.ndarray, - distribution_models: list, - larmor_freq_Hz: float, + sim: Simulator, processor: sp.SignalProcessor = None, ) -> cp.CSDM: """Helper function for generating a spectrum from a set of LMfit Parameters and @@ -605,6 +603,11 @@ def _generate_distribution_spectrum( Guess spectrum as a cp.CSDM object """ + distribution_models = sim.spin_system_models + method = sim.methods[0] + larmor_freq_Hz = method.channels[0].larmor_freq(B0=method.magnetic_flux_density) + exp_spectrum = method.experiment + guess_spectrum = exp_spectrum.copy() guess_spectrum.y[0].components[:] = 0 # Initialize guess spectrum with zeros @@ -634,11 +637,9 @@ def _generate_distribution_spectrum( def LMFIT_min_function_dist( params: Parameters, - exp_spectrum: cp.CSDM, pos: tuple, kernel: np.ndarray, - models: list, - larmor_freq_Hz: float, + sim: Simulator, processor: sp.SignalProcessor = None, ) -> np.ndarray: """The minimization routine for fitting a set of Czjzek models to an experimental @@ -666,57 +667,32 @@ def LMFIT_min_function_dist( A residual array as a numpy array. """ guess_spectrum = _generate_distribution_spectrum( - params, - exp_spectrum, - pos, - kernel, - models, - larmor_freq_Hz, - processor, + params, pos, kernel, sim, processor ) - + exp_spectrum = sim.methods[0].experiment return (exp_spectrum - guess_spectrum).y[0].components[0] def bestfit_dist( params: Parameters, - exp_spectrum: cp.CSDM, pos: tuple, kernel: np.ndarray, - models: list, - larmor_freq_Hz: float, + sim: Simulator, processor: sp.SignalProcessor = None, ) -> cp.CSDM: """Returns the best-fit spectrum as a CSDM object""" - return _generate_distribution_spectrum( - params, - exp_spectrum, - pos, - kernel, - models, - larmor_freq_Hz, - processor, - ) + return _generate_distribution_spectrum(params, pos, kernel, sim, processor) def residuals_dist( params: Parameters, - exp_spectrum: cp.CSDM, pos: tuple, kernel: np.ndarray, - models: list, - larmor_freq_Hz: float, + sim: Simulator, processor: sp.SignalProcessor = None, ) -> cp.CSDM: """Returns the residuals spectrum as a CSDM object""" - bestfit = _generate_distribution_spectrum( - params, - exp_spectrum, - pos, - kernel, - models, - larmor_freq_Hz, - processor, - ) + bestfit = _generate_distribution_spectrum(params, pos, kernel, sim, processor) + exp_spectrum = sim.methods[0].experiment return exp_spectrum - bestfit From 1dfccf3a60045191de3e5316da19f9e046bb806c Mon Sep 17 00:00:00 2001 From: Deepansh Srivastava Date: Tue, 19 Mar 2024 22:13:51 -0400 Subject: [PATCH 12/21] set model to czjzek --- .../1D_fitting/plot_7_27Al_Extended_Czjzek.py | 38 +++++++++---------- .../models/analytical_distributions.py | 2 +- src/mrsimulator/models/czjzek.py | 1 - 3 files changed, 18 insertions(+), 23 deletions(-) diff --git a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py index d617614f6..dc06f677f 100644 --- a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py +++ b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py @@ -1,18 +1,18 @@ #!/usr/bin/env python """ -Fitting an Extended Czjzek Model -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +Fitting a Czjzek Model +^^^^^^^^^^^^^^^^^^^^^^ """ # %% -# In this example, we display the capability of mrsimulator to fit an Extended Czjzek -# distribution of quadrupolar tensors to an experimental $^{27}\text{Al}$ MAS spectrum -# of a phospho-aluminosilicate glass. Setting up a least-squares minimization for an -# Extended Czjzek distribution is slightly more complicated than that of a crystalline +# In this example, we display the capability of mrsimulator to fit a lineshape from a +# Czjzek distribution of quadrupolar tensors to an experimental $^{27}\text{Al}$ MAS +# spectrum of a phospho-aluminosilicate glass. Setting up a least-squares minimization +# for a Czjzek distribution is slightly different than that of a crystalline # solid. There are 4 major steps involved in the process: # -# 1. Importing the experimental dataset -# 2. Generating a pre-computed line shape kernel for the experiment -# 3. Creating parameters for the Extended Czjzek model from an initial guess +# 1. Importing the experimental dataset, +# 2. Generating a pre-computed line shape kernel for the experiment, +# 3. Creating parameters for the Czjzek model from an initial guess, # 4. Minimizing and visualizing. # # Below, we first import the requisite libraries, functions, and classes. @@ -27,7 +27,7 @@ from mrsimulator.method.lib import BlochDecayCTSpectrum from mrsimulator.utils import get_spectral_dimensions -from mrsimulator.models.czjzek import ExtCzjzekDistribution +from mrsimulator.models.czjzek import CzjzekDistribution from mrsimulator.models.utils import generate_lineshape_kernel # sphinx_gallery_thumbnail_number = 3 @@ -36,8 +36,7 @@ # Import the experimental dataset # ------------------------------- # -# Below we import and visualize the experimental dataset. Line 3 is used to set the -# maximum intensity of the spectrum to 1, but has no other bering on the minimization. +# Below we import and visualize the experimental dataset. host = "http://ssnmr.org/sites/default/files/mrsimulator/" filename = "20K_20Al_10P_50Si_HahnEcho_27Al.csdf" exp_data = cp.load(host + filename).real @@ -55,7 +54,7 @@ # Generating a line shape kernel # ------------------------------ # -# A spectrum arising from an Extended Czjzek distribution can be modeled by drawing a +# A spectrum arising from a Czjzek distribution can be modeled by drawing a # random set of tensors, binning those tensors into a distribution on some grid of # tensors parameters, then simulating the spectra of spin systems whose abundances # correspond to the amplitude of that distribution. @@ -113,10 +112,9 @@ # reflect any changes made in the minimization. # Create initial guess ExtCzjzekDistribution -ext_cz_model = ExtCzjzekDistribution( +ext_cz_model = CzjzekDistribution( mean_isotropic_chemical_shift=60.0, # in ppm - symmetric_tensor={"Cq": 1e6, "eta": 0.3}, # Cq in Hz - eps=6.0, + sigma=1.4e6, polar=True, cache=True, ) @@ -125,9 +123,9 @@ processor = sp.SignalProcessor( operations=[ sp.IFFT(), - sp.apodization.Gaussian(FWHM="420 Hz"), + sp.apodization.Gaussian(FWHM="600 Hz"), sp.FFT(), - sp.Scale(factor=25), + sp.Scale(factor=30), ] ) @@ -140,7 +138,6 @@ processor=processor, ) - # Additional keyword arguments passed to best-fit and residual functions. addtl_sf_kwargs = dict( pos=pos, @@ -181,7 +178,6 @@ loss="soft_l1", ) - minner = Minimizer( sf.LMFIT_min_function_dist, params, @@ -215,7 +211,7 @@ for i, model in enumerate([ext_cz_model]): model.update_lmfit_params(result.params, i) -xx, yy, amp = ext_cz_model.pdf(pos=pos, size=10000000) +xx, yy, amp = ext_cz_model.pdf(pos=pos) plt.figure(figsize=(5, 5)) plt.contourf(xx / 1e6, yy / 1e6, amp) diff --git a/src/mrsimulator/models/analytical_distributions.py b/src/mrsimulator/models/analytical_distributions.py index 43ce106a5..fb8ecc200 100644 --- a/src/mrsimulator/models/analytical_distributions.py +++ b/src/mrsimulator/models/analytical_distributions.py @@ -89,4 +89,4 @@ def czjzek_polar(sigma: float, pos: list): ) hist_x_y += hist_x_y.T hist_x_y /= hist_x_y.sum() - return x, y, hist_x_y + return pos[0], pos[1], hist_x_y diff --git a/src/mrsimulator/models/czjzek.py b/src/mrsimulator/models/czjzek.py index 9cd2f8b94..fe05796f1 100644 --- a/src/mrsimulator/models/czjzek.py +++ b/src/mrsimulator/models/czjzek.py @@ -226,7 +226,6 @@ def update_lmfit_params(self, params, i): prefix = self.param_prefix() self.sigma = params[f"{prefix}_{i}_sigma"].value - self.eps = params[f"{prefix}_{i}_epsilon"].value self.mean_isotropic_chemical_shift = params[f"{prefix}_{i}_iso_shift"].value self.abundance = params[f"{prefix}_{i}_weight"].value From bd6e1d9ed9b0a49ad5628738969bde7412ce6ce8 Mon Sep 17 00:00:00 2001 From: deepanshs <21365911+deepanshs@users.noreply.github.com> Date: Tue, 19 Mar 2024 23:43:04 -0400 Subject: [PATCH 13/21] remove array copy statements --- .../1D_fitting/plot_7_27Al_Extended_Czjzek.py | 18 +++++++++--------- src/mrsimulator/utils/spectral_fitting.py | 11 ++++++----- 2 files changed, 15 insertions(+), 14 deletions(-) diff --git a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py index dc06f677f..96a2d7705 100644 --- a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py +++ b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py @@ -111,14 +111,14 @@ # many of the spectral fitting function takes in. This dictionary needs to be updated to # reflect any changes made in the minimization. -# Create initial guess ExtCzjzekDistribution -ext_cz_model = CzjzekDistribution( +# Create initial guess CzjzekDistribution +cz_model = CzjzekDistribution( mean_isotropic_chemical_shift=60.0, # in ppm sigma=1.4e6, polar=True, cache=True, ) -all_models = [ext_cz_model] +all_models = [cz_model] processor = sp.SignalProcessor( operations=[ @@ -129,7 +129,7 @@ ] ) -sim = Simulator(spin_system_models=[ext_cz_model], methods=[method]) +sim = Simulator(spin_system_models=all_models, methods=[method]) sim.config.number_of_sidebands = 8 # Make the Parameters object @@ -150,7 +150,7 @@ guess = sf.bestfit_dist(params=params, **addtl_sf_kwargs) residuals = sf.residuals_dist(params=params, **addtl_sf_kwargs) -plt.figure(figsize=(5, 4)) +plt.figure(figsize=(4, 3)) ax = plt.subplot(projection="csdm") ax.plot(exp_data, "k", alpha=0.5, label="Experiment") ax.plot(guess, "r", alpha=0.3, label="Guess") @@ -193,7 +193,7 @@ bestfit = sf.bestfit_dist(params=result.params, **addtl_sf_kwargs) residuals = sf.residuals_dist(params=result.params, **addtl_sf_kwargs) -plt.figure(figsize=(5, 4)) +plt.figure(figsize=(4, 3)) ax = plt.subplot(projection="csdm") ax.plot(exp_data, "k", alpha=0.5, label="Experiment") ax.plot(bestfit, "r", alpha=0.3, label="Guess") @@ -208,12 +208,12 @@ # Plot the best-fit distribution # ------------------------------ # -for i, model in enumerate([ext_cz_model]): +for i, model in enumerate(all_models): model.update_lmfit_params(result.params, i) -xx, yy, amp = ext_cz_model.pdf(pos=pos) +xx, yy, amp = cz_model.pdf(pos=pos) -plt.figure(figsize=(5, 5)) +plt.figure(figsize=(4, 3)) plt.contourf(xx / 1e6, yy / 1e6, amp) plt.xlabel("x / MHz ") plt.ylabel("y / MHz") diff --git a/src/mrsimulator/utils/spectral_fitting.py b/src/mrsimulator/utils/spectral_fitting.py index 603f8ee74..930772c09 100644 --- a/src/mrsimulator/utils/spectral_fitting.py +++ b/src/mrsimulator/utils/spectral_fitting.py @@ -610,23 +610,24 @@ def _generate_distribution_spectrum( guess_spectrum = exp_spectrum.copy() guess_spectrum.y[0].components[:] = 0 # Initialize guess spectrum with zeros + dv = cp.as_dependent_variable(np.empty(guess_spectrum.y[0].components.size)) for i, model in enumerate(distribution_models): model.update_lmfit_params(params, i) _, _, amp = model.pdf(pos) # Dot amplitude with kernel, then package as CSDM object - spec_tmp = guess_spectrum.copy() - spec_tmp.y[0].components[0] = np.dot(amp.flatten(), kernel) + spec_tmp = cp.CSDM(dimensions=exp_spectrum.x, dependent_variables=[dv]) + spec_tmp.y[0].components[0] = np.dot(amp.ravel(), kernel) # Apply isotropic chemical shift to distribution using FFT shift theorem spec_tmp = _apply_iso_shift( csdm_obj=spec_tmp, iso_shift_ppm=model.mean_isotropic_chemical_shift, larmor_freq_Hz=larmor_freq_Hz, - ) - - guess_spectrum += spec_tmp.real * model.abundance + ).real + spec_tmp *= model.abundance + guess_spectrum.y[0].components[0] += spec_tmp.y[0].components[0] if processor is not None: _update_processors_from_LMFIT_params(params, [processor]) From 15cc7b6b5fd7e439d01e4315d870d5fd854f9532 Mon Sep 17 00:00:00 2001 From: deepanshs <21365911+deepanshs@users.noreply.github.com> Date: Sun, 24 Mar 2024 22:42:47 -0400 Subject: [PATCH 14/21] more simplification --- .../1D_fitting/plot_7_27Al_Extended_Czjzek.py | 126 ++++++++---------- src/mrsimulator/models/czjzek.py | 4 +- src/mrsimulator/models/utils.py | 87 +++++++----- src/mrsimulator/utils/spectral_fitting.py | 106 +++++++++------ 4 files changed, 178 insertions(+), 145 deletions(-) diff --git a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py index 96a2d7705..ae4587237 100644 --- a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py +++ b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py @@ -4,23 +4,23 @@ ^^^^^^^^^^^^^^^^^^^^^^ """ # %% -# In this example, we display the capability of mrsimulator to fit a lineshape from a -# Czjzek distribution of quadrupolar tensors to an experimental $^{27}\text{Al}$ MAS -# spectrum of a phospho-aluminosilicate glass. Setting up a least-squares minimization -# for a Czjzek distribution is slightly different than that of a crystalline -# solid. There are 4 major steps involved in the process: +# In this example, we illustrate an application of mrsimulator in fitting a lineshape +# from a Czjzek distribution of quadrupolar tensors to an experimental +# :math:`^{27}\text{Al}` MAS spectrum of a phospho-aluminosilicate glass. Setting up +# the least-squares minimization for a distribution of spin systems is slightly +# different than that of a crystalline solid. + +# There are 4 steps involved in the processs: +# - Importing the experimental dataset, +# - Generating a pre-computed line shape kernel of subspectra, +# - Creating parameters for the Czjzek distribution model from an initial guess, +# - Minimizing and visualizing. # -# 1. Importing the experimental dataset, -# 2. Generating a pre-computed line shape kernel for the experiment, -# 3. Creating parameters for the Czjzek model from an initial guess, -# 4. Minimizing and visualizing. -# -# Below, we first import the requisite libraries, functions, and classes. import numpy as np import csdmpy as cp import matplotlib.pyplot as plt import mrsimulator.signal_processor as sp -from mrsimulator.simulator import Simulator +from mrsimulator.simulator.config import ConfigSimulator import mrsimulator.utils.spectral_fitting as sf from lmfit import Minimizer @@ -28,7 +28,7 @@ from mrsimulator.utils import get_spectral_dimensions from mrsimulator.models.czjzek import CzjzekDistribution -from mrsimulator.models.utils import generate_lineshape_kernel +from mrsimulator.models.utils import LineShapeKernel # sphinx_gallery_thumbnail_number = 3 @@ -54,27 +54,31 @@ # Generating a line shape kernel # ------------------------------ # -# A spectrum arising from a Czjzek distribution can be modeled by drawing a -# random set of tensors, binning those tensors into a distribution on some grid of -# tensors parameters, then simulating the spectra of spin systems whose abundances -# correspond to the amplitude of that distribution. +# The Czjzek distribution is a statistical model that represents a five-dimensional +# multivariate normal distribution of second-rank tensor components. These random +# second-rank tensors are converted into corresponding anisotropy and asymmetry +# parameters using the Haeberlen convention. The resulting parameters are then binned +# on a two-dimensional grid, forming the Czjzek probability distribution. The Czjzek +# spectra is which are the probability weighted sum of lineshapes at each grid point. +# +# However, simulating the spectra of spin systems at each minimization step for the +# least-squares fitting is computationally expensive. A more efficient way is to +# pre-define a grid system for the tensor parameters, simulate a library of sub-spectra +# for each grid point, and only update the probability distribution during each +# minimization step. # -# For least-squares fitting, however, simulating the spectra of many spin systems during -# each minimization step would be computationally costly, especially since a -# least-squares minimization could take hundreds of steps to complete. -# A more efficient way would be to define a grid for the tensor parameters, simulate a -# spectrum for each point on the grid before minimization, and only update the -# probability distribution on that grid during each minimization step. +# To simulate the spectra of the given experiment, we create a Method object. Using +# this method, we generate a kernel of lineshape sub-spectra defined on a polar grid by +# using the LineShapeKernel class. The argument "method" is the mrsimulator method +# object used for the simulation, "pos" is a tuple of coordinates defining the +# two-dimensional grid, "polar" defines a polar parameter coordinate space, and +# "config" is the Simulator config object used in lineshape simulation. Here, we use +# the "config" argument to set the number of sidebands as 8. # -# Below, we create a `Method` object for simulating the spectra of the given experiment, -# define a polar grid for generating the line shape kernel, and finally call the kernel -# generation method -- `generate_lineshape_kernel`. -# The grid we defined is in polar coordinates, so we pass `polar=True` to the kernel -# generation method; if the grid were defined for the Haeberlen components of the -# tensor, then pass `polar=False`. -# We are also interested in the distribution of the EFG (quadrupolar) tensor, hence -# `tensor_type="quadrupolar"`, but a symmetric shielding lineshape kernel can also be -# generated by passing `tensor_type="shielding"`. +# The kernel is generated with the "generate_lineshape()" function, where we pass the +# "tensor_type='quadrupolar'" argument to specify a quadrupolar parameter grid. A +# symmetric shielding lineshape kernel can also be generated by specifying + # Create a Method object to simulate the spectrum spectral_dimensions = get_spectral_dimensions(exp_data) method = BlochDecayCTSpectrum( @@ -89,34 +93,23 @@ y = np.linspace(0, 2e7, num=36) pos = (x, y) -# Call the kernel generation function with the correct arguments -kernel = generate_lineshape_kernel( - pos=pos, method=method, polar=True, tensor_type="quadrupolar" -) +# Generate the kernel +sim_config = ConfigSimulator(number_of_sidebands=8) +quad_kernel = LineShapeKernel(method=method, pos=pos, polar=True, config=sim_config) +quad_kernel.generate_lineshape(tensor_type="quadrupolar") print("Desired Kernel shape: ", (x.size * y.size, spectral_dimensions[0]["count"])) -print("Actual Kernel shape: ", kernel.shape) +print("Actual Kernel shape: ", quad_kernel.kernel.shape) # %% # Create a Parameters object # -------------------------- -# Next, we create an instance of the `ExtCzjzekDistribution` class with initial guess -# values along with a `SignalProcessor` object. A set of minimization parameters are -# then created from this initial guess. -# -# The initial guess and residual spectrum is also plotted to judge the quality of the -# initial guess. -# -# Note that the variable `addtl_sf_kwargs` holds some additional keyword arguments that -# many of the spectral fitting function takes in. This dictionary needs to be updated to -# reflect any changes made in the minimization. +# Next, we create an instance of the `CzjzekDistribution` class with initial guess +# values along with a `SignalProcessor` object. # Create initial guess CzjzekDistribution cz_model = CzjzekDistribution( - mean_isotropic_chemical_shift=60.0, # in ppm - sigma=1.4e6, - polar=True, - cache=True, + mean_isotropic_chemical_shift=60.0, sigma=1.4e6, polar=True ) all_models = [cz_model] @@ -129,26 +122,23 @@ ] ) -sim = Simulator(spin_system_models=all_models, methods=[method]) -sim.config.number_of_sidebands = 8 - -# Make the Parameters object -params = sf.make_LMFIT_distribution_params( - distribution_models=all_models, - processor=processor, -) +# %% +# Make the Parameters object and simulate a guess spectrum. +# Note that the variable `sf_kwargs` holds some additional keyword arguments that +# many of the spectral fitting function takes in. This dictionary needs to be updated to +# reflect any changes made in the minimization. +params = sf.make_LMFIT_params(spin_system_models=all_models, processors=[processor]) # Additional keyword arguments passed to best-fit and residual functions. -addtl_sf_kwargs = dict( - pos=pos, - kernel=kernel, - sim=sim, +sf_kwargs = dict( + kernel=quad_kernel, + spin_system_models=all_models, processor=processor, ) # Make a guess and residuals spectrum from the initial guess -guess = sf.bestfit_dist(params=params, **addtl_sf_kwargs) -residuals = sf.residuals_dist(params=params, **addtl_sf_kwargs) +guess = sf.bestfit_dist(params=params, **sf_kwargs) +residuals = sf.residuals_dist(params=params, **sf_kwargs) plt.figure(figsize=(4, 3)) ax = plt.subplot(projection="csdm") @@ -181,7 +171,7 @@ minner = Minimizer( sf.LMFIT_min_function_dist, params, - fcn_kws=addtl_sf_kwargs, + fcn_kws=sf_kwargs, **scipy_minimization_kwargs, ) result = minner.minimize(method="least_squares") @@ -190,8 +180,8 @@ # %% # Plot the best-fit spectrum # -------------------------- -bestfit = sf.bestfit_dist(params=result.params, **addtl_sf_kwargs) -residuals = sf.residuals_dist(params=result.params, **addtl_sf_kwargs) +bestfit = sf.bestfit_dist(params=result.params, **sf_kwargs) +residuals = sf.residuals_dist(params=result.params, **sf_kwargs) plt.figure(figsize=(4, 3)) ax = plt.subplot(projection="csdm") diff --git a/src/mrsimulator/models/czjzek.py b/src/mrsimulator/models/czjzek.py index fe05796f1..fc68215cd 100644 --- a/src/mrsimulator/models/czjzek.py +++ b/src/mrsimulator/models/czjzek.py @@ -176,7 +176,7 @@ def __init__( mean_isotropic_chemical_shift: float = 0.0, abundance: float = 1.0, polar=False, - cache=False, + cache=True, ): super().__init__( cache_tensors=cache, @@ -279,7 +279,7 @@ def __init__( mean_isotropic_chemical_shift: float = 0.0, abundance: float = 1.0, polar=False, - cache=False, + cache=True, ): super().__init__( cache_tensors=cache, diff --git a/src/mrsimulator/models/utils.py b/src/mrsimulator/models/utils.py index cc96c6b71..91708d919 100644 --- a/src/mrsimulator/models/utils.py +++ b/src/mrsimulator/models/utils.py @@ -1,8 +1,11 @@ +from dataclasses import dataclass + import numpy as np from mrsimulator import Method from mrsimulator import Simulator from mrsimulator import Site from mrsimulator import SpinSystem +from mrsimulator.simulator import ConfigSimulator from mrsimulator.spin_system.isotope import Isotope __author__ = "Deepansh J. Srivastava" @@ -87,21 +90,24 @@ def x_y_to_zeta_eta(x, y): return zeta, eta -def _simulate_spectra_over_zeta_and_eta(ZZ, ee, method, tensor_type): +def _simulate_spectra_over_zeta_and_eta(ZZ, ee, method, tensor_type, config): """Helper function to generate the kernel""" isotope = Isotope.parse( method.channels[0] ).symbol # Grab isotope from Method object spin_systems = [ - SpinSystem(sites=[Site(isotope=isotope, quadrupolar=dict(Cq=Z, eta=e))]) - if tensor_type == "quadrupolar" - else SpinSystem( - sites=[Site(isotope=isotope, shielding_symmetric=dict(zeta=Z, eta=e))] + ( + SpinSystem(sites=[Site(isotope=isotope, quadrupolar=dict(Cq=Z, eta=e))]) + if tensor_type == "quadrupolar" + else SpinSystem( + sites=[Site(isotope=isotope, shielding_symmetric=dict(zeta=Z, eta=e))] + ) ) for Z, e in zip(ZZ.ravel(), ee.ravel()) ] sim = Simulator(spin_systems=spin_systems, methods=[method]) + sim.config = config sim.config.decompose_spectrum = "spin_system" sim.run(pack_as_csdm=False) @@ -109,37 +115,52 @@ def _simulate_spectra_over_zeta_and_eta(ZZ, ee, method, tensor_type): return amp -def generate_lineshape_kernel( - pos: tuple, method: Method, polar: bool, tensor_type: str = "shielding" -) -> np.ndarray: - """Pre-compute a lineshape kernel used during least-squares fitting of an - experimental spectrum. The isotope for the spin system is grabbed from the isotope - at index zero of the channels method - - Note: The lineshape kernel is currently limited to simulating the spectrum of an - uncoupled spin system for a single Method over a range of either chemical shielding - or quadrupolar tensors. Functionality may expand in the future +@dataclass +class LineShapeKernel: + """lineshape kernel object Arguments: - (tuple) pos: A set of numpy arrays defining the grid space - (mrsimulator.Method) method: The :py:class:`~mrsimulator.method.Method` used to - simulate the spectra. - (bool) polar: Weather the grid is defined in polar coordinates - (str) tensor_type: A string enumeration literal describing which type of tensor - to use in the simulation. The allowed values are `shielding` and - `quadrupolar`. - - Returns: - (np.ndarray) kernel: The simulated lineshape kernel + (tuple) pos: A tuple of numpy arrays defining the 2D grid space. + (bool) polar: If true, the grid is defined in polar coordinates. + (mrsimulator.Method) method: The :py:class:`~mrsimulator.method.Method` + used to simulate the spectra. + (ConfigSimulator) config: Simulator config to be used in simulation. """ - if tensor_type not in {"shielding", "quadrupolar"}: - raise ValueError(f"Unrecognized value of {tensor_type} for `tensor_type.") - # If in polar coordinates, then (ZZ, ee) right now is really (xx, yy) - ZZ, ee = np.meshgrid(pos[0], pos[1], indexing="xy") + pos: list[np.ndarray] + method: Method + kernel: np.ndarray = None + polar: bool = False + config: ConfigSimulator = ConfigSimulator() + + def generate_lineshape(self, tensor_type: str = "shielding") -> np.ndarray: + """Pre-compute a lineshape kernel to use for the least-squares fitting of an + experimental spectrum. The isotope for the spin system is the isotope + at index zero of the method's channel. + + Note: The lineshape kernel is currently limited to simulating the spectrum of + an uncoupled spin system for a single Method over a range of nuclear shielding + or quadrupolar tensors. Functionality may expand in the future. + + Arguments: - # Convert polar coords to cartesian coords - if polar: - ZZ, ee = x_y_to_zeta_eta(ZZ.ravel(), ee.ravel()) + (str) tensor_type: A string enumeration literal describing the type of + tensor to use in the simulation. The allowed values are `shielding` and + `quadrupolar`. - return _simulate_spectra_over_zeta_and_eta(ZZ, ee, method, tensor_type) + Returns: + (np.ndarray) kernel: The simulated lineshape kernel + """ + if tensor_type not in {"shielding", "quadrupolar"}: + raise ValueError(f"Unrecognized value of {tensor_type} for `tensor_type.") + + # If in polar coordinates, then (ZZ, ee) right now is really (xx, yy) + ZZ, ee = np.meshgrid(self.pos[0], self.pos[1], indexing="xy") + + # Convert polar coords to cartesian coords + if self.polar: + ZZ, ee = x_y_to_zeta_eta(ZZ.ravel(), ee.ravel()) + + self.kernel = _simulate_spectra_over_zeta_and_eta( + ZZ, ee, self.method, tensor_type, self.config + ) diff --git a/src/mrsimulator/utils/spectral_fitting.py b/src/mrsimulator/utils/spectral_fitting.py index 930772c09..79cc16679 100644 --- a/src/mrsimulator/utils/spectral_fitting.py +++ b/src/mrsimulator/utils/spectral_fitting.py @@ -3,6 +3,7 @@ import numpy as np from lmfit import Parameters from mrsimulator import Simulator +from mrsimulator.models.utils import LineShapeKernel __author__ = ["Maxwell C Venetos", "Deepansh Srivastava"] __email__ = ["maxvenetos@gmail.com", "srivastava.89@osu.edu"] @@ -219,7 +220,9 @@ def _get_simulator_object_value(sim, string): return obj -def make_simulator_params(sim: Simulator, include={}): +def make_simulator_params( + sim: Simulator = Simulator(), spin_system_models: list = [], include={} +): """Parse the Simulator object for a list of LMFIT parameters. Args: @@ -232,13 +235,24 @@ def make_simulator_params(sim: Simulator, include={}): temp_list = _traverse_dictionaries(_list_of_dictionaries(sim.spin_systems)) # get total abundance scaling factor - length = len(sim.spin_systems) - abundance_scale = 100 / sum(sys.abundance for sys in sim.spin_systems) + sys_length = len(sim.spin_systems) + sys_model_length = len(spin_system_models) + total_abundance = sum(sys.abundance for sys in sim.spin_systems) + total_abundance += sum(sys.abundance for sys in spin_system_models) + abundance_scale = 100 / total_abundance # expression for the last abundance. - last_abundance = f"{length - 1}_abundance" - expression = "-".join([f"{START}{i}_abundance" for i in range(length - 1)]) - expression = "100" if expression == "" else f"100-{expression}" + if sys_length > 0: + last_abundance = f"{sys_length - 1}_abundance" + expression = "-".join([f"{START}{i}_abundance" for i in range(sys_length - 1)]) + expression = "100" if expression == "" else f"100-{expression}" + + skip_last = sys_length > 0 + if sys_model_length > 0: + param_dist = make_distribution_params( + spin_system_models, norm=abundance_scale, skip_last=skip_last + ) + _ = params.update(param_dist) for items in temp_list: value = _get_simulator_object_value(sim, items) @@ -297,7 +311,12 @@ def make_LMFIT_parameters(sim: Simulator, processors: list = None, include={}): return make_LMFIT_params(sim, processors) -def make_LMFIT_params(sim: Simulator, processors: list = None, include={}): +def make_LMFIT_params( + sim: Simulator = Simulator(), + processors: list = None, + spin_system_models: list = [], + include={}, +): r"""Parse the Simulator and PostSimulator objects for a list of LMFIT parameters. Args: @@ -325,7 +344,11 @@ def make_LMFIT_params(sim: Simulator, processors: list = None, include={}): LMFIT Parameters object. """ params = Parameters() - params.update(make_simulator_params(sim, include)) + params.update( + make_simulator_params( + sim=sim, spin_system_models=spin_system_models, include=include + ) + ) proc = make_signal_processor_params(processors) if processors is not None else None _ = params.update(proc) if proc is not None else None @@ -537,28 +560,28 @@ def _apply_iso_shift(csdm_obj, iso_shift_ppm, larmor_freq_Hz): return csdm_obj -def make_LMFIT_distribution_params( - distribution_models: list, processor: sp.SignalProcessor = None +def make_distribution_params( + spin_system_models: list, norm: float, skip_last: bool = False ) -> Parameters: """Generate LMfit Parameters object for spin system distribution model. The distribution has the following parameters: """ - n_dist = len(distribution_models) + n_dist = len(spin_system_models) - norm = 0.0 - for model in distribution_models: - norm += model.abundance + # norm = 0.0 + # for model in distribution_models: + # norm += model.abundance expr_terms = [] params = Parameters() - for i, model in enumerate(distribution_models): + for i, model in enumerate(spin_system_models): # normalize the abundance model.abundance /= norm params = model.get_lmfit_params(params, i) # Set last weight parameter as a non-free variable model_prefix = model.param_prefix() - if i < n_dist - 1: + if i < n_dist - 1 or skip_last: expr_terms.append(f"{model_prefix}_{i}_weight") else: expr = "-".join(expr_terms) @@ -566,18 +589,17 @@ def make_LMFIT_distribution_params( params[f"{model_prefix}_{i}_weight"].vary = False params[f"{model_prefix}_{i}_weight"].expr = expr - # Add SignalProcessor parameters, if requested - if processor is not None: - _make_params_single_processor(params, processor, 0) + # # Add SignalProcessor parameters, if requested + # if processor is not None: + # _make_params_single_processor(params, processor, 0) return params def _generate_distribution_spectrum( params: Parameters, - pos: tuple, - kernel: np.ndarray, - sim: Simulator, + kernel: LineShapeKernel, + spin_system_models: list, processor: sp.SignalProcessor = None, ) -> cp.CSDM: """Helper function for generating a spectrum from a set of LMfit Parameters and @@ -603,8 +625,7 @@ def _generate_distribution_spectrum( Guess spectrum as a cp.CSDM object """ - distribution_models = sim.spin_system_models - method = sim.methods[0] + method = kernel.method larmor_freq_Hz = method.channels[0].larmor_freq(B0=method.magnetic_flux_density) exp_spectrum = method.experiment @@ -612,13 +633,13 @@ def _generate_distribution_spectrum( guess_spectrum.y[0].components[:] = 0 # Initialize guess spectrum with zeros dv = cp.as_dependent_variable(np.empty(guess_spectrum.y[0].components.size)) - for i, model in enumerate(distribution_models): + for i, model in enumerate(spin_system_models): model.update_lmfit_params(params, i) - _, _, amp = model.pdf(pos) + _, _, amp = model.pdf(kernel.pos) # Dot amplitude with kernel, then package as CSDM object spec_tmp = cp.CSDM(dimensions=exp_spectrum.x, dependent_variables=[dv]) - spec_tmp.y[0].components[0] = np.dot(amp.ravel(), kernel) + spec_tmp.y[0].components[0] = np.dot(amp.ravel(), kernel.kernel) # Apply isotropic chemical shift to distribution using FFT shift theorem spec_tmp = _apply_iso_shift( @@ -638,9 +659,8 @@ def _generate_distribution_spectrum( def LMFIT_min_function_dist( params: Parameters, - pos: tuple, - kernel: np.ndarray, - sim: Simulator, + kernel: LineShapeKernel, + spin_system_models: list, processor: sp.SignalProcessor = None, ) -> np.ndarray: """The minimization routine for fitting a set of Czjzek models to an experimental @@ -668,32 +688,34 @@ def LMFIT_min_function_dist( A residual array as a numpy array. """ guess_spectrum = _generate_distribution_spectrum( - params, pos, kernel, sim, processor + params, kernel, spin_system_models, processor ) - exp_spectrum = sim.methods[0].experiment + exp_spectrum = kernel.method.experiment return (exp_spectrum - guess_spectrum).y[0].components[0] def bestfit_dist( params: Parameters, - pos: tuple, - kernel: np.ndarray, - sim: Simulator, + kernel: LineShapeKernel, + spin_system_models: list, processor: sp.SignalProcessor = None, ) -> cp.CSDM: """Returns the best-fit spectrum as a CSDM object""" - return _generate_distribution_spectrum(params, pos, kernel, sim, processor) + return _generate_distribution_spectrum( + params, kernel, spin_system_models, processor + ) def residuals_dist( params: Parameters, - pos: tuple, - kernel: np.ndarray, - sim: Simulator, + kernel: LineShapeKernel, + spin_system_models: list, processor: sp.SignalProcessor = None, ) -> cp.CSDM: """Returns the residuals spectrum as a CSDM object""" - bestfit = _generate_distribution_spectrum(params, pos, kernel, sim, processor) + bestfit_ = _generate_distribution_spectrum( + params, kernel, spin_system_models, processor + ) - exp_spectrum = sim.methods[0].experiment - return exp_spectrum - bestfit + exp_spectrum = kernel.method.experiment + return exp_spectrum - bestfit_ From be79bc3bb0cbc782a544076052b29eb2b66af8a9 Mon Sep 17 00:00:00 2001 From: deepanshs <21365911+deepanshs@users.noreply.github.com> Date: Sun, 24 Mar 2024 23:01:15 -0400 Subject: [PATCH 15/21] fix dataclass field issue --- src/mrsimulator/models/utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/mrsimulator/models/utils.py b/src/mrsimulator/models/utils.py index 91708d919..35821cddc 100644 --- a/src/mrsimulator/models/utils.py +++ b/src/mrsimulator/models/utils.py @@ -1,4 +1,5 @@ from dataclasses import dataclass +from dataclasses import field import numpy as np from mrsimulator import Method @@ -127,11 +128,11 @@ class LineShapeKernel: (ConfigSimulator) config: Simulator config to be used in simulation. """ - pos: list[np.ndarray] + pos: list method: Method kernel: np.ndarray = None polar: bool = False - config: ConfigSimulator = ConfigSimulator() + config: ConfigSimulator = field(default_factory=ConfigSimulator()) def generate_lineshape(self, tensor_type: str = "shielding") -> np.ndarray: """Pre-compute a lineshape kernel to use for the least-squares fitting of an From 3c31a0a814afa74b32f6013b5c24634e9c95f653 Mon Sep 17 00:00:00 2001 From: deepanshs <21365911+deepanshs@users.noreply.github.com> Date: Mon, 25 Mar 2024 00:11:13 -0400 Subject: [PATCH 16/21] fix example html render --- fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py index ae4587237..15221922b 100644 --- a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py +++ b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py @@ -9,7 +9,7 @@ # :math:`^{27}\text{Al}` MAS spectrum of a phospho-aluminosilicate glass. Setting up # the least-squares minimization for a distribution of spin systems is slightly # different than that of a crystalline solid. - +# # There are 4 steps involved in the processs: # - Importing the experimental dataset, # - Generating a pre-computed line shape kernel of subspectra, From 9bd20ce40b5ef7248c6b8c21c0220262b1aed9da Mon Sep 17 00:00:00 2001 From: deepanshs <21365911+deepanshs@users.noreply.github.com> Date: Fri, 29 Mar 2024 06:06:36 -0400 Subject: [PATCH 17/21] speed up fitting example --- .../1D_fitting/plot_6_139La_Ext_Czjzek.py | 41 ++++++++++++------- .../1D_fitting/plot_7_27Al_Extended_Czjzek.py | 2 +- 2 files changed, 27 insertions(+), 16 deletions(-) diff --git a/fitting_source/1D_fitting/plot_6_139La_Ext_Czjzek.py b/fitting_source/1D_fitting/plot_6_139La_Ext_Czjzek.py index 0ba9d73f4..5e37bb278 100644 --- a/fitting_source/1D_fitting/plot_6_139La_Ext_Czjzek.py +++ b/fitting_source/1D_fitting/plot_6_139La_Ext_Czjzek.py @@ -33,7 +33,7 @@ experiment.x[0].to("ppm", "nmr_frequency_ratio") experiment /= experiment.max() -plt.figure(figsize=(6, 4)) +plt.figure(figsize=(4, 3)) ax = plt.subplot(projection="csdm") ax.plot(experiment, "k", alpha=0.5) plt.grid() @@ -95,6 +95,7 @@ def make_kernel(pos, method, isotope): for cq, e in zip(Cq.ravel(), eta.ravel()) ] sim = Simulator(spin_systems=spin_systems, methods=[method]) + sim.config.number_of_sidebands = 4 sim.config.decompose_spectrum = "spin_system" sim.run(pack_as_csdm=False) # Will return spectrum as numpy array, not CSDM object @@ -148,13 +149,13 @@ def make_spectrum_from_parameters(params, kernel, processor, pos, distribution): iso_shift = values["dist_iso_shift"] # Setup model object and get the amplitude - distribution.symmetric_tensor = {"Cq": Cq, "eta": eta} + distribution.symmetric_tensor.Cq = Cq + distribution.symmetric_tensor.eta = eta distribution.eps = eps - # model_quad = ExtCzjzekDistribution({"Cq": Cq, "eta": eta}, eps) _, _, amp = distribution.pdf(pos=pos) # Create spectra by dotting the amplitude distribution with the kernel - dist = kernel.dot(amp.ravel()) + dist = np.dot(kernel, amp.ravel()) # Pack numpy array as csdm object and apply signal processing guess_dataset = cp.CSDM( @@ -213,13 +214,14 @@ def residuals(exp_spectra, simulated_spectra): ) residual_spectrum = residuals(experiment, initial_guess) -plt.figure(figsize=(6, 4)) +plt.figure(figsize=(4, 3)) ax = plt.subplot(projection="csdm") -ax.plot(experiment.real, "k", alpha=0.5) -ax.plot(initial_guess.real, "r", alpha=0.75, label="Guess") -ax.plot(residual_spectrum.real, "b", alpha=0.75, label="Residuals") +ax.plot(experiment.real, "k", alpha=0.5, label="Experiment") +ax.plot(initial_guess.real, "r", alpha=0.3, label="Guess") +ax.plot(residual_spectrum.real, "b", alpha=0.3, label="Residuals") plt.legend() plt.grid() +plt.title("Initial Guess") plt.tight_layout() plt.show() @@ -242,17 +244,25 @@ def minimization_function( # %% -# Create the Minimization object with all the functions and variables previously -# created. +scipy_minimization_kwargs = dict( + diff_step=1e-3, # Increase step size + gtol=1e-15, # Increase global convergence requirement (default 1e-8) + xtol=1e-15, # Increase variable convergence requirement (default 1e-8) + verbose=2, # Print minimization info during each step + loss="soft_l1", +) + minner = Minimizer( minimization_function, params, fcn_args=(experiment, processor, kernel, pos, distribution), + **scipy_minimization_kwargs, ) -result = minner.minimize(method="powell") +result = minner.minimize(method="least_squares") best_fit_params = result.params # Grab the Parameters object from the best fit result + # %% # Plot the best-fit solution # -------------------------- @@ -263,14 +273,15 @@ def minimization_function( ) residual_spectrum = residuals(experiment, final_fit) -plt.figure(figsize=(6, 4)) +plt.figure(figsize=(4, 3)) ax = plt.subplot(projection="csdm") -ax.plot(experiment, "k", alpha=0.5) -ax.plot(final_fit, "r", alpha=0.75, label="Fit") -ax.plot(residual_spectrum, "b", alpha=0.75, label="Residuals") +ax.plot(experiment, "k", alpha=0.5, label="Experiment") +ax.plot(final_fit, "r", alpha=0.3, label="Fit") +ax.plot(residual_spectrum, "b", alpha=0.3, label="Residuals") plt.legend() ax.set_xlim(-11000, 9000) plt.grid() +plt.title("Best Fit") plt.tight_layout() plt.show() diff --git a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py index 15221922b..39ed498a9 100644 --- a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py +++ b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py @@ -186,7 +186,7 @@ plt.figure(figsize=(4, 3)) ax = plt.subplot(projection="csdm") ax.plot(exp_data, "k", alpha=0.5, label="Experiment") -ax.plot(bestfit, "r", alpha=0.3, label="Guess") +ax.plot(bestfit, "r", alpha=0.3, label="Fit") ax.plot(residuals, "b", alpha=0.3, label="Residuals") plt.legend() plt.grid() From f1df00dda894c19b9af84e1c982f6c933175a62b Mon Sep 17 00:00:00 2001 From: Deepansh Srivastava Date: Fri, 29 Mar 2024 11:37:55 -0400 Subject: [PATCH 18/21] add csdm option for distributions --- .../spin_system_distributions/czjzek.rst | 3 +- .../plot_5_czjzek_distribution.py | 6 ++- .../plot_6_extended_czjzek.py | 8 ++-- .../1D_fitting/plot_6_139La_Ext_Czjzek.py | 38 ++++++++++++---- .../1D_fitting/plot_7_27Al_Extended_Czjzek.py | 9 ++-- .../models/analytical_distributions.py | 2 +- src/mrsimulator/models/czjzek.py | 45 +++++++++++++------ .../tests/test_czjzek_and_extended_czjzek.py | 6 +-- 8 files changed, 81 insertions(+), 36 deletions(-) diff --git a/docs/user_guide/spin_system_distributions/czjzek.rst b/docs/user_guide/spin_system_distributions/czjzek.rst index 96622f1d9..53aaf272d 100644 --- a/docs/user_guide/spin_system_distributions/czjzek.rst +++ b/docs/user_guide/spin_system_distributions/czjzek.rst @@ -233,7 +233,8 @@ tensor parameters follow the Czjzek distribution. eta_range = np.linspace(0, 1, num=50) # Create (Cq, eta) grid points and amplitude - Cq_grid, eta_grid, amp = cz_model.pdf(pos=[Cq_range, eta_range]) + 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", diff --git a/examples_source/1D_simulation(macro_amorphous)/plot_5_czjzek_distribution.py b/examples_source/1D_simulation(macro_amorphous)/plot_5_czjzek_distribution.py index 16588cfc3..6d17e9914 100644 --- a/examples_source/1D_simulation(macro_amorphous)/plot_5_czjzek_distribution.py +++ b/examples_source/1D_simulation(macro_amorphous)/plot_5_czjzek_distribution.py @@ -33,7 +33,8 @@ # The range of zeta and eta coordinates over which the distribution is sampled. z_range = np.arange(100) - 50 # in ppm e_range = np.arange(21) / 20 -z_dist, e_dist, amp = CzjzekDistribution(sigma=3.1415).pdf(pos=[z_range, e_range]) +z_dist, e_dist = np.meshgrid(z_range, e_range) +_, _, amp = CzjzekDistribution(sigma=3.1415).pdf(pos=[z_range, e_range]) # %% # Here ``z_range`` and ``e_range`` are the coordinates along the :math:`\zeta` and @@ -94,7 +95,8 @@ # The range of Cq and eta coordinates over which the distribution is sampled. cq_range = np.arange(100) * 0.6 - 30 # in MHz e_range = np.arange(21) / 20 -cq_dist, e_dist, amp = CzjzekDistribution(sigma=2.3).pdf(pos=[cq_range, e_range]) +cq_dist, e_dist = np.meshgrid(cq_range, e_range) +_, _, amp = CzjzekDistribution(sigma=2.3).pdf(pos=[cq_range, e_range]) # The following is the contour plot of the Czjzek distribution. plt.figure(figsize=(4.25, 3.0)) diff --git a/examples_source/1D_simulation(macro_amorphous)/plot_6_extended_czjzek.py b/examples_source/1D_simulation(macro_amorphous)/plot_6_extended_czjzek.py index e31693e33..e05c6d457 100644 --- a/examples_source/1D_simulation(macro_amorphous)/plot_6_extended_czjzek.py +++ b/examples_source/1D_simulation(macro_amorphous)/plot_6_extended_czjzek.py @@ -36,7 +36,8 @@ e_lim = np.arange(21) / 20 dominant = {"zeta": 60, "eta": 0.3} -z_dist, e_dist, amp = ExtCzjzekDistribution(dominant, eps=0.14).pdf(pos=[z_lim, e_lim]) +z_dist, e_dist = np.meshgrid(z_lim, e_lim) +_, _, amp = ExtCzjzekDistribution(dominant, eps=0.14).pdf(pos=[z_lim, e_lim]) # %% # The following is the plot of the extended Czjzek distribution. @@ -93,9 +94,8 @@ e_lim = np.arange(21) / 20 dominant = {"Cq": 6.1, "eta": 0.1} -cq_dist, e_dist, amp = ExtCzjzekDistribution(dominant, eps=0.25).pdf( - pos=[cq_lim, e_lim] -) +cq_dist, e_dist = np.meshgrid(cq_lim, e_lim) +_, _, amp = ExtCzjzekDistribution(dominant, eps=0.25).pdf(pos=[cq_lim, e_lim]) # %% # The following is the plot of the extended Czjzek distribution. diff --git a/fitting_source/1D_fitting/plot_6_139La_Ext_Czjzek.py b/fitting_source/1D_fitting/plot_6_139La_Ext_Czjzek.py index 5e37bb278..94e517f44 100644 --- a/fitting_source/1D_fitting/plot_6_139La_Ext_Czjzek.py +++ b/fitting_source/1D_fitting/plot_6_139La_Ext_Czjzek.py @@ -22,7 +22,7 @@ import matplotlib.pyplot as plt -# sphinx_gallery_thumbnail_number = 3 +# sphinx_gallery_thumbnail_number = 4 # %% # Import the dataset @@ -91,8 +91,10 @@ def make_kernel(pos, method, isotope): Cq, eta = np.meshgrid(pos[0], pos[1], indexing="xy") spin_systems = [ - SpinSystem(sites=[Site(isotope=isotope, quadrupolar=dict(Cq=cq * 1e6, eta=e))]) - for cq, e in zip(Cq.ravel(), eta.ravel()) + SpinSystem( + sites=[Site(isotope=isotope, quadrupolar=dict(Cq=cq_ * 1e6, eta=e_))] + ) + for cq_, e_ in zip(Cq.ravel(), eta.ravel()) ] sim = Simulator(spin_systems=spin_systems, methods=[method]) sim.config.number_of_sidebands = 4 @@ -104,7 +106,7 @@ def make_kernel(pos, method, isotope): # Create ranges to construct cq and eta grid points -cq_range = np.linspace(0, 100, num=100) * 0.8 + 25 +cq_range = np.linspace(0, 50, num=50) * 1.6 + 25 eta_range = np.arange(21) / 20 pos = [cq_range, eta_range] kernel = make_kernel(pos, method, "139La") @@ -200,7 +202,7 @@ def residuals(exp_spectra, simulated_spectra): params = lmfit.Parameters() params.add("dist_Cq", value=49.5) params.add("dist_eta", value=0.55, min=0, max=1) -params.add("dist_eps", value=0.24, min=0) +params.add("dist_eps", value=0.1, min=0) params.add("dist_iso_shift", value=350) params.add("sp_scale_factor", value=3.8e3, min=0) @@ -209,15 +211,17 @@ def residuals(exp_spectra, simulated_spectra): symmetric_tensor={"Cq": params["dist_Cq"], "eta": params["dist_eta"]}, eps=params["dist_eps"], ) -initial_guess = make_spectrum_from_parameters( +guess_distribution = distribution.pdf(pos, size=100_000, pack_as_csdm=True) + +initial_guess_spectrum = make_spectrum_from_parameters( params, kernel, processor, pos, distribution ) -residual_spectrum = residuals(experiment, initial_guess) +residual_spectrum = residuals(experiment, initial_guess_spectrum) plt.figure(figsize=(4, 3)) ax = plt.subplot(projection="csdm") ax.plot(experiment.real, "k", alpha=0.5, label="Experiment") -ax.plot(initial_guess.real, "r", alpha=0.3, label="Guess") +ax.plot(initial_guess_spectrum.real, "r", alpha=0.3, label="Guess") ax.plot(residual_spectrum.real, "b", alpha=0.3, label="Residuals") plt.legend() plt.grid() @@ -226,6 +230,15 @@ def residuals(exp_spectra, simulated_spectra): plt.show() +# %% +# Guess distribution +plt.figure(figsize=(4, 3)) +ax = plt.subplot(projection="csdm") +ax.imshow(guess_distribution, interpolation="none", cmap="gist_ncar_r", aspect="auto") +plt.tight_layout() +plt.show() + + # %% # Least-squares minimization with LMFIT # ------------------------------------- @@ -285,6 +298,15 @@ def minimization_function( plt.tight_layout() plt.show() + +# %% +# Best fit probability distribution +prob = distribution.pdf(pos, pack_as_csdm=True) +plt.figure(figsize=(4, 3)) +ax = plt.subplot(projection="csdm") +ax.imshow(prob, interpolation="none", cmap="gist_ncar_r", aspect="auto") +plt.tight_layout() +plt.show() # %% # # .. [#f1] A. J. Fernández-Carrión, M. Allix, P. Florian, M. R. Suchomel, and A. I. diff --git a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py index 39ed498a9..1d5f515c0 100644 --- a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py +++ b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py @@ -201,11 +201,12 @@ for i, model in enumerate(all_models): model.update_lmfit_params(result.params, i) -xx, yy, amp = cz_model.pdf(pos=pos) +amp = cz_model.pdf(pos=pos, pack_as_csdm=True) plt.figure(figsize=(4, 3)) -plt.contourf(xx / 1e6, yy / 1e6, amp) -plt.xlabel("x / MHz ") -plt.ylabel("y / MHz") +ax = plt.subplot(projection="csdm") +ax.imshow(amp, cmap="gist_ncar_r", interpolation="none", aspect="auto") +ax.set_xlabel("x / Hz ") +ax.set_ylabel("y / Hz") plt.tight_layout() plt.show() diff --git a/src/mrsimulator/models/analytical_distributions.py b/src/mrsimulator/models/analytical_distributions.py index fb8ecc200..b1084064e 100644 --- a/src/mrsimulator/models/analytical_distributions.py +++ b/src/mrsimulator/models/analytical_distributions.py @@ -52,7 +52,7 @@ def czjzek_zeta_eta(sigma: float, pos: list): pdf_model = pdf_model.ravel() pdf_model[eta_idx] /= 2.0 pdf_model.shape = zeta.shape - return zeta, eta, pdf_model + return pos[0], pos[1], pdf_model def czjzek_polar(sigma: float, pos: list): diff --git a/src/mrsimulator/models/czjzek.py b/src/mrsimulator/models/czjzek.py index fc68215cd..3313e8e08 100644 --- a/src/mrsimulator/models/czjzek.py +++ b/src/mrsimulator/models/czjzek.py @@ -4,6 +4,7 @@ __author__ = "Deepansh J. Srivastava" __email__ = "dsrivastava@hyperfine.io" """ +import csdmpy as cp import mrsimulator.models.analytical_distributions as analytical_dist import numpy as np from mrsimulator.spin_system.tensors import SymmetricTensor @@ -12,7 +13,6 @@ from .utils import get_principal_components from .utils import zeta_eta_to_x_y - __author__ = "Deepansh J. Srivastava" __email__ = "srivastava.89@osu.edu" @@ -60,15 +60,15 @@ def _czjzek_random_distribution_tensors(sigma, n): sqrt_3_U5 = np.random.normal(0.0, sqrt_3_sigma, n) # Create N random tensors - tensors = np.zeros((n, 3, 3)) # n x 3 x 3 tensors + tensors = np.empty((n, 3, 3)) # n x 3 x 3 tensors tensors[:, 0, 0] = sqrt_3_U5 - U1 # xx - tensors[:, 0, 1] = sqrt_3_U4 # xy - tensors[:, 0, 2] = sqrt_3_U2 # xz + # tensors[:, 0, 1] = sqrt_3_U4 # xy + # tensors[:, 0, 2] = sqrt_3_U2 # xz tensors[:, 1, 0] = sqrt_3_U4 # yx tensors[:, 1, 1] = -sqrt_3_U5 - U1 # yy - tensors[:, 1, 2] = sqrt_3_U3 # yz + # tensors[:, 1, 2] = sqrt_3_U3 # yz tensors[:, 2, 0] = sqrt_3_U2 # zx tensors[:, 2, 1] = sqrt_3_U3 # zy @@ -92,7 +92,13 @@ def __init__( self.abundance = abundance self.polar = polar - def pdf(self, pos, size: int = 400000, analytical: bool = True): + def pdf( + self, + pos, + size: int = 400000, + analytical: bool = True, + pack_as_csdm: bool = False, + ): """Generates a probability distribution function by binning the random variates of length size onto the given grid system. @@ -100,21 +106,30 @@ def pdf(self, pos, size: int = 400000, analytical: bool = True): pos: A list of coordinates along the two dimensions given as NumPy arrays. size: The number of random variates drawn in generating the pdf. The default is 400000. + pack_as_csdm: If true, returns as csdm object. Returns: - A list of x and y coordinates and the corresponding amplitudes. + A list of x and y coordinates and the corresponding amplitudes if not packed + as csdm object, else a csdm object. Example: >>> import numpy as np >>> cq = np.arange(50) - 25 >>> eta = np.arange(21)/20 - >>> Cq_dist, eta_dist, amp = cz_model.pdf(pos=[cq, eta]) + >>> amp = cz_model.pdf(pos=[cq, eta]) # returns amp as a CSDM object. """ if analytical and self.model_name in ANALYTICAL_AVAILABLE: analytical_model = ANALYTICAL_AVAILABLE[self.model_name] - return analytical_model(self.sigma, pos, self.polar) + pos_a, pos_b, data_ = analytical_model(self.sigma, pos, self.polar) else: - return self.pdf_numerical(pos, size) + pos_a, pos_b, data_ = self.pdf_numerical(pos, size) + + if pack_as_csdm: + if len(pos_a.shape) == 1: + return self.pack_csdm_object([pos_a, pos_b], data_) + else: + return self.pack_csdm_object([pos_a[0, :], pos_b[:, 0]], data_) + return pos_a, pos_b, data_ def pdf_numerical(self, pos, size: int = 400000): """Generate distribution numerically""" @@ -129,9 +144,13 @@ def pdf_numerical(self, pos, size: int = 400000): hist, _, _ = np.histogram2d(zeta, eta, bins=[x_size, y_size], range=[x, y]) hist /= hist.sum() + return pos[0], pos[1], hist.T - x_, y_ = np.meshgrid(pos[0], pos[1]) - return x_, y_, hist.T + def pack_csdm_object(self, pos, data): + """Pack data and coordinates as csdm objects""" + dims = [cp.as_dimension(item) for item in pos] + dvs = cp.as_dependent_variable(data) + return cp.CSDM(dimensions=dims, dependent_variables=[dvs]) class CzjzekDistribution(AbstractDistribution): @@ -330,7 +349,7 @@ def rvs(self, size: int): norm_T0 = np.linalg.norm(T0) # the perturbation factor - rho = self.eps * norm_T0 / np.sqrt(30) + rho = self.eps * norm_T0 / 5.4772255751 # np.sqrt(30) = 5.477225575 # total tensor total_tensors = np.diag(T0) + rho * tensors diff --git a/src/mrsimulator/models/tests/test_czjzek_and_extended_czjzek.py b/src/mrsimulator/models/tests/test_czjzek_and_extended_czjzek.py index 5d58b7900..31e15e9b5 100644 --- a/src/mrsimulator/models/tests/test_czjzek_and_extended_czjzek.py +++ b/src/mrsimulator/models/tests/test_czjzek_and_extended_czjzek.py @@ -78,10 +78,10 @@ def test_czjzek_distribution(): z_range = (ran_z[1:] + ran_z[:-1]) / 2 # czjzek distribution from analytical formula - _, _, res = CzjzekDistribution(sigma).pdf(pos=[z_range, e_range]) + res = CzjzekDistribution(sigma).pdf(pos=[z_range, e_range], pack_as_csdm=True) - eta_pro = res.sum(axis=1) - zeta_pro = res.sum(axis=0) + eta_pro = res.sum(axis=0).y[0].components[0] + zeta_pro = res.sum(axis=1).y[0].components[0] # eta test message = "failed to compare eta projection for Czjzek distribution" From 17ef011c3d277d9419866711c6920aa713daad36 Mon Sep 17 00:00:00 2001 From: deepanshs <21365911+deepanshs@users.noreply.github.com> Date: Sat, 30 Mar 2024 08:57:54 -0400 Subject: [PATCH 19/21] update fit params --- .../1D_fitting/plot_6_139La_Ext_Czjzek.py | 27 +++++++++++-------- .../1D_fitting/plot_7_27Al_Extended_Czjzek.py | 14 +++++----- src/mrsimulator/utils/spectral_fitting.py | 2 +- 3 files changed, 25 insertions(+), 18 deletions(-) diff --git a/fitting_source/1D_fitting/plot_6_139La_Ext_Czjzek.py b/fitting_source/1D_fitting/plot_6_139La_Ext_Czjzek.py index 94e517f44..8f40c8bfb 100644 --- a/fitting_source/1D_fitting/plot_6_139La_Ext_Czjzek.py +++ b/fitting_source/1D_fitting/plot_6_139La_Ext_Czjzek.py @@ -91,9 +91,7 @@ def make_kernel(pos, method, isotope): Cq, eta = np.meshgrid(pos[0], pos[1], indexing="xy") spin_systems = [ - SpinSystem( - sites=[Site(isotope=isotope, quadrupolar=dict(Cq=cq_ * 1e6, eta=e_))] - ) + SpinSystem(sites=[Site(isotope=isotope, quadrupolar=dict(Cq=cq_, eta=e_))]) for cq_, e_ in zip(Cq.ravel(), eta.ravel()) ] sim = Simulator(spin_systems=spin_systems, methods=[method]) @@ -106,7 +104,7 @@ def make_kernel(pos, method, isotope): # Create ranges to construct cq and eta grid points -cq_range = np.linspace(0, 50, num=50) * 1.6 + 25 +cq_range = (np.linspace(0, 100, num=100) * 0.8 + 25) * 1e6 # in Hz eta_range = np.arange(21) / 20 pos = [cq_range, eta_range] kernel = make_kernel(pos, method, "139La") @@ -200,10 +198,10 @@ def residuals(exp_spectra, simulated_spectra): # If you adapt this example to your own dataset, make sure the initial guess is # decently good, otherwise LMFIT is likely to fall into a local minima. params = lmfit.Parameters() -params.add("dist_Cq", value=49.5) +params.add("dist_Cq", value=49.5e6) # Hz params.add("dist_eta", value=0.55, min=0, max=1) params.add("dist_eps", value=0.1, min=0) -params.add("dist_iso_shift", value=350) +params.add("dist_iso_shift", value=350) # ppm params.add("sp_scale_factor", value=3.8e3, min=0) # Plot the initial guess spectrum along with the experimental data @@ -211,7 +209,7 @@ def residuals(exp_spectra, simulated_spectra): symmetric_tensor={"Cq": params["dist_Cq"], "eta": params["dist_eta"]}, eps=params["dist_eps"], ) -guess_distribution = distribution.pdf(pos, size=100_000, pack_as_csdm=True) +guess_distribution = distribution.pdf(pos, size=400_000, pack_as_csdm=True) initial_guess_spectrum = make_spectrum_from_parameters( params, kernel, processor, pos, distribution @@ -235,6 +233,8 @@ def residuals(exp_spectra, simulated_spectra): plt.figure(figsize=(4, 3)) ax = plt.subplot(projection="csdm") ax.imshow(guess_distribution, interpolation="none", cmap="gist_ncar_r", aspect="auto") +ax.set_xlabel("Cq / Hz") +ax.set_ylabel(r"$\eta$") plt.tight_layout() plt.show() @@ -257,12 +257,15 @@ def minimization_function( # %% +# Sinice the probabilty distribution is generated from a sparsely sampled +# from a 5D second rank tensor parameter space, we increase the `diff_step` size from +# machine precession to avoid approaching local minima from noise. scipy_minimization_kwargs = dict( - diff_step=1e-3, # Increase step size - gtol=1e-15, # Increase global convergence requirement (default 1e-8) - xtol=1e-15, # Increase variable convergence requirement (default 1e-8) + diff_step=1e-4, # Increase step size from machine precession + gtol=1e-10, # Decrease global convergence requirement (default 1e-8) + xtol=1e-10, # Decrease variable convergence requirement (default 1e-8) verbose=2, # Print minimization info during each step - loss="soft_l1", + loss="linear", ) minner = Minimizer( @@ -305,6 +308,8 @@ def minimization_function( plt.figure(figsize=(4, 3)) ax = plt.subplot(projection="csdm") ax.imshow(prob, interpolation="none", cmap="gist_ncar_r", aspect="auto") +ax.set_xlabel("Cq / Hz") +ax.set_ylabel(r"$\eta$") plt.tight_layout() plt.show() # %% diff --git a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py index 1d5f515c0..5aa22de4f 100644 --- a/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py +++ b/fitting_source/1D_fitting/plot_7_27Al_Extended_Czjzek.py @@ -159,13 +159,15 @@ # ----------------------------- # Finally, a `Minimizer` object is created and a minimization run using least-squares. # The same arguments defined in the `addtl_sf_kwargs` variable are also passed to the -# minimizer. +# minimizer. Sinice the probabilty distribution is generated from a sparsely sampled +# from a 5D second rank tensor parameter space, we increase the `diff_step` size from +# machine precession to avoid approaching local minima from noise. scipy_minimization_kwargs = dict( - diff_step=1e-3, # Increase step size - gtol=1e-15, # Increase global convergence requirement (default 1e-8) - xtol=1e-15, # Increase variable convergence requirement (default 1e-8) + diff_step=1e-4, # Increase step size from machine precesion. + gtol=1e-10, # Decrease global convergence requirement (default 1e-8) + xtol=1e-10, # Decrease variable convergence requirement (default 1e-8) verbose=2, # Print minimization info during each step - loss="soft_l1", + loss="linear", ) minner = Minimizer( @@ -206,7 +208,7 @@ plt.figure(figsize=(4, 3)) ax = plt.subplot(projection="csdm") ax.imshow(amp, cmap="gist_ncar_r", interpolation="none", aspect="auto") -ax.set_xlabel("x / Hz ") +ax.set_xlabel("x / Hz") ax.set_ylabel("y / Hz") plt.tight_layout() plt.show() diff --git a/src/mrsimulator/utils/spectral_fitting.py b/src/mrsimulator/utils/spectral_fitting.py index 79cc16679..fa768e70b 100644 --- a/src/mrsimulator/utils/spectral_fitting.py +++ b/src/mrsimulator/utils/spectral_fitting.py @@ -552,7 +552,7 @@ def residuals(sim: Simulator, processors: list = None): def _apply_iso_shift(csdm_obj, iso_shift_ppm, larmor_freq_Hz): """Apply isotropic chemical shift to a CSDM object using the FFT shift theorem.""" csdm_obj = csdm_obj.fft() - time_coords = csdm_obj.x[0].coordinates.value + time_coords = csdm_obj.x[0].coordinates.to("s").value iso_shift_Hz = larmor_freq_Hz * iso_shift_ppm csdm_obj.y[0].components[0] *= np.exp(-np.pi * 2j * time_coords * iso_shift_Hz) csdm_obj = csdm_obj.fft() From fabc657af9023582a53c78482f7d68a2dca93a5f Mon Sep 17 00:00:00 2001 From: deepanshs <21365911+deepanshs@users.noreply.github.com> Date: Sat, 30 Mar 2024 13:33:06 -0400 Subject: [PATCH 20/21] add tests --- src/mrsimulator/models/czjzek.py | 66 ++++++++++--------- .../models/tests/test_distribution_params.py | 57 ++++++++++++++++ src/mrsimulator/utils/spectral_fitting.py | 21 ++---- 3 files changed, 99 insertions(+), 45 deletions(-) create mode 100644 src/mrsimulator/models/tests/test_distribution_params.py diff --git a/src/mrsimulator/models/czjzek.py b/src/mrsimulator/models/czjzek.py index 3313e8e08..de1f4e24a 100644 --- a/src/mrsimulator/models/czjzek.py +++ b/src/mrsimulator/models/czjzek.py @@ -78,6 +78,10 @@ def _czjzek_random_distribution_tensors(sigma, n): class AbstractDistribution: + """Abstract distribution""" + + model_name = "base" + def __init__( self, mean_isotropic_chemical_shift=0.0, @@ -152,6 +156,18 @@ def pack_csdm_object(self, pos, data): dvs = cp.as_dependent_variable(data) return cp.CSDM(dimensions=dims, dependent_variables=[dvs]) + def add_lmfit_params(self, params, i): + """Add lmfit params for base class""" + prefix = self.model_name + params.add(f"{prefix}_{i}_iso_shift", value=self.mean_isotropic_chemical_shift) + params.add(f"{prefix}_{i}_abundance", value=self.abundance, min=0, max=1) + + def update_lmfit_params(self, params, i): + """Update lmfit params for base class""" + prefix = self.model_name + self.mean_isotropic_chemical_shift = params[f"{prefix}_{i}_iso_shift"].value + self.abundance = params[f"{prefix}_{i}_abundance"].value + class CzjzekDistribution(AbstractDistribution): r"""A Czjzek distribution model class. @@ -230,23 +246,18 @@ def rvs(self, size: int): return get_Haeberlen_components(tensors) return zeta_eta_to_x_y(*get_Haeberlen_components(tensors)) - def param_prefix(self): - return "czjzek" - - def get_lmfit_params(self, params, i): + def add_lmfit_params(self, params, i): """Create lmfit params for index i""" - params.add(f"czjzek_{i}_sigma", value=self.sigma, min=0) - params.add(f"czjzek_{i}_iso_shift", value=self.mean_isotropic_chemical_shift) - params.add(f"czjzek_{i}_weight", value=self.abundance, min=0, max=1) + prefix = self.model_name + params.add(f"{prefix}_{i}_sigma", value=self.sigma, min=0) + super().add_lmfit_params(params, i) return params def update_lmfit_params(self, params, i): - """Create lmfit params for index i""" - prefix = self.param_prefix() - + """Update lmfit params for index i""" + prefix = self.model_name self.sigma = params[f"{prefix}_{i}_sigma"].value - self.mean_isotropic_chemical_shift = params[f"{prefix}_{i}_iso_shift"].value - self.abundance = params[f"{prefix}_{i}_weight"].value + super().update_lmfit_params(params, i) class ExtCzjzekDistribution(AbstractDistribution): @@ -289,7 +300,7 @@ class ExtCzjzekDistribution(AbstractDistribution): >>> S0 = {"Cq": 1e6, "eta": 0.3} >>> ext_cz_model = ExtCzjzekDistribution(S0, eps=0.35) """ - model_name = "extended czjzek" + model_name = "extended_czjzek" def __init__( self, @@ -348,8 +359,8 @@ def rvs(self, size: int): # 2-norm of the tensor norm_T0 = np.linalg.norm(T0) - # the perturbation factor - rho = self.eps * norm_T0 / 5.4772255751 # np.sqrt(30) = 5.477225575 + # the perturbation factor # np.sqrt(30) = 5.477225575 + rho = self.eps * norm_T0 / 5.4772255751 # total tensor total_tensors = np.diag(T0) + rho * tensors @@ -358,34 +369,29 @@ def rvs(self, size: int): return get_Haeberlen_components(total_tensors) return zeta_eta_to_x_y(*get_Haeberlen_components(total_tensors)) - def param_prefix(self): - return "ext_czjzek" - - def get_lmfit_params(self, params, i): + def add_lmfit_params(self, params, i): """Create lmfit params for index i""" - prefix = self.param_prefix() + prefix = self.model_name if self.symmetric_tensor.zeta is not None: zeta = self.symmetric_tensor.zeta else: zeta = self.symmetric_tensor.Cq - params.add(f"{prefix}_{i}_zeta0", value=zeta) - params.add(f"{prefix}_{i}_eta0", value=self.symmetric_tensor.eta, min=0, max=1) - params.add(f"{prefix}_{i}_epsilon", value=self.eps, min=0) - params.add(f"{prefix}_{i}_iso_shift", value=self.mean_isotropic_chemical_shift) - params.add(f"{prefix}_{i}_weight", value=self.abundance, min=0, max=1) + params.add(f"{prefix}_{i}_zeta", value=zeta) + params.add(f"{prefix}_{i}_eta", value=self.symmetric_tensor.eta, min=0, max=1) + params.add(f"{prefix}_{i}_epsilon", value=self.eps, min=1e-6) + super().add_lmfit_params(params, i) return params def update_lmfit_params(self, params, i): """Create lmfit params for index i""" - prefix = self.param_prefix() + prefix = self.model_name - zeta = params[f"{prefix}_{i}_zeta0"].value + zeta = params[f"{prefix}_{i}_zeta"].value if self.symmetric_tensor.zeta is not None: self.symmetric_tensor.zeta = zeta else: self.symmetric_tensor.Cq = zeta - self.symmetric_tensor.eta = params[f"{prefix}_{i}_eta0"].value + self.symmetric_tensor.eta = params[f"{prefix}_{i}_eta"].value self.eps = params[f"{prefix}_{i}_epsilon"].value - self.mean_isotropic_chemical_shift = params[f"{prefix}_{i}_iso_shift"].value - self.abundance = params[f"{prefix}_{i}_weight"].value + super().update_lmfit_params(params, i) diff --git a/src/mrsimulator/models/tests/test_distribution_params.py b/src/mrsimulator/models/tests/test_distribution_params.py new file mode 100644 index 000000000..67c904d8b --- /dev/null +++ b/src/mrsimulator/models/tests/test_distribution_params.py @@ -0,0 +1,57 @@ +from lmfit import Parameters +from mrsimulator.models import CzjzekDistribution +from mrsimulator.models import ExtCzjzekDistribution + + +def test_dist_1(): + dist = CzjzekDistribution(sigma=1.5, mean_isotropic_chemical_shift=23.4) + prefix = "czjzek" + + params = Parameters() + params = dist.add_lmfit_params(params=params, i=0) + assert len(params.keys()) == 3 + + assert params[f"{prefix}_0_sigma"] == 1.5 + assert params[f"{prefix}_0_iso_shift"] == 23.4 + assert params[f"{prefix}_0_abundance"] == 1.0 + + params[f"{prefix}_0_sigma"].value = 10 + params[f"{prefix}_0_iso_shift"].value = -63.4 + params[f"{prefix}_0_abundance"].value = 0.15 + + dist.update_lmfit_params(params=params, i=0) + assert dist.sigma == 10.0 + assert dist.mean_isotropic_chemical_shift == -63.4 + assert dist.abundance == 0.15 + + +def test_dist_2(): + dist = ExtCzjzekDistribution( + symmetric_tensor={"zeta": 4.1, "eta": 0.2}, + eps=0.3, + mean_isotropic_chemical_shift=0.0, + ) + prefix = "extended_czjzek" + + params = Parameters() + params = dist.add_lmfit_params(params=params, i=0) + assert len(params.keys()) == 5 + + assert params[f"{prefix}_0_zeta"] == 4.1 + assert params[f"{prefix}_0_eta"] == 0.2 + assert params[f"{prefix}_0_epsilon"] == 0.3 + assert params[f"{prefix}_0_iso_shift"] == 0.0 + assert params[f"{prefix}_0_abundance"] == 1.0 + + params[f"{prefix}_0_zeta"].value = -30.2 + params[f"{prefix}_0_eta"].value = 0.65 + params[f"{prefix}_0_epsilon"].value = 0.6 + params[f"{prefix}_0_iso_shift"].value = -3.4 + params[f"{prefix}_0_abundance"].value = 0.9 + + dist.update_lmfit_params(params=params, i=0) + assert dist.symmetric_tensor.zeta == -30.2 + assert dist.symmetric_tensor.eta == 0.65 + assert dist.eps == 0.6 + assert dist.mean_isotropic_chemical_shift == -3.4 + assert dist.abundance == 0.9 diff --git a/src/mrsimulator/utils/spectral_fitting.py b/src/mrsimulator/utils/spectral_fitting.py index fa768e70b..c17d667af 100644 --- a/src/mrsimulator/utils/spectral_fitting.py +++ b/src/mrsimulator/utils/spectral_fitting.py @@ -568,31 +568,22 @@ def make_distribution_params( """ n_dist = len(spin_system_models) - # norm = 0.0 - # for model in distribution_models: - # norm += model.abundance - expr_terms = [] params = Parameters() for i, model in enumerate(spin_system_models): # normalize the abundance model.abundance /= norm - params = model.get_lmfit_params(params, i) + params = model.add_lmfit_params(params, i) - # Set last weight parameter as a non-free variable - model_prefix = model.param_prefix() + # Set last abundance parameter as a non-free variable + model_prefix = model.model_name if i < n_dist - 1 or skip_last: - expr_terms.append(f"{model_prefix}_{i}_weight") + expr_terms.append(f"{model_prefix}_{i}_abundance") else: expr = "-".join(expr_terms) expr = "1" if expr == "" else f"1-{expr}" - params[f"{model_prefix}_{i}_weight"].vary = False - params[f"{model_prefix}_{i}_weight"].expr = expr - - # # Add SignalProcessor parameters, if requested - # if processor is not None: - # _make_params_single_processor(params, processor, 0) - + params[f"{model_prefix}_{i}_abundance"].vary = False + params[f"{model_prefix}_{i}_abundance"].expr = expr return params From 545e4ddaba82d4faad58d9051b377b4ceee4ad54 Mon Sep 17 00:00:00 2001 From: deepanshs <21365911+deepanshs@users.noreply.github.com> Date: Sat, 30 Mar 2024 14:15:11 -0400 Subject: [PATCH 21/21] more tests --- src/mrsimulator/models/czjzek.py | 36 ++++++++++++------- .../models/tests/test_distribution_params.py | 26 +++++++------- src/mrsimulator/utils/spectral_fitting.py | 9 ++--- .../utils/tests/test_spectral_fitting.py | 26 ++++++++++++++ 4 files changed, 67 insertions(+), 30 deletions(-) diff --git a/src/mrsimulator/models/czjzek.py b/src/mrsimulator/models/czjzek.py index de1f4e24a..dfca75b14 100644 --- a/src/mrsimulator/models/czjzek.py +++ b/src/mrsimulator/models/czjzek.py @@ -85,7 +85,7 @@ class AbstractDistribution: def __init__( self, mean_isotropic_chemical_shift=0.0, - abundance=1.0, + abundance=100.0, polar=False, cache_tensors=False, ): @@ -159,13 +159,18 @@ def pack_csdm_object(self, pos, data): def add_lmfit_params(self, params, i): """Add lmfit params for base class""" prefix = self.model_name - params.add(f"{prefix}_{i}_iso_shift", value=self.mean_isotropic_chemical_shift) - params.add(f"{prefix}_{i}_abundance", value=self.abundance, min=0, max=1) + params.add( + f"{prefix}_{i}_mean_isotropic_chemical_shift", + value=self.mean_isotropic_chemical_shift, + ) + params.add(f"{prefix}_{i}_abundance", value=self.abundance, min=0, max=100) def update_lmfit_params(self, params, i): """Update lmfit params for base class""" prefix = self.model_name - self.mean_isotropic_chemical_shift = params[f"{prefix}_{i}_iso_shift"].value + self.mean_isotropic_chemical_shift = params[ + f"{prefix}_{i}_mean_isotropic_chemical_shift" + ].value self.abundance = params[f"{prefix}_{i}_abundance"].value @@ -209,7 +214,7 @@ def __init__( self, sigma: float, mean_isotropic_chemical_shift: float = 0.0, - abundance: float = 1.0, + abundance: float = 100.0, polar=False, cache=True, ): @@ -300,14 +305,14 @@ class ExtCzjzekDistribution(AbstractDistribution): >>> S0 = {"Cq": 1e6, "eta": 0.3} >>> ext_cz_model = ExtCzjzekDistribution(S0, eps=0.35) """ - model_name = "extended_czjzek" + model_name = "ext_czjzek" def __init__( self, symmetric_tensor: SymmetricTensor, eps: float, mean_isotropic_chemical_shift: float = 0.0, - abundance: float = 1.0, + abundance: float = 100.0, polar=False, cache=True, ): @@ -376,9 +381,14 @@ def add_lmfit_params(self, params, i): zeta = self.symmetric_tensor.zeta else: zeta = self.symmetric_tensor.Cq - params.add(f"{prefix}_{i}_zeta", value=zeta) - params.add(f"{prefix}_{i}_eta", value=self.symmetric_tensor.eta, min=0, max=1) - params.add(f"{prefix}_{i}_epsilon", value=self.eps, min=1e-6) + params.add(f"{prefix}_{i}_symmetric_tensor_zeta", value=zeta) + params.add( + f"{prefix}_{i}_symmetric_tensor_eta", + value=self.symmetric_tensor.eta, + min=0, + max=1, + ) + params.add(f"{prefix}_{i}_eps", value=self.eps, min=1e-6) super().add_lmfit_params(params, i) return params @@ -386,12 +396,12 @@ def update_lmfit_params(self, params, i): """Create lmfit params for index i""" prefix = self.model_name - zeta = params[f"{prefix}_{i}_zeta"].value + zeta = params[f"{prefix}_{i}_symmetric_tensor_zeta"].value if self.symmetric_tensor.zeta is not None: self.symmetric_tensor.zeta = zeta else: self.symmetric_tensor.Cq = zeta - self.symmetric_tensor.eta = params[f"{prefix}_{i}_eta"].value - self.eps = params[f"{prefix}_{i}_epsilon"].value + self.symmetric_tensor.eta = params[f"{prefix}_{i}_symmetric_tensor_eta"].value + self.eps = params[f"{prefix}_{i}_eps"].value super().update_lmfit_params(params, i) diff --git a/src/mrsimulator/models/tests/test_distribution_params.py b/src/mrsimulator/models/tests/test_distribution_params.py index 67c904d8b..0358d4c24 100644 --- a/src/mrsimulator/models/tests/test_distribution_params.py +++ b/src/mrsimulator/models/tests/test_distribution_params.py @@ -12,11 +12,11 @@ def test_dist_1(): assert len(params.keys()) == 3 assert params[f"{prefix}_0_sigma"] == 1.5 - assert params[f"{prefix}_0_iso_shift"] == 23.4 - assert params[f"{prefix}_0_abundance"] == 1.0 + assert params[f"{prefix}_0_mean_isotropic_chemical_shift"] == 23.4 + assert params[f"{prefix}_0_abundance"] == 100.0 params[f"{prefix}_0_sigma"].value = 10 - params[f"{prefix}_0_iso_shift"].value = -63.4 + params[f"{prefix}_0_mean_isotropic_chemical_shift"].value = -63.4 params[f"{prefix}_0_abundance"].value = 0.15 dist.update_lmfit_params(params=params, i=0) @@ -31,22 +31,22 @@ def test_dist_2(): eps=0.3, mean_isotropic_chemical_shift=0.0, ) - prefix = "extended_czjzek" + prefix = "ext_czjzek" params = Parameters() params = dist.add_lmfit_params(params=params, i=0) assert len(params.keys()) == 5 - assert params[f"{prefix}_0_zeta"] == 4.1 - assert params[f"{prefix}_0_eta"] == 0.2 - assert params[f"{prefix}_0_epsilon"] == 0.3 - assert params[f"{prefix}_0_iso_shift"] == 0.0 - assert params[f"{prefix}_0_abundance"] == 1.0 + assert params[f"{prefix}_0_symmetric_tensor_zeta"] == 4.1 + assert params[f"{prefix}_0_symmetric_tensor_eta"] == 0.2 + assert params[f"{prefix}_0_eps"] == 0.3 + assert params[f"{prefix}_0_mean_isotropic_chemical_shift"] == 0.0 + assert params[f"{prefix}_0_abundance"] == 100.0 - params[f"{prefix}_0_zeta"].value = -30.2 - params[f"{prefix}_0_eta"].value = 0.65 - params[f"{prefix}_0_epsilon"].value = 0.6 - params[f"{prefix}_0_iso_shift"].value = -3.4 + params[f"{prefix}_0_symmetric_tensor_zeta"].value = -30.2 + params[f"{prefix}_0_symmetric_tensor_eta"].value = 0.65 + params[f"{prefix}_0_eps"].value = 0.6 + params[f"{prefix}_0_mean_isotropic_chemical_shift"].value = -3.4 params[f"{prefix}_0_abundance"].value = 0.9 dist.update_lmfit_params(params=params, i=0) diff --git a/src/mrsimulator/utils/spectral_fitting.py b/src/mrsimulator/utils/spectral_fitting.py index c17d667af..f32568140 100644 --- a/src/mrsimulator/utils/spectral_fitting.py +++ b/src/mrsimulator/utils/spectral_fitting.py @@ -237,9 +237,10 @@ def make_simulator_params( # get total abundance scaling factor sys_length = len(sim.spin_systems) sys_model_length = len(spin_system_models) - total_abundance = sum(sys.abundance for sys in sim.spin_systems) + total_abundance = 0.0 + total_abundance += sum(sys.abundance for sys in sim.spin_systems) total_abundance += sum(sys.abundance for sys in spin_system_models) - abundance_scale = 100 / total_abundance + abundance_scale = 100.0 / total_abundance # expression for the last abundance. if sys_length > 0: @@ -572,7 +573,7 @@ def make_distribution_params( params = Parameters() for i, model in enumerate(spin_system_models): # normalize the abundance - model.abundance /= norm + model.abundance *= norm params = model.add_lmfit_params(params, i) # Set last abundance parameter as a non-free variable @@ -581,7 +582,7 @@ def make_distribution_params( expr_terms.append(f"{model_prefix}_{i}_abundance") else: expr = "-".join(expr_terms) - expr = "1" if expr == "" else f"1-{expr}" + expr = "100" if expr == "" else f"100 - {expr}" params[f"{model_prefix}_{i}_abundance"].vary = False params[f"{model_prefix}_{i}_abundance"].expr = expr return params diff --git a/src/mrsimulator/utils/tests/test_spectral_fitting.py b/src/mrsimulator/utils/tests/test_spectral_fitting.py index 9ae49f5c3..031e95d97 100644 --- a/src/mrsimulator/utils/tests/test_spectral_fitting.py +++ b/src/mrsimulator/utils/tests/test_spectral_fitting.py @@ -10,6 +10,8 @@ from mrsimulator.method.lib import BlochDecaySpectrum from mrsimulator.method.lib import SSB2D from mrsimulator.method.lib import ThreeQ_VAS +from mrsimulator.models import CzjzekDistribution +from mrsimulator.models import ExtCzjzekDistribution def test_str_encode(): @@ -357,3 +359,27 @@ def test_array(): # params = sf.make_LMFIT_params(sim, processor) # a = sf.LMFIT_min_function(params, sim, processor) # np.testing.assert_almost_equal(-a.sum(), dataset.sum().real, decimal=8) + + +def test_distribution(): + dist1 = CzjzekDistribution(sigma=0.5) + dist2 = CzjzekDistribution(sigma=1.4, mean_isotropic_chemical_shift=4.7) + dist3 = ExtCzjzekDistribution(symmetric_tensor={"Cq": 1e6, "eta": 0.1}, eps=0.5) + + models = [dist1, dist2, dist3] + + params = sf.make_LMFIT_params(spin_system_models=models) + + assert params["czjzek_0_sigma"] == 0.5 + assert params["czjzek_0_mean_isotropic_chemical_shift"] == 0.0 + assert np.allclose(params["czjzek_0_abundance"], 100.0 / 3.0) + + assert params["czjzek_1_sigma"] == 1.4 + assert params["czjzek_1_mean_isotropic_chemical_shift"] == 4.7 + assert np.allclose(params["czjzek_1_abundance"], 100.0 / 3.0) + + assert params["ext_czjzek_2_symmetric_tensor_zeta"] == 1e6 + assert params["ext_czjzek_2_symmetric_tensor_eta"] == 0.1 + assert params["ext_czjzek_2_eps"] == 0.5 + assert params["ext_czjzek_2_mean_isotropic_chemical_shift"] == 0.0 + assert np.allclose(params["ext_czjzek_2_abundance"], 100.0 / 3.0)