%config InlineBackend.figure_format = 'retina' 

"""
Explanation: For high dpi displays.
End of explanation
"""


import numpy as np
import matplotlib.pyplot as plt
import pytheos as eos

"""
Explanation: 0. General note


Through high pressure experiments, we can measure the molar volume of minerals at high pressure.  The data can be then fit to static equations of state, in order to obtain bulk modulus $(K_0)$ and pressure derivative of bulk modulus $(K'_0)$.  


pytheos uses lmfit to provide equation of state fitting for static compression data.  It has a few convenient features for the flexibility, such as fixing/varying particular parameters, propagating uncertainties, etc.  This notebook shows a few examples.


1. Global setup
End of explanation
"""


v0 = 11.244; k0 = 160.; k0p = 4.0

"""
Explanation: 2. Setup fitting models
Using MgO, we create models for three equations for static EOS.  Volume can be either in unit-cell volume or molar volume.  But v and v0 should be in the same unit.
End of explanation
"""


exp_bm3 = eos.BM3Model()
exp_vinet = eos.VinetModel()
exp_kunc = eos.KuncModel(order=5) # this way I can handle order

"""
Explanation: Three different equations for static EOS of MgO.
End of explanation
"""


params = exp_bm3.make_params(v0=v0, k0=k0, k0p=k0p)

"""
Explanation: Assign initial values to the parameters.
End of explanation
"""


v_data = v0 * np.linspace(0.99,0.6,20)

"""
Explanation: 3. Synthesize data
We make data with random noise.
End of explanation
"""


p_bm3 = exp_bm3.eval(params, v=v_data)
p_vinet = exp_vinet.eval(params, v=v_data)
p_kunc = exp_kunc.eval(params, v=v_data)

"""
Explanation: Calculate pressures from three different equations.
End of explanation
"""


sp = np.random.normal(0.,2.,20)
p_data = p_bm3 + sp

"""
Explanation: Create random noise (noise) and add it to pressure value to simulate the data.
End of explanation
"""


plt.plot(p_data, v_data, 'ko', label='data')
plt.plot(p_bm3, v_data, label='bm3')
plt.plot(p_vinet, v_data, label='vinet')
plt.plot(p_kunc, v_data, label='kunc')
plt.xlabel('Pressure (GPa)'); plt.ylabel('Molar volume (cm$^3$/mol)')
plt.legend();

"""
Explanation: Plot the synthetic data together with true values.
End of explanation
"""


plt.plot(p_bm3, p_vinet-p_bm3, label='Vinet - BM3')
plt.plot(p_bm3, p_kunc-p_bm3, label='Kunc - BM3')
plt.xlabel('Pressure (GPa)'); plt.ylabel('$\Delta P$')
plt.axhline(0, c='k', ls=':')
plt.legend();

"""
Explanation: The cell below shows the systematic differences between the equations.
End of explanation
"""


params['v0'].vary = False
#params['k0'].vary = False

"""
Explanation: 4. EOS fitting
We fix $V_0$ as it is typically well known.  We can also fix any parameters in lmfit.
End of explanation
"""


weight = None # 1./sp#None

"""
Explanation: It is common practice not to include weight for the $P$-$V$ fit because it will de-emphasize the high pressure data points.  But if you need, you can set weight in the following way.  None performs unweighted fitting
End of explanation
"""


fitresult_bm3 = exp_bm3.fit(p_data, params, v=v_data, weights=weight)

"""
Explanation: The cell below performs fitting with BM3 equation.
End of explanation
"""


print(fitresult_bm3.fit_report())

"""
Explanation: You may print out the fitting result.
End of explanation
"""


fitresult_vinet = exp_vinet.fit(p_data, params, v=v_data, weights=weight)
print(fitresult_vinet.fit_report())

fitresult_kunc = exp_kunc.fit(p_data, params, v=v_data, weights=weight)
print(fitresult_kunc.fit_report())

"""
Explanation: Fitting can be performed for other equations.
End of explanation
"""


fitresult_vinet.params['k0p'].value

"""
Explanation: 5. Play with the results
The example below shows how to get individual parameters from the fitting result.
End of explanation
"""


fitresult_vinet.params['k0p'].stderr

"""
Explanation: You can also get estimated uncertainties.
End of explanation
"""


v0_r = fitresult_vinet.params['v0']
k0_r = fitresult_vinet.params['k0']
k0p_r = fitresult_vinet.params['k0p']

"""
Explanation: The cells below show how to calculate fitted pressure values at the data points.  First, we get the fit parameters.
End of explanation
"""


f = fitresult_vinet.model

"""
Explanation: Then we get the function used for the fitting.
End of explanation
"""


f.eval(params, v=v_data)

f.func(v_data, v0, k0, k0p)

"""
Explanation: eval conducts calculation for given volume value.
End of explanation
"""


v_fitline = np.linspace(v0,v_data.min(),100)

p_fitline = f.func(v_fitline, v0_r, k0_r, k0p_r)

plt.plot(p_fitline, v_fitline)
plt.plot(p_data, v_data, 'k.', label='data')
plt.xlabel('Pressure (GPa)'); plt.ylabel('Molar volume (cm$^3$/mol)');

"""
Explanation: If you want to get array of (p, v) for plotting smooth fitting curve:
End of explanation
"""


fitresult_bm3.params

"""
Explanation: The cell below shows how to get the fitting results with uncertainties.
End of explanation
"""


from uncertainties import ufloat
p_err = f.func(v_fitline, ufloat(fitresult_bm3.params['v0'].value, 0.),
         ufloat(fitresult_bm3.params['k0'].value, fitresult_bm3.params['k0'].stderr),
         ufloat(fitresult_bm3.params['k0p'].value, fitresult_bm3.params['k0p'].stderr))

"""
Explanation: The cell below shows how to get the uncertainties of the fit curve.
End of explanation
"""


import pandas as pd
df = pd.DataFrame()
df['p'] = p_err
df.head()

"""
Explanation: The cell below shows the results in a table.  It only shows the first five data points, but p_err contains 100 points as we setup that way above.
End of explanation
"""


eos.plot.static_fit_result(fitresult_bm3, p_err=sp)

"""
Explanation: pytheos has internal plot script, which takes care of the plotting the fitting results with useful information.
End of explanation
"""


eos.plot.static_fit_result(fitresult_vinet)

eos.plot.static_fit_result(fitresult_kunc)

"""
Explanation: How about the fit results for the other functions?
End of explanation
"""


