For perse

import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt
import math

from SYSTEMMODELS.Functions.system_properties import *
from SYSTEMMODELS.Functions.h_desublimation_CO2 import *
from SYSTEMMODELS.Functions.h_vaporisation_H2O import *
from SYSTEMMODELS.Functions.pressure_liquefaction_H2O import *
from SYSTEMMODELS.Functions.pressure_sublimation_CO2 import *
from SYSTEMMODELS.Functions.molar_desublimation_rate_CO2 import *
from SYSTEMMODELS.Functions.molar_liquefaction_rate_H2O import *
from SYSTEMMODELS.Functions.q_Ib import *
from SYSTEMMODELS.Functions.q_gI import *
from SYSTEMMODELS.Functions.q_star_CO2 import *
from SYSTEMMODELS.Functions.q_star_H2O import *
from SYSTEMMODELS.Functions.U_star import *

#Model:
#Flow of CO2/H2O/N2 through a cold packed bed with constant mass and heat transfer coefficients

#initial conditions
rho0CO2 = 1.1
rho0H2O = 1.1
y0CO2 = 0.1
y0H2O = 0.1
Tg0 = 200
Tb0 = 200
M0CO2 = 0 #solid moles deposited
M0H2O = 0
n = 5
initial_constants = n, rho0CO2, rho0H2O, y0CO2, y0H2O, Tg0, Tb0, M0CO2, M0H2O

##global constants (can remove actual constants and rename as name in system_properties. The non real constants will just be deleted/replaced with the functions)
A_C = AREA_CROSS_SECTION
A_s = SURFACE_AREA_SPECIFIC
h_vap_h2o = h_vaporisation_h2o
h_desublimation_CO2 = h_desublimation_CO2
k_s = k_alumina
dz = dz
U_b = 100
U_g = 100
cp_CO2 = 45
cp_H2O = 45
ki_co2 = 8 #mol/m2.s
ki_h2o = 16 #mol/m2.s
m_co2 = rho0CO2*dz*A_C
m_h2o = rho0H2O*dz*A_C
M_al = rho_alumina/n ##check this

#Boundary conditions
mol_in = 0.5
y_co2_in = 0.1
y_h2o_in  = 0.01
T_in = 293

system_constants = A_C, A_s, h_vap_h2o, h_desublimation_CO2, k_s, dz, U_b, U_g, cp_CO2, cp_H2O, ki_co2, ki_h2o,m_co2,m_h2o, M_al, mol_in,y_co2_in,y_h2o_in,T_in

def initialize_system(*initial_constants):
    ##sets the appropriate initial conditions (column vector of constants)
    n, rho0CO2, rho0H2O, y0CO2, y0H2O, Tg0, Tb0, M0CO2, M0H2O = initial_constants  # unpack initial conditions

    # Initial conditions values are the same at all positions except boundary at t=0
    constants_array = np.asarray(initial_constants)
    constants_array = np.delete(constants_array,0)
    x0 = np.reshape(np.array([constants_array]*n),-1)
    # [rho0CO2, rho0H2O, y0CO2, y0H2O, Tg0, Tb0, M0CO2,...rho0CO2, rho0H2O, ynCO2, ynH2O, Tgn, Tbn, MnCO2, MnH2O]
    return x0


def df(t,x,*system_constants):
    ##the time derivative of x
    A_C, A_s, h_vap_h2o, h_desublimation_CO2, k_s, dz, U_b, U_g, cp_CO2, cp_H2O, ki_co2, ki_h2o, m_co2,m_h2o, M_al, mol_in, y_co2_in, y_h2o_in, T_in = system_constants #unpack system constants
    ##setting up solution array
    x = x.reshape(n,6)
    dxdt = np.zeros_like(x)

    ##Set up boundary conditions here (n=0, t=t i.e. at position 0, the inlet, at all times)
    #note: the reason I set the boundary conditions here rather than in initialize_system
    # is because I need to reference the solution array (initialize_system does not take
    # x as an argument meaning I could not do this
    MDR_CO2 = MDR(ki_co2, x[0][4], psubCO2(x[0][4]))
    MLR_H2O = MLR(ki_h2o, x[0][4], pliqH2O(x[0][4]))

    qstar = (h_desublimation_CO2*MDR_CO2) + (MLR_H2O*h_vaporisation_h2o)
    ustar = U_star(U_g,U_b)
    QGI = (-U_g*qstar/(U_g + U_b)) + (ustar*(x[0][4]-x[0][5]))
    QIB = (U_b*qstar/(U_g + U_b)) + (ustar*(x[0][4]-x[0][5]))

    # CO2 mole fraction (dyi/dt)
    dxdt[0][0] = (mol_in * (y_co2_in - x[0][0]) - (MDR_CO2 * A_s) + (x[0][0]*A_s * (MDR_CO2+MLR_H2O))) / m_co2
    # H2O mole fraction (dyi/dt)
    dxdt[0][1] = (mol_in * (y_h2o_in - x[1][0]) - (MLR_H2O * A_s) + (x[1][0] * A_s*(MDR_CO2+MLR_H2O))) / m_h2o
    # CO2 solid deposition (dMi/dt)
    dxdt[0][2] = MDR_CO2 * A_s / (A_C * dz)
    # H2O solid deposition (or liq) (dMi/dt)
    dxdt[0][3] = MLR_H2O * A_s / (A_C * dz)
    # Gas phase heat balance (dTg/dt)
    dxdt[0][4] = (mol_in * cp_CO2*(T_in - x0[4][0]) - QGI * A_s) / (m_co2 * cp_CO2)
    # Bed heat balance (dTb/dt)
    dxdt[0][5] = (QIB * A_s) / (M_al * cp_alumina)
    # mj
    mdot = mol_in + (rho0CO2 / x[0][4]) * A_C * dz * dxdt[0][4] - (MDR_CO2 * A_s)

    ## internal nodes
    for j in range(1,n): ##start at index 1 (to avoid BCs), go to n (last node)
        ## CO2 mole fraction (dyi/dt)
        MDR_CO2 = MDR(ki_co2, x[j][4], psubCO2(x[j][4]))
        MLR_H2O = MLR(ki_h2o, x[j][4], pliqH2O(x[j][4]))

        qstar = MDR_CO2 + MLR_H2O
        ustar = U_star(U_g, U_b)
        QGI = (-U_g * qstar / (U_g + U_b)) + (ustar * (x[j][4] - x[j][5]))
        QIB = (U_b * qstar / (U_g + U_b)) + (ustar * (x[j][4] - x[j][5]))

        dxdt[j][0] = (mdot*(x[j-1][0] - x[j][0]) - MDR_CO2*A_s + x[j][0]*1)/m_co2
        ## H2O mole fraction (dyi/dt)
        dxdt[j][1] = (mdot*(x[j-1][1] - x[j][1]) - MDR_CO2*A_s + x[j][1]*1)/m_co2
        ## CO2 solid deposition (dMi/dt)
        dxdt[j][2] = MDR_CO2*A_s/(A_C*dz)
        ## H2O solid deposition (or liq) (dMi/dt)
        dxdt[j][3] = MLR_H2O*A_s/(A_C*dz)
        ## Gas phase heat balance (dTg/dt)
        dxdt[j][4] = (mdot*cp_CO2*(x[j-1][4] - x[j][4]) - QGI * A_s)/(m_co2*cp_CO2)
        ## Bed heat balance (dTb/dt)
        dxdt[j][5] = (QIB * A_s)/(M_al*cp_alumina)
        ## mj
        mdot = mdot + (rho0CO2/x[j][4])*A_C*dz*dxdt[j][4] - (MDR_CO2*A_s)

    return dxdt.reshape(-1)

#create your initial conditions
t0 = 0
tf = 100 #how long you want to integrate for
x0 = initialize_system(*initial_constants)

#solve the system
sol = solve_ivp(df, [t0,tf], x0, dense_output = True) #dense_output gives you an interpolant I think (3rd/4th order polynomial)

#decide at which times you'd like to check your system
t_eval = [0,10,20,30,40,50,60,70,80,90,100] ##better way to increment probably

x_eval = [sol.sol(t) for t in t_eval]