In [ ]:
from kineticlib import particles, omegaint, loaddata, ratesvibr
from scipy import constants
from scipy.integrate import odeint
from scipy.optimize import fsolve
import numpy as np
import matplotlib.pyplot as plt

In [ ]:
%matplotlib inline

In [ ]:
d_T_step = 2.

In [ ]:
global global_print
global_print = True

In [ ]:
def eta(T, omint_22):
    return 5. * constants.k * T / (8. * omint_22)

In [ ]:
def deta_dtau(T, omint_22, domint_22_dT):
    return T_wall_2 * 5. * constants.k * (omint_22 - T * domint_22_dT) / (8. * omint_22**2)

In [ ]:
def lambd(T, omint_22, omint_11):
    return constants.k**2 * T * (75. / (32. * mol.mass * omint_22) + 3. / (8. * mol.mass * omint_11))

In [ ]:
def dlambd_dtau(T, omint_22, domint_22_dT, omint_11, domint_11_dT):
    return T_wall_2 * constants.k**2 * (75. * (omint_22 - T * domint_22_dT) / (32. * mol.mass * omint_22**2) + 
                                        3. * (omint_11 - T * domint_11_dT) / (8. * mol.mass * omint_11**2))

In [ ]:
def D_coeff(T, omint_11, n):
    return (3. * constants.k * T) / (8. * n * mol.mass * omint_11)

In [ ]:
def dD_coeff_dx(T, omint_11, n):
    return -3. * constants.k * T * n_wall_2 / (8. * n**2 * mol.mass * omint_11)

In [ ]:
def dD_coeff_dtau(T, omint_11, domint_11_dT, n):
    return T_wall_2 * 3. * constants.k * (omint_11 - T * domint_11_dT) / (8. * n * mol.mass * omint_11**2)

In [ ]:
def make_vt_rates(T):
    res = np.zeros(mol.num_vibr)  # i + 1 -> i
    for i in range(mol.num_vibr):
        res[i] = ratesvibr.vt_rate_ssh(T, d, mol, i+1, -1)
    return res

In [ ]:
def make_reverse_vt_rates(T, rates_vec):
    # vec[i] = i + 1 -> i
    # need to return i -> i + 1
    return rates_vec * np.exp((mol.vibr[:mol.num_vibr] - mol.vibr[1:mol.num_vibr+1]) / (constants.k * T))

In [ ]:
def make_reverse_vv_rates(rates_mat):
    # mat[i, k] = (i + 1) + k -> i + (k + 1)
    # need to return i + (k + 1)  -> (i + 1) + k
    res = np.zeros(mol.num_vibr, mol.num_vibr) 
    for i in range(mol.num_vibr):
        res[i, :] = rates_mat[i, :] * np.exp((mol.vibr[i] + mol.vibr[1:mol.num_vibr+1]
                                              - mol.vibr[i+1] - mol.vibr[:mol.num_vibr]) / (constants.k * T))
    return res

In [ ]:
def LHS_xi_tau_mat(tau, xi_vec, x_tot, D, lambd_val):
    lhs_mat = np.zeros([mol.num_vibr + 2, mol.num_vibr + 2])
    lhs_mat[:mol.num_vibr, :mol.num_vibr+1] = 1. / x_tot
    for i in range(mol.num_vibr):
        for j in range(mol.num_vibr):
            if i == j:
                lhs_mat[i, j] -= 1. / xi_vec[i]  # TODO - redo via digaonal, needs numpy 1.10
    lhs_mat[mol.num_vibr, :mol.num_vibr+1] = tau
    lhs_mat[mol.num_vibr, mol.num_vibr+1] = x_tot
    lhs_mat[mol.num_vibr+1, :mol.num_vibr+1] = mol.vibr
    lhs_mat[mol.num_vibr+1, :mol.num_vibr+1] -= np.dot(mol.vibr, xi_vec) / x_tot
    lhs_mat[mol.num_vibr+1, :mol.num_vibr+1] *= D * n_wall_2 / lambd_val
    lhs_mat[mol.num_vibr+1, mol.num_vibr+1] = T_wall_2
#     print(lhs_mat)
#     global global_print
#     if global_print:
#         for i in range(mol.num_vibr + 2):
#             print('MAT', i, np.max(lhs_mat[i, :]))
    return lhs_mat

In [ ]:
def RHS_xi_tau_vec(tau, dtau, xi_vec, dxi_vec, x_tot, dx_tot, dvel_val,# flow parameters
                   D, dD_dtau, dD_dx, eta_val, deta_val_dtau, lambd_val, dlambd_val_dtau,  # transport coefficients
                   R_i_vibr_vec):
    rhs_vec = np.zeros(mol.num_vibr + 2)
    for i in range(mol.num_vibr):
        rhs_vec[i] = 0.
    rhs_vec[:mol.num_vibr] = dxi_vec[:mol.num_vibr] * (dxi_vec[:mol.num_vibr] / xi_vec[:mol.num_vibr]**2
                                                       - dx_tot / (xi_vec[:mol.num_vibr] * x_tot))
    rhs_vec[:mol.num_vibr] += (dD_dx * dx_tot + dD_dtau * dtau) * (dxi_vec[:mol.num_vibr] / xi_vec[:mol.num_vibr]
                                                                   - dx_tot / x_tot) / D
    rhs_vec[:mol.num_vibr] += (dx_tot / x_tot)**2 - (dxi_vec[:mol.num_vibr] / xi_vec[:mol.num_vibr])**2
    global global_print
    if global_print:
        for i in range(5):
            print('RHS - 1', i, rhs_vec[i])
    rhs_vec[:mol.num_vibr] += Y**2 * R_i_vibr_vec[:mol.num_vibr] / (D * xi_vec[:mol.num_vibr] * n_wall_2)
    rhs_vec[mol.num_vibr] = -2 * dtau * dx_tot
    if global_print:
        for i in range(5):
            print('RHS - 2', i, rhs_vec[i])
    
#     rhs_tau = 0.
#     for i in range(mol.num_vibr + 1):
#         rhs_tau -= mol.vibr[i] * ((dD_dx * d_x_tot + dD_dtau * d_tau) * (d_xi_vec[i] - xi_vec[i] * d_x_tot / x_tot)
#                                   + D * (xi_vec[i] * (d_x_tot / x_tot)**2) - d_xi_vec[i] * d_x_tot / x_tot)
    rhs_tau = -np.dot(mol.vibr, ((dD_dx * dx_tot + dD_dtau * dtau) * (dxi_vec - xi_vec * dx_tot / x_tot)
                                 + D * (xi_vec * (dx_tot / x_tot)**2 - dxi_vec * dx_tot / x_tot)))
    rhs_tau *= n_wall_2 / lambd_val
    rhs_tau -= eta_val * (v_wall_2 * dvel_val)**2 / lambd_val + T_wall_2 * dtau**2 * dlambd_val_dtau / lambd_val
    rhs_vec[mol.num_vibr + 1] = rhs_tau
#     global global_print
#     if global_print:
#         for i in range(mol.num_vibr + 2):
#             print(i, rhs_vec[i])
    return rhs_vec

In [ ]:
def second_derivatives_value_xi_tau_vec(tau, dtau, xi_vec, dxi_vec, x_tot, dx_tot, dvel_val,
                                        D, dD_dtau, dD_dx, eta_val, deta_val_dtau, lambd_val, dlambd_val_dtau,
                                        R_i_vibr_vec):
    """
    Calculate the values of the second derivatives of variables x_{i} and tau
    """
    lhs_mat = LHS_xi_tau_mat(tau, xi_vec, x_tot, D, lambd_val)
    rhs_vec = RHS_xi_tau_vec(tau, dtau, xi_vec, dxi_vec, x_tot, dx_tot, dvel_val,
                             D, dD_dtau, dD_dx, eta_val, deta_val_dtau, lambd_val, dlambd_val_dtau,
                             R_i_vibr_vec)
    
    return np.linalg.solve(lhs_mat, rhs_vec)

In [ ]:
def sysvec(y, t):
    """
    Right-hand side of the Couette ODE system in dimensionless variables
    d_xi - [0,..,num_vibr]
    d2_xi - [num_vibr+1,..,2*num_vibr+1]
    d_v - [2*num_vibr+2]
    d2_v - [2*num_vibr+3]
    d_tau - [2*num_vibr+4]
    d2_tau - [2*num_vibr+5]
    """
    rhs_vec = np.zeros(2*mol.num_vibr+6)
    R_i_vibr_vec = np.zeros(mol.num_vibr+1)
#     x1, x2, dx1, dx2 = y
    xi_vec = y[:mol.num_vibr+1]
    dxi_vec = y[mol.num_vibr+1:2*mol.num_vibr+2]
    vel_val = y[2*mol.num_vibr+2]
    dvel_val = y[2*mol.num_vibr+3]
    tau = y[2*mol.num_vibr+4]
    dtau = y[2*mol.num_vibr+5]
    
    rhs_vec[:mol.num_vibr+1] = dxi_vec
    rhs_vec[2*mol.num_vibr+2] = dvel_val
    rhs_vec[2*mol.num_vibr+4] = dtau
    
    T = tau * T_wall_2
    x_tot = np.sum(xi_vec)
    dx_tot = np.sum(dxi_vec)
    n = n_wall_2 * x_tot
    
    omint_11 = omegaint.omega(T, 1, 1, d)
    domint_11 = (omegaint.omega(T + d_step, 1, 1, d, dimensional=True)
                 - omegaint.omega(T - d_step, 1, 1, d, dimensional=True)) / (2. * d_step)
    omint_22 = omegaint.omega(T, 2, 2, d)
    domint_22 = (omegaint.omega(T + d_step, 2, 2, d, dimensional=True)
                 - omegaint.omega(T - d_step, 2, 2, d, dimensional=True)) / (2. * d_step)
    
    eta_val = eta(T, omint_22)
    deta_val_dtau = deta_dtau(T, omint_22, domint_22)
    lambd_val = lambd(T, omint_22, omint_11)
    dlambd_val_dtau = dlambd_dtau(T, omint_22, domint_22, omint_11, domint_11)
    
    D = D_coeff(T, omint_11, n)
    dD_dx = dD_coeff_dx(T, omint_11, n)
    dD_dtau = dD_coeff_dtau(T, omint_11, domint_11, n)
    
#     vt_rate_down = np.zeros(mol.num_vibr)  # i + 1 -> i
    # TODO ADD RATE CALCULATION
    vt_rate_down = make_vt_rates(T)
    vt_rate_up = make_reverse_vt_rates(T, vt_rate_down)  # i -> i + 1
    
#     vv_rate_down = np.zeros(mol.num_vibr, mol.num_vibr)  # (i + 1) + k -> i + (k + 1)
    # TODO ADD RATE CALCULATION
#     vv_rate_up = make_reverse_vv_rates(T, vv_rate_down)   # i + (k + 1)  -> (i + 1) + k
    
#     for i in range(mol.num_vibr-1):
#         R_i_vibr_vec[i+1] = xi_vec[i] * vt_rate_up[i] - xi_vec[i+1] * vt_rate_down[i]
#         R_i_vibr_vec[i+1] += xi_vec[i+2] * vt_rate_down[i+1] - xi_vec[i+1] * vt_rate_up[i+1]
        
    R_i_vibr_vec[1:mol.num_vibr] = xi_vec[:mol.num_vibr-1] * vt_rate_up[:mol.num_vibr-1]
    R_i_vibr_vec[1:mol.num_vibr] -= xi_vec[1:mol.num_vibr] * vt_rate_down[:mol.num_vibr-1]
    R_i_vibr_vec[1:mol.num_vibr] += xi_vec[2:mol.num_vibr+1] * vt_rate_down[1:mol.num_vibr]
    R_i_vibr_vec[1:mol.num_vibr] -= xi_vec[1:mol.num_vibr] * vt_rate_up[1:mol.num_vibr]
    
    R_i_vibr_vec[0] = xi_vec[1] * vt_rate_down[0] - xi_vec[0] * vt_rate_up[0]
    R_i_vibr_vec[mol.num_vibr] = xi_vec[mol.num_vibr - 1] * vt_rate_up[mol.num_vibr - 1]
    R_i_vibr_vec[mol.num_vibr] -= xi_vec[mol.num_vibr] * vt_rate_down[mol.num_vibr - 1]
    R_i_vibr_vec *= n * n_wall_2  # VT transitions
    
    # TODO – VV transitions
    
    d2xi_d2_tau_vec = second_derivatives_value_xi_tau_vec(tau, dtau, xi_vec, dxi_vec, x_tot, dx_tot, dvel_val,
                                                          D, dD_dtau, dD_dx, eta_val, deta_val_dtau,
                                                          lambd_val, dlambd_val_dtau,
                                                          R_i_vibr_vec)
    rhs_vec[mol.num_vibr+1:2*mol.num_vibr+2] = d2xi_d2_tau_vec[:mol.num_vibr+1]
    rhs_vec[2*mol.num_vibr+3] = - dvel_val * dtau * deta_val_dtau / eta_val
    rhs_vec[2*mol.num_vibr+5] = d2xi_d2_tau_vec[mol.num_vibr+1]
    
    global global_print
    if global_print:
        for i in range(2*mol.num_vibr + 6):
            print(i, rhs_vec[i])
        global_print = False
        print('R_I', np.sum(R_i_vibr_vec), R_i_vibr_vec[0:10])
    return rhs_vec

In [ ]:
def give_solution_dimensionless(start_d_values):
    """
    Returns a solution of the Couette ODE system in dimensionless variables,
    given the values of their derivatives on the lower wall
    """
    full_cauchy_vec = np.zeros(2*mol.num_vibr+6)
    full_cauchy_vec[:mol.num_vibr+1] = xi_wall_1[:]  # xi
    full_cauchy_vec[mol.num_vibr+1:2*mol.num_vibr+2] = start_d_values[:mol.num_vibr+1]  # dxi/dy
    full_cauchy_vec[2*mol.num_vibr+2] = 0.  # the velocity is zero
    full_cauchy_vec[2*mol.num_vibr+3] = start_d_values[mol.num_vibr+1]
    full_cauchy_vec[2*mol.num_vibr+4] = T_wall_1 / T_wall_2
    full_cauchy_vec[2*mol.num_vibr+5] = start_d_values[mol.num_vibr+2]
    return odeint(sysvec, full_cauchy_vec, y_span)

In [ ]:
def dimensionless_obj(start_d_values):
    """
    The objective function for which we need to find the roots to solve the Couette ODE system, dimensionless version
    param_vec is a list containing an approximation of the velocity and temperature derivatives
    """
    res = give_solution_dimensionless(start_d_values)
    tmp = np.zeros(mol.num_vibr+3)
    
#     for i in range(mol.num_vibr+1):
#         tmp[i] = res[-1, i] - xi_wall_2[i]
    tmp[:mol.num_vibr+1] = res[-1, :mol.num_vibr+1] - xi_wall_2[:]
    tmp[mol.num_vibr+1] = res[-1, 2*mol.num_vibr+2] - 1
    tmp[mol.num_vibr+2] = res[-1, 2*mol.num_vibr+4] - 1
    return tmp

In [ ]:
def full_solve_dimensionless():
    omint_22 = omegaint.omega(T_wall_1, 2, 2, d)
    omint_11 = omegaint.omega(T_wall_1, 2, 2, d)
    eta_lw = eta(T_wall_1, omint_22)
    lamb_lw = lambd(T_wall_1, omint_22, omint_11)
    start_cond_T = (T_wall_2 - T_wall_1 + 0.5 * eta_lw / lamb_lw) / T_wall_1
    
    start_d_values = np.zeros(mol.num_vibr+3)
    start_d_values[:mol.num_vibr+1] = 0.
    start_d_values[mol.num_vibr+1] = 1.
    start_d_values[mol.num_vibr+2] = start_cond_T
    start_d_values = fsolve(dimensionless_obj, start_d_values)
    return give_solution_dimensionless(start_d_values)

In [ ]:
def convert_to_dimensional(sol_xi, sol_vx, sol_T):
    return [sol_xi * n_wall_2, sol_vx * v_wall_2, sol_T * T_wall_2]

In [ ]:
def create_boltzmann_distribution(T):
    v_exp = np.exp(-mol.vibr[:] / (constants.k * T))
    return v_exp / np.sum(v_exp)

In [ ]:
def test_boltzmann_Ri():
    xi_v = create_boltzmann_distribution(5000.)
    vt_rd = make_vt_rates(5000.)
    n = 10. ** 23
    vt_ru = make_reverse_vt_rates(5000., vt_rd)  # i -> i + 1
    
    R_i_vibr_vec = np.zeros(mol.num_vibr+1)
    
    R_i_vibr_vec[1:mol.num_vibr] = xi_v[:mol.num_vibr-1] * vt_ru[:mol.num_vibr-1]
    R_i_vibr_vec[1:mol.num_vibr] -= xi_v[1:mol.num_vibr] * vt_rd[:mol.num_vibr-1]
    
#     R_i_vibr_vec[1:mol.num_vibr] = xi_v[1:mol.num_vibr] * vt_ru[:mol.num_vibr-1] * (xi_v[:mol.num_vibr-1]
#                                                                                     / xi_v[1:mol.num_vibr]
#                                                                                     - vt_rd[:mol.num_vibr-1]
#                                                                                     / vt_ru[:mol.num_vibr-1])
#     R_i_vibr_vec[1:mol.num_vibr] += xi_v[1:mol.num_vibr] * vt_rd[1:mol.num_vibr] * (xi_v[2:mol.num_vibr+1]
#                                                                                     / xi_v[1:mol.num_vibr]
#                                                                                     - vt_ru[1:mol.num_vibr]
#                                                                                     / vt_rd[1:mol.num_vibr]) 
    R_i_vibr_vec[1:mol.num_vibr] += xi_v[2:mol.num_vibr+1] * vt_rd[1:mol.num_vibr]
    R_i_vibr_vec[1:mol.num_vibr] -= xi_v[1:mol.num_vibr] * vt_ru[1:mol.num_vibr]
    
    R_i_vibr_vec[0] = vt_rd[0] * xi_v[1] - vt_ru[0] * xi_v[0]
#     R_i_vibr_vec[0] = vt_rd[0] * xi_v[0] * (xi_v[1] / xi_v[0] - vt_ru[0] / vt_rd[0])
#     R_i_vibr_vec[mol.num_vibr] = xi_v[mol.num_vibr - 1] * vt_rd[mol.num_vibr - 1] * (vt_ru[mol.num_vibr - 1]
#                                                                                      / vt_rd[mol.num_vibr - 1]
#                                                                                      - xi_v[mol.num_vibr]
#                                                                                      / xi_v[mol.num_vibr - 1])
    R_i_vibr_vec[mol.num_vibr] = xi_v[mol.num_vibr - 1] * vt_ru[mol.num_vibr - 1]
    R_i_vibr_vec[mol.num_vibr] -= xi_v[mol.num_vibr] * vt_rd[mol.num_vibr - 1]
    R_i_vibr_vec *= n * n
    return R_i_vibr_vec

In [ ]:
mol = particles.Molecule('N2', 'harmonic')
d = loaddata.load_elastic_parameters(mol, mol)

In [ ]:
dddd = test_boltzmann_Ri()
print(dddd[0:6], dddd[-1], np.sum(dddd))
print(dddd[0])

r10 = ratesvibr.vt_rate_ssh(5000., d, mol, 1, -1)
r01 = r10 * np.exp(-(mol.vibr[1] - mol.vibr[0])/ (constants.k * 5000.))

t_xi = create_boltzmann_distribution(5000.)
print((r10 * t_xi[1] - r01 * t_xi[0]) * 10**46,
      t_xi[0] * r10 * (t_xi[1] / t_xi[0] - r01 / r10) * 10**46)

In [ ]:
v_wall_2 = 300.  # the velocity of the moving wall
T_wall_2 = 2000.  # the temperature of the moving wall
T_wall_1 = 1000.  # the temperature of the stationary wall
n_wall_1 = 10.**21
n_wall_2 = 10.**21
xi_wall_1 = create_boltzmann_distribution(T_wall_1) * n_wall_1 / n_wall_2
xi_wall_2 = create_boltzmann_distribution(T_wall_2)
Y = 1.  # the distance between the walls

In [ ]:
d_step = 5.
points = 100.
points_per_meter = 100
y_span = np.linspace(0, 1., points)

In [ ]:
global_print = True
sol_dimless = full_solve_dimensionless()

In [ ]:
sol_dimless.shape

In [ ]:
s2 = convert_to_dimensional(sol_dimless[:, :mol.num_vibr+1],
                            sol_dimless[:, 2*mol.num_vibr+2],
                            sol_dimless[:, 2*mol.num_vibr+4])

In [ ]: