Perfect model \(k-\varepsilon\) calibration on Kato-Phillips case
In this notebook, we will demonstrate how get started with Tunax to run a forward model and do a perfect model calibration. Our approach will use the \(k-\varepsilon\) closure and will be based on the idealized Kato-Phillips [1] case. This case is characterized by the absence of heat flux and the presence of uniform zonal wind forcing. In a perfect-model framework, the “observations” used for calibration are outputs of a forward model run, generated using a specific set of \(k-\varepsilon\) parameters. The goal is for Tunax to successfully retrieve these original parameters through the calibration process.
[1]:
import jax.numpy as jnp
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.pyplot import Figure, Axes
from typing import List, Tuple, TypeAlias
subplot_1D_type: TypeAlias = Tuple[Figure, List[Axes]]
subplot_2D_type: TypeAlias = Tuple[Figure, List[List[Axes]]]
plt.rc('text', usetex=True)
plt.rcParams.update({
'text.usetex': True,
'figure.figsize': (8, 5),
'axes.titlesize': 18,
'figure.titlesize': 18,
'axes.labelsize': 14,
'xtick.labelsize': 12,
'ytick.labelsize': 12,
'legend.fontsize': 12,
'lines.linewidth': 2,
'lines.markersize': 6
})
Forward model
In this part we will run the single column model that goes with Tunax, with the \(k-\varepsilon\) closure with some initial parameters and on the Kato-Phillips case to create our database of “observations”. To make a model, Tunax needs a geometry, a initial state, a physical case and some time parameters.
Geometry
First we have to define the geometry of the water column for our model with the class Grid. We use a simple linear geometry with the class method linear to create a regular grid with 100 points on a depth of 50 meters. This object contains in part the vector of the depths of the centers of the cells zr and the vector of the vector of the interfaces zw.
[2]:
from tunax import Grid
grid = Grid.linear(50, 50)
grid
[2]:
Grid(nz=50, hbot=f32[], zr=f32[50], zw=f32[51], hz=f32[50])
Initial state
The initial state for the Kato-Phillips idealized case is defind by zero velocities, a fully stratified temperature and a constant salinity. The class State represent these 4 variables state of the water column. To built our initial state we follow these steps 1. we initialize a State object defined on our grid, every variable is set to 0 2. we reinitialize the temparature and the salinity to our specific case with init_t and init_s (note that we use the default slope of the
stratification for the temperature) We are obliged to create a new instance of the State object everytime that we want to modify an attribute as long as JAX doesn’t allow in place modification. It will be the case for all Tunax usage.
[3]:
from tunax import State
init_state = State.zeros(grid)
init_state = init_state.init_t(hmxl=0., t_sfc=16.)
init_state = init_state.init_s(hmxl=100)
init_state
[3]:
State(
grid=Grid(nz=50, hbot=f32[], zr=f32[50], zw=f32[51], hz=f32[50]),
u=f32[50],
v=f32[50],
t=f32[50],
s=f32[50]
)
[4]:
zr = grid.zr
sp: subplot_1D_type = plt.subplots(1, 4, sharey=True, figsize=(8, 4))
fig, [ax_t, ax_s, ax_u, ax_v] = sp
fig.tight_layout(rect=[0, 0.0, 1, 0.94])
fig.subplots_adjust(wspace=0.1)
ax_t.plot(init_state.t, zr, 'k', label='initial state')
ax_s.plot(init_state.s, zr, 'k')
ax_u.plot(init_state.u, zr, 'k')
ax_v.plot(init_state.v, zr, 'k')
ax_t.set_xlabel(r'$T [{}^\circ\mathrm C]$')
ax_s.set_xlabel(r'$S \ [\mathrm{psu}]$')
ax_u.set_xlabel(r'$u \left[\mathrm m \cdot \mathrm s^{-1}\right]$')
ax_v.set_xlabel(r'$v \left[\mathrm m \cdot \mathrm s^{-1}\right]$')
ax_t.set_ylabel(r'$z [\mathrm m]$')
ax_t.legend()
fig.suptitle('Initial state of Kato-Phillips case')
plt.show()
Physical case of Kato-Phillips
Here we will initialize the physical case of our Kato-Phillips experiments, which as we said correspond to no heat flux and a constant zonal wind of \(0.01 \text m \cdot \text s ^{-1}\). The class Case contains all the physical parameters (forcings and physical constants) for a run of the model. Calling the constructor of this class create an instance with default values in it, they need to be modified to add the forcings.
[5]:
from tunax import Case
import equinox as eqx
kp_case = Case()
kp_case = kp_case.set_u_wind(0.01)
kp_case
[5]:
Case(
rho0=1024.0,
grav=9.81,
cp=3985.0,
alpha=0.0002,
beta=0.0008,
t_rho_ref=0.0,
s_rho_ref=35.0,
vkarmn=0.384,
fcor=0.0,
ustr_sfc=0.0001,
ustr_btm=0.0,
vstr_sfc=0.0,
vstr_btm=0.0,
tflx_sfc=0.0,
tflx_btm=0.0,
sflx_sfc=0.0,
sflx_btm=0.0,
rflx_sfc_max=0.0
)
Initialization of the model
Now that we have a initial condition, a grid and a physical case, we can defined the forward model instance with the class SingleColumnModel. For that we need to add the lentght of our model time_frame, the duration of one time-step dt the time between 2 outputs out_dt, and last but not least, the name of the closure that we will use, here it’s 'k-epsilon'. Here we do a simulation of \(30 \text h\).
[6]:
from tunax import SingleColumnModel
time_frame = 30.
dt = 30.
out_dt = 300.
model = SingleColumnModel(time_frame, dt, out_dt, init_state, kp_case, 'k-epsilon')
model
[6]:
SingleColumnModel(
nt=3600,
dt=30.0,
n_out=10,
init_state=State(
grid=Grid(nz=50, hbot=f32[], zr=f32[50], zw=f32[51], hz=f32[50]),
u=f32[50],
v=f32[50],
t=f32[50],
s=f32[50]
),
case=Case(
rho0=1024.0,
grav=9.81,
cp=3985.0,
alpha=0.0002,
beta=0.0008,
t_rho_ref=0.0,
s_rho_ref=35.0,
vkarmn=0.384,
fcor=0.0,
ustr_sfc=0.0001,
ustr_btm=0.0,
vstr_sfc=0.0,
vstr_btm=0.0,
tflx_sfc=0.0,
tflx_btm=0.0,
sflx_sfc=0.0,
sflx_btm=0.0,
rflx_sfc_max=0.0
),
closure=Closure(
name='k-epsilon',
parameters_class=<class 'tunax.closures.k_epsilon.KepsParameters'>,
state_class=<class 'tunax.closures.k_epsilon.KepsState'>,
step_fun=<wrapped function keps_step>
)
)
Closure parameters
The \(k-\varepsilon\) closure is included in the Tunax sources. One can find it in src/closures/k_epsilon.py. Here we initialize the parameters of \(k-\varepsilon\) with their default values, by calling the constructor of KepsParameters class which contains all these parameters.
[7]:
from tunax.closures import KepsParameters
keps_default_params = KepsParameters()
keps_default_params
[7]:
KepsParameters(
c1=5.0,
c2=0.8,
c3=1.968,
c4=1.136,
c5=0.0,
c6=0.4,
cb1=5.95,
cb2=0.6,
cb3=1.0,
cb4=0.0,
cb5=0.3333,
cbb=0.72,
c_mu0=0.5477,
sig_k=1.0,
sig_eps=1.3,
c_eps1=1.44,
c_eps2=1.92,
c_eps3m=-0.4,
c_eps3p=1.0,
chk_grav=1400.0,
galp=0.53,
z0s_min=0.01,
z0b_min=0.01,
z0b=1e-14,
akt_min=1e-05,
akv_min=0.0001,
tke_min=1e-06,
eps_min=1e-12,
c_mu_min=0.1,
c_mu_prim_min=0.1,
dir_sfc=False,
dir_btm=True,
gls_p=3,
gls_m=1.5,
gls_n=-1,
sf_d0=19913.90625,
sf_d1=5087.41065,
sf_d2=571.8541632000001,
sf_d3=103.99979766988804,
sf_d4=172.80647999999997,
sf_d5=-0.6715392000000003,
sf_n0=2124.1500000265514,
sf_n1=345.3003360057121,
sf_n2=-2.40000000003,
sf_nb0=2231.25,
sf_nb1=90.0,
sf_nb2=17.66630399938841,
lim_am0=42300123.961466245,
lim_am1=17682701.851643827,
lim_am2=2123751.491340626,
lim_am3=59670.13560796435,
lim_am4=1214704.020776464,
lim_am5=418372.60492247937,
lim_am6=35911.165079938415
)
Run the model
Now we can run our model with these default values of \(k-\varepsilon\) parameters. The run don’t modify the model instance (as long as we use JAX which don’t allow in place modification) but returns a object of the class Trajectory. This class is simply the set of the timeseries of the velocities and the tracers. It’s like a pandas dataframe but defined inside Tunax.
[8]:
traj_obs = model.compute_trajectory_with(keps_default_params)
traj_obs
[8]:
Trajectory(
grid=Grid(nz=50, hbot=f32[], zr=f32[50], zw=f32[51], hz=f32[50]),
time=f32[360],
t=f32[360,50],
s=f32[360,50],
u=f32[360,50],
v=f32[360,50]
)
[9]:
zr = grid.zr
sp: subplot_1D_type = plt.subplots(1, 4, sharey=True, figsize=(8, 4))
fig, [ax_t, ax_s, ax_u, ax_v] = sp
fig.tight_layout(rect=[0, 0.0, 1, 0.94])
fig.subplots_adjust(wspace=0.1)
ax_t.plot(traj_obs.t[0, :], zr, 'k', linewidth=1)
ax_t.plot(traj_obs.t[-1, :], zr, 'r')
ax_s.plot(traj_obs.s[0, :], zr, 'k', linewidth=1)
ax_s.plot(traj_obs.s[-1, :], zr, 'r')
ax_u.plot(traj_obs.u[0, :], zr, 'k', linewidth=1)
ax_u.plot(traj_obs.u[-1, :], zr, 'r')
ax_v.plot(traj_obs.v[0, :], zr, 'k', linewidth=1, label='initial state')
ax_v.plot(traj_obs.v[-1, :], zr, 'r', label='final state')
ax_t.set_xlabel(r'$T [{}^\circ\mathrm C]$')
ax_s.set_xlabel(r'$S [\mathrm{psu}]$', labelpad=15)
ax_u.set_xlabel(r'$u \left[\mathrm m \cdot \mathrm s^{-1}\right]$')
ax_v.set_xlabel(r'$v \left[\mathrm m \cdot \mathrm s^{-1}\right]$')
ax_t.set_ylabel(r'$z [\mathrm m]$')
ax_v.legend(framealpha=1.)
fig.suptitle(r'Initial and final states of Kato-Phillips case run with $k-\varepsilon$ closure')
plt.show()
Here we can notice the Stokes drift effect on the zonal velocity \(u\) from the wind forcing. We notice that there is no effect on \(v\) because we don’t consider rotation and Coriolis effect (possible to change it by changing the latitude in Case). Then we notice that the Stokes drift has an effect on the temperature a create a mixed layer with almost constant temperature.
Perfect model calibration
Here we will show how to use the calibration part of Tunax for a problem of a perfect model calibration : we will try to find back the parameters of \(k-\varepsilon\) that we just used for the forward model. To configure a experience of calibration, Tunax needs a database, a configuration on the parameters of the closure to calibrate, a loss function and some parameters for the optimizer algorithm.
Database
In Tunax, a Database is a list of obervations, and an observation (Obs class) is the joint of a Trajetory object (time-series of the evolution of the state) and a Case object that represent a physical situation. These obervations are the basis on what our calibration algorithm is going to make the model and the closure fit on. Typically these observations are extracted from measurments or Large Eddy Simulations (LES), but here, as we are working in a perfect model
framework to be more simple, the observation will be created from the output of the model that we just ran.
[10]:
from tunax import Obs, Database
obs = Obs(traj_obs, kp_case)
database = Database([obs])
database
[10]:
Database(
observations=[
Obs(
trajectory=Trajectory(
grid=Grid(nz=50, hbot=f32[], zr=f32[50], zw=f32[51], hz=f32[50]),
time=f32[360],
t=f32[360,50],
s=f32[360,50],
u=f32[360,50],
v=f32[360,50]
),
case=Case(
rho0=1024.0,
grav=9.81,
cp=3985.0,
alpha=0.0002,
beta=0.0008,
t_rho_ref=0.0,
s_rho_ref=35.0,
vkarmn=0.384,
fcor=0.0,
ustr_sfc=0.0001,
ustr_btm=0.0,
vstr_sfc=0.0,
vstr_btm=0.0,
tflx_sfc=0.0,
tflx_btm=0.0,
sflx_sfc=0.0,
sflx_btm=0.0,
rflx_sfc_max=0.0
)
)
]
)
Configurations of parameters to calibrate
To define what parameters of our closure we are going to calibrate, we use the class FittableParameter which describes for one parameter if we fit it or not and the initial value in this case. The class FittableParametersSet represents the set of all the parameters of the closure that we want to calibrate. Here we calibrate only the parameter \(c_1\) of \(k-\varepsilon\), the initial value during the calibration is 15. The other parameters are set to the default ones. The default
value of \(c_1\) used previously in the perfect model is 5, so we expect the value of 15 decrease to 5 during the calibration.
[11]:
from tunax import FittableParameter, FittableParametersSet
c1_par = FittableParameter(True, 15.)
coef_fit_params = FittableParametersSet({'c1': c1_par}, 'k-epsilon')
coef_fit_params
[11]:
FittableParametersSet(
coef_fit_dict={'c1': FittableParameter(do_fit=True, val=15.0)},
closure=Closure(
name='k-epsilon',
parameters_class=<class 'tunax.closures.k_epsilon.KepsParameters'>,
state_class=<class 'tunax.closures.k_epsilon.KepsState'>,
step_fun=<wrapped function keps_step>
)
)
Loss function
Then we have to define the loss function used in the calibration. This loss function must have the signature Callable[[List[Trajectory], Database], float] and the output should be positive. Here we take the squared \(L_2\) norm of the difference of the temperature between the obervations and the model at every hour and every depth. This loss function will be wrapped to be a function of the parameters to calibrate. To be more specific, let’s note \(\theta\) the set of parameters of
the closure that we want to calibrate and \(U_i\) the vectors of the state at each iteration of the model. Then we can write \(\mathcal M_\theta\) the operator of the model that passes from one step to another one \(U_{i+1} = \mathcal M_\theta (U_i)\). The perfect model that we ran previously with the parameters \(\theta_p\) can be written
where \(N\) is the number of iterations. Now we can write the loss function that we define of a set of the \(k-\varepsilon\) parameters \(\theta\)
where \(T\) is only the projection that keeps the temperature part of a state \(U\), \(h=-50\) m is the depth of our column and \(H\) is the set of index \(i\) so that the time \(t_i\) is a multiple of one hour.
[12]:
from tunax import Trajectory
from typing import List
def loss(trajectories: List[Trajectory], database: Database):
t_obs = database.observations[0].trajectory.t
t_scm = trajectories[0].t
return jnp.sum((t_scm[::10]-t_obs[::10])**2)
Here trajectories represent the different run of the model with the configuration of every case in the database (here we just have one). We can visualize there the location of the observations space with the vertical red lines.
[13]:
time = traj_obs.time/3600
x, y = jnp.meshgrid(time, traj_obs.grid.zr)
fig, ax = plt.subplots()
fig.tight_layout()
pcm = ax.pcolormesh(x, y, traj_obs.t.transpose(), shading='nearest')
fig.colorbar(pcm, label=r'$T [{}^\circ\mathrm C]$')
ax.set_xlabel(r'time $[\mathrm h]$')
ax.set_ylabel(r'$z [\mathrm m]$')
ax.set_title('Temperature evolution for the Kato-Phillips\n case and location of the observations')
for t in time[::12]:
ax.axvline(x=t, color='r', linewidth=1)
plt.show()
Definition and call of the fitter
Adding some parametes for the optmizer, we can define the instance of the class Fitter that will execute the calibration. An obect of this class contains the informations of how to do the calibration. The call of this object will try to find the best parameters of the closure to minmize the loss functions that we defined just before. First, the fitter compute the gradient function of the loss function, which is doable because JAX is a differentiable langage. Then it does a loop where at each
steps it computes the gradient function on a point of the space of the parameters, then an optimization algorithm described by the Optax package computes the next point to explore.
[14]:
from tunax import Fitter
nloop = 50
nit_loss = 1
learning_rate = .5
dt_cal = 300.
output_path = 'k-epsilon_Kato-Phillips/c1_evolution.npz'
f = Fitter(coef_fit_params, database, dt_cal, loss, nloop, nit_loss=nit_loss, learning_rate=learning_rate, output_path=output_path)
keps_params_calibrated = f()
loop 0
x [15.]
grads [0.3391961]
loop 1
x [14.500004]
grads [0.3480592]
loop 2
x [13.999714]
grads [0.36043194]
loop 3
x [13.498813]
grads [0.40461874]
loop 4
x [12.996218]
grads [0.41749904]
loop 5
x [12.492009]
grads [0.43715474]
loop 6
x [11.985976]
grads [0.45326394]
loop 7
x [11.47802]
grads [0.47675622]
loop 8
x [10.967768]
grads [0.51993597]
loop 9
x [10.4543495]
grads [0.5668275]
loop 10
x [9.937118]
grads [0.5592757]
loop 11
x [9.417005]
grads [0.63479054]
loop 12
x [8.892428]
grads [0.66029876]
loop 13
x [8.363532]
grads [0.6891761]
loop 14
x [7.830332]
grads [0.6885269]
loop 15
x [7.293636]
grads [0.69416535]
loop 16
x [6.75394]
grads [0.7283362]
loop 17
x [6.2107224]
grads [0.666131]
loop 18
x [5.6669736]
grads [0.60154927]
loop 19
x [5.1257954]
grads [0.39132735]
loop 20
x [4.599365]
grads [0.05478675]
loop 21
x [4.1147532]
grads [-0.43280727]
loop 22
x [3.725699]
grads [-0.4189247]
loop 23
x [3.4201264]
grads [-0.55001545]
loop 24
x [3.2049959]
grads [-0.727044]
loop 25
x [3.0894427]
grads [-0.72363585]
loop 26
x [3.0588133]
grads [-0.7552156]
loop 27
x [3.1042118]
grads [-0.7055095]
loop 28
x [3.210773]
grads [-0.72228235]
loop 29
x [3.3719087]
grads [-0.5612357]
loop 30
x [3.5679653]
grads [-0.6095512]
loop 31
x [3.7989492]
grads [-0.45418835]
loop 32
x [4.048833]
grads [-0.4836539]
loop 33
x [4.3182073]
grads [-0.25574923]
loop 34
x [4.5864506]
grads [0.04923618]
loop 35
x [4.8261194]
grads [0.20352697]
loop 36
x [5.0247016]
grads [0.3079912]
loop 37
x [5.1753783]
grads [0.45818654]
loop 38
x [5.2675147]
grads [0.42482185]
loop 39
x [5.31006]
grads [0.4522025]
loop 40
x [5.3051]
grads [0.45692396]
loop 41
x [5.2567787]
grads [0.41007257]
loop 42
x [5.1736465]
grads [0.44523564]
loop 43
x [5.0556965]
grads [0.33189687]
loop 44
x [4.916692]
grads [0.26215872]
loop 45
x [4.764993]
grads [0.15116331]
loop 46
x [4.6122513]
grads [0.04555369]
loop 47
x [4.4688406]
grads [-0.09177426]
loop 48
x [4.3477554]
grads [-0.20555884]
loop 49
x [4.258763]
grads [-0.31099942]
The output of the fitter is the final value of the calibrated parameters of the closure (here only \(c_1\) has changed).
[15]:
keps_params_calibrated.c1
[15]:
Array(4.209973, dtype=float32)
The evolution of the calibrated parameters and their gradients have been written and we can visualize them.
[17]:
data = np.load(output_path, allow_pickle=True)
c1_ev = data['x'][0, :]
c1_grad_ev = data['grads'][0, :]
loss_it = data['loss_it']
loss_values = data['loss_values']
sp: subplot_1D_type = plt.subplots(3, 1, sharex=True, figsize=(8, 6))
fig, [ax_x, ax_g, ax_l] = sp
plt.tight_layout(rect=[0, 0.03, 1, 0.98])
ax_x.plot(c1_ev, '.', label='$c_1$ calibration')
ax_x.axhline([5.], color='k', linewidth=1, label='$c_1$ obs')
ax_g.plot(c1_grad_ev, '.', label=r'$\partial_{c_1} \mathcal L$')
ax_g.axhline(0, color='k', linewidth=1)
ax_l.plot(loss_it, loss_values, '.', label=r'$\mathcal L$')
ax_l.set_xlabel('calibration iterations')
ax_x.legend()
ax_g.legend()
ax_l.legend()
fig.suptitle(r'$c_1$ calibration evolution')
plt.show()
Check of the results
One can see that the value of \(c_1\) has decreased from 15 to around 5 as expected. We can run the model with these calibrated parameters.
[18]:
traj_calibrated = model.compute_trajectory_with(keps_params_calibrated)
keps_params_uncalibrated = coef_fit_params.fit_to_closure(coef_fit_params.gen_init_val())
traj_uncalibrated = model.compute_trajectory_with(keps_params_uncalibrated)
Now let’s visualize these results on the final state. Remember that here we did the calibration onlt on the temperature, it’s a reason why the calibration on the meridional speed \(u\) seems to work less.
[31]:
sp: subplot_1D_type = plt.subplots(1, 4, sharey=True, figsize=(8, 4))
fig, [ax_t, ax_s, ax_u, ax_v] = sp
fig.tight_layout(rect=[0, 0.0, 1, 0.94])
fig.subplots_adjust(wspace=0.1)
ax_t.plot(traj_obs.t[0, :], zr, 'k', linewidth=1)
ax_t.plot(traj_obs.t[-1, :], zr, 'r')
ax_t.plot(traj_uncalibrated.t[-1, :], zr, 'g--')
ax_t.plot(traj_calibrated.t[-1, :], zr, 'b--')
ax_s.plot(traj_obs.s[0, :], zr, 'k', linewidth=1)
ax_s.plot(traj_obs.s[-1, :], zr, 'r')
ax_s.plot(traj_uncalibrated.s[-1, :], zr, 'g--')
ax_s.plot(traj_calibrated.s[-1, :], zr, 'b--')
ax_u.plot(traj_obs.u[0, :], zr, 'k', linewidth=1)
ax_u.plot(traj_obs.u[-1, :], zr, 'r')
ax_u.plot(traj_uncalibrated.u[-1, :], zr, 'g--')
ax_u.plot(traj_calibrated.u[-1, :], zr, 'b--')
ax_v.plot(traj_obs.v[0, :], zr, 'k', linewidth=1, label='initial state')
ax_v.plot(traj_obs.v[-1, :], zr, 'r', label='observations')
ax_v.plot(traj_uncalibrated.v[-1, :], zr, 'g--', label='uncalibrated')
ax_v.plot(traj_calibrated.v[-1, :], zr, 'b--', label='calibrated')
ax_t.set_xlim(14, 15.4)
ax_t.set_ylim(-40, 0)
ax_t.set_xlabel(r'$T [{}^\circ\mathrm C]$')
ax_s.set_xlabel(r'$S [\mathrm{psu}]$', labelpad=15)
ax_u.set_xlabel(r'$u \left[\mathrm m \cdot \mathrm s^{-1}\right]$')
ax_v.set_xlabel(r'$v \left[\mathrm m \cdot \mathrm s^{-1}\right]$')
ax_t.set_ylabel(r'$z [\mathrm m]$')
ax_v.legend(framealpha=1.)
plt.suptitle('Final states before and after the calibration of the parameter $c_1$')
plt.show()
References
Kato H, Phillips OM. On the penetration of a turbulent layer into stratified fluid. Journal of Fluid Mechanics. 1969;37(4):643-655. doi : 10.1017/S0022112069000784