Expand lattice strain example to include fitting a profile #85

Author: morganjwilliams

docs/source/gallery/examples/geochem/mineral_lattice.py

@@ -134,6 +134,45 @@
 ax.legend(bbox_to_anchor=(1.05, 1))
 fig
 ########################################################################################
+# Given the lattice strain model and a partitioning profile (for e.g. REE data),
+# we can also fit a model to a given curve; here we fit to our REE data above,
+# for which we have some known parameters to compare to:
+#
+from pyrolite.mineral.lattice import fit_lattice_strain, _lattice_opt_function
+
+_ri, _tk, _D = fit_lattice_strain(np.array(site3radii), site3Ds, z=3)
+########################################################################################
+# We can compare the results of this fit to our source parameters - the ionic radius of
+# La, 900°C and estimated :math:`D_{La}`:
+#
+import pandas as pd
+
+pd.DataFrame(
+    [(_ri, rLa), (Tk, _tk), (D_La, _D)],
+    index=["radii", "T", "D"],
+    columns=["Source Parameters", "Fit Parameters"],
+).T
+
+########################################################################################
+# We can also compare the curves visually:
+#
+from pyrolite.plot.spider import REE_v_radii
+
+ax = REE_v_radii()
+
+ax.plot(site3radii, site3Ds, label="True", color="k")
+ax.plot(
+    site3radii,
+    _lattice_opt_function(site3radii, site3radii.mean(), 298 + 1000, 1, z=3),
+    label="Starting Guess",
+    ls="--",
+    color="0.5",
+)
+
+ax.plot(site3radii, fit_lattice_strain._opt(site3radii, _ri, _tk, _D, z=3), label="Fit")
+
+ax.figure
+########################################################################################
 # .. [#ref_1] Blundy, J., Wood, B., 1994. Prediction of crystal–melt partition coefficients
 #             from elastic moduli. Nature 372, 452.
 #             doi: `10.1038/372452A0 <https://doi.org/10.1038/372452A0>`__

pyrolite/mineral/lattice.py

@@ -28,7 +28,11 @@
 
 """
 
+from functools import partial
+
 import numpy as np
+import pandas as pd
+from scipy.optimize import curve_fit
 
 from ..util.log import Handle
 from ..util.meta import sphinx_doi_link, update_docstring_references
@@ -104,7 +108,7 @@ def strain_coefficient(ri, rx, r0=None, E=None, T=298.15, z=None, **kwargs):
     """
     n = 6.023 * 10**23
     if r0 is None:
-        logger.warn("Use fictive ideal cation radii where possible.")
+        logger.warning("Use fictive ideal cation radii where possible.")
         r0 = ri
     ri, rx, r0 = ri / 10**10, rx / 10**10, r0 / 10**10  # convert to meters
     E = E or youngs_modulus_approximation(z, r0)  # use E if defined, else try to calc
@@ -182,6 +186,60 @@ def youngs_modulus_approximation(z, r):
     return E
 
 
+def _lattice_opt_function(xs, ri, Tk, D, z=3, E=None):
+    return D * strain_coefficient(
+        ri, xs, r0=ri, E=E or youngs_modulus_approximation(z, ri), T=Tk
+    )
+
+
+def fit_lattice_strain(
+    radii,
+    ys,
+    E=None,
+    z=3,
+    bounds=[(0.1, 2.2), (273.15, 273.15 + 2700), (0, np.inf)],
+    r0=None,
+    t0=273.15 + 500,
+    d0=1.0,
+    **kwargs,
+):
+    """
+    Fit a lattice strain model to a given set of abundances.
+
+    Parameters
+    ----------
+    radii : :class:`numpy.ndarray`
+        Radii to fit against.
+    ys : :class:`numpy.ndarray`
+        Partition coefficients for given elemental data.
+    E : :class:`float`, :code:`None`
+        Young's modulus (stiffness) for the site, in pascals (Pa). Will be estimated using
+        :func:`youngs_modulus_approximation` if none is given.
+    z : :class:`int`
+        Optional specification of cationic valence, for calcuation of approximate
+        Young's modulus using :func:`youngs_modulus_approximation`,
+        where the modulus is not specified.
+    bounds : :class:`list`
+        List of tuples specifying bounds on parameters `ri`, `T` and `D`.
+
+    Returns
+    -------
+    ri, tk, D : :class:`float`
+        Radius, temperature and partition coefficeint describing the
+        lattice strain fit.
+    """
+
+    popt, pcov = curve_fit(
+        partial(_lattice_opt_function, z=z, E=E),
+        radii,
+        ys,
+        p0=[r0 if r0 is not None else radii.mean(), t0, d0],
+        bounds=pd.DataFrame(bounds).T.values,
+        **kwargs,
+    )
+    return popt
+
+
 __doc__ = __doc__.format(
     brice1975=sphinx_doi_link("10.1016/0022-0248(75)90241-9"),
     blundy1994=sphinx_doi_link("10.1038/372452a0"),