Transient absorption spectroscopy fits with rate equation (fake data)
This notebook is a demonstration of how to fit transient absorption spectroscopy (TAS) with the following rate equations:
\[\frac{dn}{dt} = G - k_{trap} n - k_{direct} n^2\]
where \(n\) is the charge carrier density, \(G\) is the generation rate in m⁻³ s⁻¹, ktrap is the trapping rate in s⁻¹, and kdirect is the bimolecular/band-to-bad recombination rate in m³ s⁻¹.
The rate equation described here is the same as employed in the paper by Haffner‐Schirmer et al. 2017 to fit organic solar cell TAS data.
[1]:
# Import necessary libraries
import warnings, os, sys, shutil
# remove warnings from the output
os.environ["PYTHONWARNINGS"] = "ignore"
warnings.filterwarnings(action='ignore', category=FutureWarning)
warnings.filterwarnings(action='ignore', category=UserWarning)
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from numpy.random import default_rng
import torch, copy, uuid
import pySIMsalabim as sim
from pySIMsalabim.experiments.JV_steady_state import *
import ax, logging
from ax.utils.notebook.plotting import init_notebook_plotting, render
init_notebook_plotting() # for Jupyter notebooks
try:
from optimpv import *
from optimpv.RateEqfits.RateEqAgent import RateEqAgent
from optimpv.RateEqfits.RateEqModel import *
from optimpv.RateEqfits.Pumps import *
except Exception as e:
sys.path.append('../') # add the path to the optimpv module
from optimpv import *
from optimpv.RateEqfits.RateEqAgent import RateEqAgent
from optimpv.RateEqfits.RateEqModel import *
from optimpv.RateEqfits.Pumps import *
[INFO 08-14 10:06:00] ax.utils.notebook.plotting: Injecting Plotly library into cell. Do not overwrite or delete cell.
[INFO 08-14 10:06:00] ax.utils.notebook.plotting: Please see
(https://ax.dev/tutorials/visualizations.html#Fix-for-plots-that-are-not-rendering)
if visualizations are not rendering.
Define the parameters for the simulation
[2]:
# Define the parameters to be fitted
# Note: in general for log spaced value it is better to use the foce_log option when doing the Bayesian inference
params = []
k_direct = FitParam(name = 'k_direct', value = 1e-16, bounds = [1e-18,1e-14], log_scale = True, rescale = True, value_type = 'float', type='range', display_name=r'$k_{\text{direct}}$', unit='m$^{3}$ s$^{-1}$', axis_type = 'log',force_log=True)
params.append(k_direct)
k_trap = FitParam(name = 'k_trap', value = 1e6, bounds = [1e5,1e7], log_scale = True, rescale = True, value_type = 'float', type='range', display_name=r'$k_{\text{trap}}$', unit='s$^{-1}$', axis_type = 'log',force_log=True)
params.append(k_trap)
cross_section = FitParam(name = 'cross_section', value = 6e-21, bounds = [1e-21,1e-20], log_scale = True, rescale = True, value_type = 'float', type='range', display_name=r'$\sigma$', unit='m$^{2}$', axis_type = 'log',force_log=True)
params.append(cross_section)
# original values
params_orig = copy.deepcopy(params)
dum_dic = {}
for i in range(len(params)):
if params[i].force_log:
dum_dic[params[i].name] = np.log10(params[i].value)
else:
dum_dic[params[i].name] = params[i].value/params[i].fscale
# we need this just to run the model to generate some fake data
Generate some fake data
Here we generate some fake data to fit. The data is generated using the same model as the one used for the fitting, so it is a good test of the fitting procedure.
[3]:
# Define the necessary input parameters for the model
metric = 'nrmse'
loss = 'linear'
threshold = 10
exp_format = 'TAS'
P = 0.0039 # W
wavelegnth = 850e-9 # m
fpu = 1e5 # Hz
Area = 0.3*0.3*1e-4 # effective pump area in m^-2
penetration_depth = 1e-5*1e-2 # penetration depth in m
pulse_width = 0.2*(1/fpu) # width of the pump pulse in seconds
t0 = 0 # time offset of the pump pulse in seconds
background = 0 # background pump density in m^-3
# Calculate the average absorbed density in photons m^-3
flux, density = get_flux_density(P, wavelegnth, fpu, Area, penetration_depth)
pump_args = {'fpu': fpu , 'pulse_width': pulse_width, 't0': t0, 'P': density, 'background': background, }
fixed_model_args = {'L':100e-9}
t = np.linspace(0,1/fpu,1000)
t = t[:-1]
Gfracs = [0.1,0.5,1]
# concatenate the time and Gfracs
X = None
for Gfrac in Gfracs:
if X is None:
X = np.array([t,Gfrac*np.ones(len(t))]).T
else:
X = np.concatenate((X,np.array([t,Gfrac*np.ones(len(t))]).T),axis=0)
y_ = X # dummy data
RateEq_fake = RateEqAgent(params, [X], [y_], model = BT_model, pump_model = square_pump, pump_args = pump_args, fixed_model_args = fixed_model_args, metric = metric, loss = loss, threshold=threshold,minimize=True,exp_format=exp_format)
y = RateEq_fake.run(dum_dic,exp_format=exp_format)
noise = 1e-7
from numpy.random import default_rng
rng = default_rng()
y += rng.standard_normal(len(X))*noise
plt.figure()
for idx, Gfrac in enumerate(Gfracs):
plt.plot(X[X[:,1]==Gfrac,0], y[X[:,1]==Gfrac],'o',label=str(Gfrac),)
plt.xlabel('Time [s]')
plt.ylabel(exp_format + ' [a.u.]')
plt.show()
Run the optimization
[4]:
# Define the Agent and the target metric/loss function
metric = 'nrmse' # can be 'nrmse', 'mse', 'mae'
loss = 'linear' # can be 'linear', 'huber', 'soft_l1'
RateEq = RateEqAgent(params, [X], [y], model = BT_model, pump_model = square_pump, pump_args = pump_args, fixed_model_args = fixed_model_args, metric = metric, loss = loss, threshold=threshold,minimize=True,exp_format=exp_format)
MCMC Bayesian Inference for Parameter Fitting
We’ll use the emcee package to perform Markov Chain Monte Carlo sampling to find the posterior distribution of our model parameters.
[5]:
from optimpv.axBOtorch.axBOtorchOptimizer import axBOtorchOptimizer
from botorch.acquisition.logei import qLogNoisyExpectedImprovement
from ax.adapter.transforms.standardize_y import StandardizeY
from ax.adapter.transforms.unit_x import UnitX
from ax.adapter.transforms.remove_fixed import RemoveFixed
from ax.adapter.transforms.log import Log
from ax.generators.torch.botorch_modular.utils import ModelConfig
from ax.generators.torch.botorch_modular.surrogate import SurrogateSpec
from gpytorch.kernels import MaternKernel
from gpytorch.kernels import ScaleKernel
from botorch.models import SingleTaskGP
model_gen_kwargs_list = None
parameter_constraints = None
model_kwargs_list = [{},{"torch_device":torch.device("cuda" if torch.cuda.is_available() else "cpu"),'botorch_acqf_class':qLogNoisyExpectedImprovement,'transforms':[RemoveFixed, Log,UnitX, StandardizeY],'surrogate_spec':SurrogateSpec(model_configs=[ModelConfig(botorch_model_class=SingleTaskGP,covar_module_class=ScaleKernel, covar_module_options={'base_kernel':MaternKernel(nu=2.5, ard_num_dims=len(params))})])}]
optimizer = axBOtorchOptimizer(params = params, agents = RateEq, models = ['SOBOL','BOTORCH_MODULAR'],n_batches = [1,45], batch_size = [10,2], ax_client = None, max_parallelism = 100, model_kwargs_list = model_kwargs_list, model_gen_kwargs_list = model_gen_kwargs_list, name = 'ax_opti',parameter_constraints = parameter_constraints,)
[6]:
# optimizer.optimize() # run the optimization with ax
optimizer.optimize_turbo() # run the optimization with turbo
[INFO 08-14 10:06:01] optimpv.axBOtorchOptimizer: Starting optimization with 46 batches and a total of 100 trials
[INFO 08-14 10:06:01] optimpv.axBOtorchOptimizer: Starting Sobol batch 1 with 10 trials
[INFO 08-14 10:06:02] optimpv.axBOtorchOptimizer: Finished Sobol
[INFO 08-14 10:06:03] optimpv.axBOtorchOptimizer: Finished Turbo batch 2 with 2 trials with current best value: 8.42e-02, TR length: 8.00e-01
[INFO 08-14 10:06:03] optimpv.axBOtorchOptimizer: Finished Turbo batch 3 with 2 trials with current best value: 4.75e-02, TR length: 8.00e-01
[INFO 08-14 10:06:04] optimpv.axBOtorchOptimizer: Finished Turbo batch 4 with 2 trials with current best value: 2.99e-02, TR length: 8.00e-01
[INFO 08-14 10:06:04] optimpv.axBOtorchOptimizer: Finished Turbo batch 5 with 2 trials with current best value: 2.99e-02, TR length: 8.00e-01
[INFO 08-14 10:06:04] optimpv.axBOtorchOptimizer: Finished Turbo batch 6 with 2 trials with current best value: 2.15e-02, TR length: 8.00e-01
[INFO 08-14 10:06:04] optimpv.axBOtorchOptimizer: Finished Turbo batch 7 with 2 trials with current best value: 2.04e-02, TR length: 8.00e-01
[INFO 08-14 10:06:04] optimpv.axBOtorchOptimizer: Finished Turbo batch 8 with 2 trials with current best value: 1.53e-02, TR length: 1.60e+00
[INFO 08-14 10:06:05] optimpv.axBOtorchOptimizer: Finished Turbo batch 9 with 2 trials with current best value: 1.53e-02, TR length: 1.60e+00
[INFO 08-14 10:06:05] optimpv.axBOtorchOptimizer: Finished Turbo batch 10 with 2 trials with current best value: 1.53e-02, TR length: 1.60e+00
[INFO 08-14 10:06:05] optimpv.axBOtorchOptimizer: Finished Turbo batch 11 with 2 trials with current best value: 1.53e-02, TR length: 1.60e+00
[INFO 08-14 10:06:05] optimpv.axBOtorchOptimizer: Finished Turbo batch 12 with 2 trials with current best value: 1.53e-02, TR length: 1.60e+00
[INFO 08-14 10:06:06] optimpv.axBOtorchOptimizer: Finished Turbo batch 13 with 2 trials with current best value: 1.53e-02, TR length: 1.60e+00
[INFO 08-14 10:06:06] optimpv.axBOtorchOptimizer: Finished Turbo batch 14 with 2 trials with current best value: 1.53e-02, TR length: 1.60e+00
[INFO 08-14 10:06:06] optimpv.axBOtorchOptimizer: Finished Turbo batch 15 with 2 trials with current best value: 1.53e-02, TR length: 1.60e+00
[INFO 08-14 10:06:06] optimpv.axBOtorchOptimizer: Finished Turbo batch 16 with 2 trials with current best value: 1.29e-02, TR length: 1.60e+00
[INFO 08-14 10:06:06] optimpv.axBOtorchOptimizer: Finished Turbo batch 17 with 2 trials with current best value: 1.29e-02, TR length: 1.60e+00
[INFO 08-14 10:06:06] optimpv.axBOtorchOptimizer: Finished Turbo batch 18 with 2 trials with current best value: 1.29e-02, TR length: 1.60e+00
[INFO 08-14 10:06:07] optimpv.axBOtorchOptimizer: Finished Turbo batch 19 with 2 trials with current best value: 1.29e-02, TR length: 1.60e+00
[INFO 08-14 10:06:07] optimpv.axBOtorchOptimizer: Finished Turbo batch 20 with 2 trials with current best value: 1.29e-02, TR length: 1.60e+00
[INFO 08-14 10:06:07] optimpv.axBOtorchOptimizer: Finished Turbo batch 21 with 2 trials with current best value: 1.29e-02, TR length: 1.60e+00
[INFO 08-14 10:06:07] optimpv.axBOtorchOptimizer: Finished Turbo batch 22 with 2 trials with current best value: 1.29e-02, TR length: 1.60e+00
[INFO 08-14 10:06:07] optimpv.axBOtorchOptimizer: Finished Turbo batch 23 with 2 trials with current best value: 1.29e-02, TR length: 1.60e+00
[INFO 08-14 10:06:07] optimpv.axBOtorchOptimizer: Finished Turbo batch 24 with 2 trials with current best value: 1.29e-02, TR length: 1.60e+00
[INFO 08-14 10:06:08] optimpv.axBOtorchOptimizer: Finished Turbo batch 25 with 2 trials with current best value: 1.29e-02, TR length: 1.60e+00
[INFO 08-14 10:06:08] optimpv.axBOtorchOptimizer: Finished Turbo batch 26 with 2 trials with current best value: 1.29e-02, TR length: 1.60e+00
[INFO 08-14 10:06:08] optimpv.axBOtorchOptimizer: Finished Turbo batch 27 with 2 trials with current best value: 1.29e-02, TR length: 1.60e+00
[INFO 08-14 10:06:08] optimpv.axBOtorchOptimizer: Finished Turbo batch 28 with 2 trials with current best value: 1.29e-02, TR length: 1.60e+00
[INFO 08-14 10:06:08] optimpv.axBOtorchOptimizer: Finished Turbo batch 29 with 2 trials with current best value: 1.29e-02, TR length: 1.60e+00
[INFO 08-14 10:06:08] optimpv.axBOtorchOptimizer: Finished Turbo batch 30 with 2 trials with current best value: 1.14e-02, TR length: 1.60e+00
[INFO 08-14 10:06:09] optimpv.axBOtorchOptimizer: Finished Turbo batch 31 with 2 trials with current best value: 1.14e-02, TR length: 1.60e+00
[INFO 08-14 10:06:09] optimpv.axBOtorchOptimizer: Finished Turbo batch 32 with 2 trials with current best value: 1.14e-02, TR length: 1.60e+00
[INFO 08-14 10:06:09] optimpv.axBOtorchOptimizer: Finished Turbo batch 33 with 2 trials with current best value: 1.14e-02, TR length: 1.60e+00
[INFO 08-14 10:06:09] optimpv.axBOtorchOptimizer: Finished Turbo batch 34 with 2 trials with current best value: 1.14e-02, TR length: 1.60e+00
[INFO 08-14 10:06:09] optimpv.axBOtorchOptimizer: Finished Turbo batch 35 with 2 trials with current best value: 9.80e-03, TR length: 1.60e+00
[INFO 08-14 10:06:10] optimpv.axBOtorchOptimizer: Finished Turbo batch 36 with 2 trials with current best value: 9.80e-03, TR length: 1.60e+00
[INFO 08-14 10:06:10] optimpv.axBOtorchOptimizer: Finished Turbo batch 37 with 2 trials with current best value: 9.80e-03, TR length: 1.60e+00
[INFO 08-14 10:06:10] optimpv.axBOtorchOptimizer: Finished Turbo batch 38 with 2 trials with current best value: 9.80e-03, TR length: 1.60e+00
[INFO 08-14 10:06:10] optimpv.axBOtorchOptimizer: Finished Turbo batch 39 with 2 trials with current best value: 9.80e-03, TR length: 1.60e+00
[INFO 08-14 10:06:10] optimpv.axBOtorchOptimizer: Finished Turbo batch 40 with 2 trials with current best value: 9.80e-03, TR length: 1.60e+00
[INFO 08-14 10:06:11] optimpv.axBOtorchOptimizer: Finished Turbo batch 41 with 2 trials with current best value: 9.80e-03, TR length: 1.60e+00
[INFO 08-14 10:06:11] optimpv.axBOtorchOptimizer: Finished Turbo batch 42 with 2 trials with current best value: 9.80e-03, TR length: 1.60e+00
[INFO 08-14 10:06:11] optimpv.axBOtorchOptimizer: Finished Turbo batch 43 with 2 trials with current best value: 9.80e-03, TR length: 1.60e+00
[INFO 08-14 10:06:11] optimpv.axBOtorchOptimizer: Finished Turbo batch 44 with 2 trials with current best value: 9.80e-03, TR length: 1.60e+00
[INFO 08-14 10:06:11] optimpv.axBOtorchOptimizer: Finished Turbo batch 45 with 2 trials with current best value: 9.80e-03, TR length: 1.60e+00
[INFO 08-14 10:06:12] optimpv.axBOtorchOptimizer: Finished Turbo batch 46 with 2 trials with current best value: 9.80e-03, TR length: 1.60e+00
[INFO 08-14 10:06:12] optimpv.axBOtorchOptimizer: Finished Turbo batch 47 with 2 trials with current best value: 7.99e-03, TR length: 1.60e+00
[INFO 08-14 10:06:12] optimpv.axBOtorchOptimizer: Turbo is terminated as the max number (100) of trials is reached
[7]:
# get the best parameters and update the params list in the optimizer and the agent
ax_client = optimizer.ax_client # get the ax client
optimizer.update_params_with_best_balance() # update the params list in the optimizer with the best parameters
RateEq.params = optimizer.params # update the params list in the agent with the best parameters
# print the best parameters
print('Best parameters:')
for p,po in zip(optimizer.params, params_orig):
print(p.name, 'fitted value:', p.value, 'original value:', po.value)
Best parameters:
k_direct fitted value: 8.943333078659956e-17 original value: 1e-16
k_trap fitted value: 938021.1109050045 original value: 1000000.0
cross_section fitted value: 5.721490735439374e-21 original value: 6e-21
[8]:
# Plot optimization results
data = ax_client.summarize()
all_metrics = optimizer.all_metrics
plt.figure()
plt.plot(np.minimum.accumulate(data[all_metrics]), label="Best value seen so far")
plt.yscale("log")
plt.xlabel("Iteration")
plt.ylabel("Target metric: " + metric + " with " + loss + " loss")
plt.legend()
plt.title("Best value seen so far")
print("Best value seen so far is ", min(data[all_metrics[0]]), "at iteration ", int(data[all_metrics[0]].idxmin()))
plt.show()
Best value seen so far is 0.007989130617933934 at iteration 101
[9]:
# rerun the simulation with the best parameters
yfit = RateEq.run(parameters={},exp_format=exp_format) # run the simulation with the best parameters
viridis = plt.get_cmap('viridis', len(Gfracs))
plt.figure(figsize=(10,10))
linewidth = 2
for idx, Gfrac in enumerate(Gfracs[::-1]):
plt.plot(X[X[:,1]==Gfrac,0],y[X[:,1]==Gfrac],label='Gfrac = '+str(Gfrac),color=viridis(idx),alpha=0.5,linewidth=linewidth)
plt.plot(X[X[:,1]==Gfrac,0],yfit[X[:,1]==Gfrac],label='Gfrac = '+str(Gfrac)+' fit',linestyle='--',color=viridis(idx),linewidth=linewidth)
plt.xlabel('Time [s]')
plt.ylabel(exp_format + ' [a.u.]')
plt.legend()
plt.show()
