BASTET/main.py

import sys
import warnings
import numpy as np
import xarray as xr
import pandas as pd
import cartopy.crs as ccrs
import astropy.units as unit
import matplotlib.pyplot as plt
from dask import delayed
from datetime import datetime
from netCDF4 import Dataset
from astropy.time import Time
from scipy import interpolate
from scipy.spatial import cKDTree
from scipy.integrate import solve_ivp
from input.user_input import *
from input.natural_constants import *
from models.gravity import grav
from models.valving import valving
from models.ballasting import ballasting
from models.sun import sun_angles_analytical, tau
from models.drag import drag, cd_PalumboLow, cd_Palumbo, cd_PalumboHigh, cd_sphere
from models.transformation import visible_cells, transform, radii

starttime = datetime.now()

if not sys.warnoptions:
    warnings.simplefilter("ignore")

data = pd.read_excel(r'C:\Users\marcel\PycharmProjects\MasterThesis\Data_PoGo2016.xls', sheet_name='Tabelle3')  # Tabelle3

comp_time = pd.DataFrame(data, columns=['Time']).to_numpy().squeeze()
comp_height = pd.DataFrame(data, columns=['Height']).to_numpy().squeeze()
comp_lat = pd.DataFrame(data, columns=['Latitude']).to_numpy().squeeze()
comp_lon = pd.DataFrame(data, columns=['Longitude']).to_numpy().squeeze()

print("")
print("Initialising simulation...")
print("Launch location : longitude %.4f, latitude: %.4f" % (start_lon, start_lat))
print("Launch time: " + str(start_utc) + " (UTC)")
print("")
print("Reading ERA5-datasets, please wait.")

ascend_data = xr.open_dataset("ascend_2019_kiruna.nc")
#float_data = xr.open_dataset("float_2019.nc")

df = xr.open_mfdataset(['float1_2019.nc', 'float2_2019.nc', 'float3_2019.nc', 'float4_2019.nc'], combine='by_coords', concat_dim="time", parallel=True)

startNC = Dataset('float1_2019.nc')
endNC = Dataset('float4_2019.nc')
start = int(startNC.variables['time'][0])
end = int(endNC.variables['time'][-1])

float_data = df.assign_coords(time=np.linspace(start, end, (end - start) + 1))

single_data = xr.open_dataset("single_2019.nc")

ERAtime = float_data.variables['time'][:]  # time


# ERA5 MULTI-LEVEL ASCENT (WIND + ATMOSPHERIC DATA DURING ASCENT)

ERAlat0 = ascend_data.variables['latitude'][:].values       # latitude [deg]
ERAlon0 = ascend_data.variables['longitude'][:].values      # longitude [deg]
ERAz_ascend = ascend_data.variables['z'][:].values / g      # geopotential [m^-2/s^-2] to geopotential height [m]
ERApress_ascend = ascend_data.variables['level'][:].values  # pressure level [-]
ERAtemp_ascend = ascend_data.variables['t'][:].values       # air temperature in K
vw_x_ascend = ascend_data.variables['u'][:].values          # v_x in [m/s]
vw_y_ascend = ascend_data.variables['v'][:].values          # v_y in [m/s]
vw_z_ascend = ascend_data.variables['w'][:].values          # v_z in [m/s]


# ERA5 MULTI-LEVEL FLOAT (WIND + ATMOSPHERIC DATA DURING FLOAT)

ERAlat1 = float_data.variables['latitude'][:].values      # latitude [deg]
ERAlon1 = float_data.variables['longitude'][:].values     # longitude [deg]
ERAz_float = float_data.variables['z'][:].values / g      # geopotential [m^-2/s^-2] to geopotential height [m]
ERApress_float = float_data.variables['level'][:].values  # pressure level [-]
ERAtemp_float = float_data.variables['t'][:].values       # air temperature in [K]
vw_x_float = float_data.variables['u'][:].values          # v_x in [m/s]
vw_y_float = float_data.variables['v'][:].values          # v_y in [m/s]
vw_z_float = float_data.variables['w'][:].values          # v_z in [m/s]


# ERA5 SINGLE-LEVEL (RADIATIVE ENVIRONMENT)

ERAlat2 = single_data.variables['latitude'][:].values   # latitude [deg]
ERAlon2 = single_data.variables['longitude'][:].values  # longitude [deg]
ERAtcc = single_data.variables['tcc'][:].values         # total cloud cover [-]
ERAskt = single_data.variables['skt'][:].values         # skin (ground) temperature in [K]
ERAcbh = single_data.variables['cbh'][:].values         # cloud base height in [m]
ERAlcc = single_data.variables['lcc'][:].values         # low cloud cover [-]
ERAmcc = single_data.variables['mcc'][:].values         # medium cloud cover [-]
ERAhcc = single_data.variables['hcc'][:].values         # high cloud cover [-]
ERAssr = single_data.variables['ssr'][:].values         # hourly accumulated surface net solar radiation [J/m^2]
ERAstrn = single_data.variables['str'][:].values        # hourly accumulated surface net thermal radiation [J/m^2]
ERAstrd = single_data.variables['strd'][:].values       # hourly accumulated surface thermal radiation downwards [J/m^2]
ERAssrd = single_data.variables['ssrd'][:].values       # hourly accumulated surface solar radiation downwards [J/m^2]
ERAtsr = single_data.variables['tsr'][:].values         # hourly accumulated top net solar radiation [J/m^2]
ERAttr = single_data.variables['ttr'][:].values         # hourly accumulated top net thermal radiation [J/m^2]
ERAtisr = single_data.variables['tisr'][:].values       # hourly accumulated TOA incident solar radiation [J/m^2]
ERAstrdc = single_data.variables['strdc'][:].values     # hourly accumulated surface thermal radiation downward clear-sky [J/m^2]
ERAsp = single_data.variables['sp'][:].values           # surface pressure in [Pa]

print("Finished reading ERA5-datasets.")

lon_era2d0, lat_era2d0 = np.meshgrid(ERAlon0, ERAlat0)
lon_era2d1, lat_era2d1 = np.meshgrid(ERAlon1, ERAlat1)
lon_era2d2, lat_era2d2 = np.meshgrid(ERAlon2, ERAlat2)

xs0, ys0, zs0 = transform(lon_era2d0.flatten(), lat_era2d0.flatten())
xs1, ys1, zs1 = transform(lon_era2d1.flatten(), lat_era2d1.flatten())
xs2, ys2, zs2 = transform(lon_era2d2.flatten(), lat_era2d2.flatten())

print("")
tree0 = cKDTree(np.column_stack((xs0, ys0, zs0)))
tree1 = cKDTree(np.column_stack((xs1, ys1, zs1)))
tree2 = cKDTree(np.column_stack((xs2, ys2, zs2)))
print("Built kd-trees")


def ERA5Data(lon, lat, h, t, deltaT_ERA):
    t_epoch = deltaT_ERA + t / 3600

    t_pre = int(t_epoch)
    t_pre_ind = t_pre - int(ERAtime[0])
    t_post_ind = t_pre_ind + 1

    xt, yt, zt = transform(lon, lat)  # current coordinates

    d0, inds0 = tree0.query(np.column_stack((xt, yt, zt)), k=4)
    d1, inds1 = tree1.query(np.column_stack((xt, yt, zt)), k=4)  # longitude, latitude
    d2, inds2 = tree2.query(np.column_stack((xt, yt, zt)), k=visible_cells(h))  # longitude, latitude  visible_cells(h)

    w0 = 1.0 / d0[0]
    w1 = 1.0 / d1[0]
    w2 = 1.0 / d2[0] ** 2

    lat_ind0 = np.unravel_index(inds0[0], lon_era2d0.shape)[0]
    lon_ind0 = np.unravel_index(inds0[0], lon_era2d0.shape)[1]
    lat_ind1 = np.unravel_index(inds1[0], lon_era2d1.shape)[0]
    lon_ind1 = np.unravel_index(inds1[0], lon_era2d1.shape)[1]
    lat_ind2 = np.unravel_index(inds2[0], lon_era2d2.shape)[0]
    lon_ind2 = np.unravel_index(inds2[0], lon_era2d2.shape)[1]

    if h >= 30000:
        try:
            interp4d_temp_pre = np.ma.dot(w1, ERAtemp_float[t_pre_ind, :, lat_ind1, lon_ind1]) / np.sum(w1)
            interp4d_temp_post = np.ma.dot(w1, ERAtemp_float[t_post_ind, :, lat_ind1, lon_ind1]) / np.sum(w1)
            interp4d_temp = (interp4d_temp_post - interp4d_temp_pre) * (t_epoch - t_pre) + interp4d_temp_pre

            interp4d_height_pre = np.ma.dot(w1, ERAz_float[t_pre_ind, :, lat_ind1, lon_ind1]) / np.sum(w1)
            interp4d_height_post = np.ma.dot(w1, ERAz_float[t_post_ind, :, lat_ind1, lon_ind1]) / np.sum(w1)
            interp4d_height = (interp4d_height_post - interp4d_height_pre) * (t_epoch - t_pre) + interp4d_height_pre

            interp4d_vw_x_pre = np.ma.dot(w1, vw_x_float[t_pre_ind, :, lat_ind1, lon_ind1]) / np.sum(w1)
            interp4d_vw_x_post = np.ma.dot(w1, vw_x_float[t_post_ind, :, lat_ind1, lon_ind1]) / np.sum(w1)
            interp4d_vw_x = (interp4d_vw_x_post - interp4d_vw_x_pre) * (t_epoch - t_pre) + interp4d_vw_x_pre

            interp4d_vw_y_pre = np.ma.dot(w1, vw_y_float[t_pre_ind, :, lat_ind1, lon_ind1]) / np.sum(w1)
            interp4d_vw_y_post = np.ma.dot(w1, vw_y_float[t_post_ind, :, lat_ind1, lon_ind1]) / np.sum(w1)
            interp4d_vw_y = (interp4d_vw_y_post - interp4d_vw_y_pre) * (t_epoch - t_pre) + interp4d_vw_y_pre

            interp4d_vw_z_pre = np.ma.dot(w1, vw_z_float[t_pre_ind, :, lat_ind1, lon_ind1]) / np.sum(w1)
            interp4d_vw_z_post = np.ma.dot(w1, vw_z_float[t_post_ind, :, lat_ind1, lon_ind1]) / np.sum(w1)
            interp4d_vw_z = (interp4d_vw_z_post - interp4d_vw_z_pre) * (t_epoch - t_pre) + interp4d_vw_z_pre

            pressure_hPa = np.array([1, 2, 3, 5, 7, 10, 20])

            pressure = 100 * pressure_hPa

            temp_interp1d = interpolate.interp1d(interp4d_height, interp4d_temp)
            press_interp1d = interpolate.interp1d(interp4d_height, pressure)
            vw_x_interp1d = interpolate.interp1d(interp4d_height, interp4d_vw_x)
            vw_y_interp1d = interpolate.interp1d(interp4d_height, interp4d_vw_y)
            vw_z_interp1d = interpolate.interp1d(interp4d_height, interp4d_vw_z)

        except IndexError:
            print("Error: Please check time range of ERA5 data!")
    elif np.abs(lat - start_lat) <= 10.0 and np.abs(lon - start_lon) <= 10.0:
        try:
            interp4d_temp_pre = np.ma.dot(w0, ERAtemp_ascend[t_pre_ind, :, lat_ind0, lon_ind0]) / np.sum(w0)
            interp4d_temp_post = np.ma.dot(w0, ERAtemp_ascend[t_post_ind, :, lat_ind0, lon_ind0]) / np.sum(w0)
            interp4d_temp = (interp4d_temp_post - interp4d_temp_pre) * (t_epoch - t_pre) + interp4d_temp_pre

            interp4d_height_pre = np.ma.dot(w0, ERAz_ascend[t_pre_ind, :, lat_ind0, lon_ind0]) / np.sum(w0)
            interp4d_height_post = np.ma.dot(w0, ERAz_ascend[t_post_ind, :, lat_ind0, lon_ind0]) / np.sum(w0)
            interp4d_height = (interp4d_height_post - interp4d_height_pre) * (t_epoch - t_pre) + interp4d_height_pre

            interp4d_vw_x_pre = np.ma.dot(w0, vw_x_ascend[t_pre_ind, :, lat_ind0, lon_ind0]) / np.sum(w0)
            interp4d_vw_x_post = np.ma.dot(w0, vw_x_ascend[t_post_ind, :, lat_ind0, lon_ind0]) / np.sum(w0)
            interp4d_vw_x = (interp4d_vw_x_post - interp4d_vw_x_pre) * (t_epoch - t_pre) + interp4d_vw_x_pre

            interp4d_vw_y_pre = np.ma.dot(w0, vw_y_ascend[t_pre_ind, :, lat_ind0, lon_ind0]) / np.sum(w0)
            interp4d_vw_y_post = np.ma.dot(w0, vw_y_ascend[t_post_ind, :, lat_ind0, lon_ind0]) / np.sum(w0)
            interp4d_vw_y = (interp4d_vw_y_post - interp4d_vw_y_pre) * (t_epoch - t_pre) + interp4d_vw_y_pre

            interp4d_vw_z_pre = np.ma.dot(w0, vw_z_ascend[t_pre_ind, :, lat_ind0, lon_ind0]) / np.sum(w0)
            interp4d_vw_z_post = np.ma.dot(w0, vw_z_ascend[t_post_ind, :, lat_ind0, lon_ind0]) / np.sum(w0)
            interp4d_vw_z = (interp4d_vw_z_post - interp4d_vw_z_pre) * (t_epoch - t_pre) + interp4d_vw_z_pre

            pressure_hPa = np.array(
                [1, 2, 3, 5, 7, 10, 20, 30, 50, 70, 100, 125, 150, 175, 200, 225, 250, 300, 350, 400,
                 450, 500, 550, 600, 650, 700, 750, 775, 800, 825, 850, 875, 900, 925, 950, 975, 1000])

            pressure = 100 * pressure_hPa

            temp_interp1d = interpolate.interp1d(interp4d_height, interp4d_temp)
            press_interp1d = interpolate.interp1d(interp4d_height, pressure)
            vw_x_interp1d = interpolate.interp1d(interp4d_height, interp4d_vw_x)
            vw_y_interp1d = interpolate.interp1d(interp4d_height, interp4d_vw_y)
            vw_z_interp1d = interpolate.interp1d(interp4d_height, interp4d_vw_z)

        except IndexError:
            print("Error: Check time range of ERA5 data!")
    else:
        pass

    tcc_pre = np.ma.dot(w2, ERAtcc[t_pre_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    tcc_post = np.ma.dot(w2, ERAtcc[t_post_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    tcc = (tcc_post - tcc_pre) * (t_epoch - t_pre) + tcc_pre

    if isinstance(tcc, float) != True:
        print("WARNING: Corrupt ERA5 Data for parameter \"tcc\"!")
        print("Assuming simplified value for parameter \"tcc\".")
        tcc = cc

    cbh_pre = np.ma.dot(w2, ERAcbh[t_pre_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    cbh_post = np.ma.dot(w2, ERAcbh[t_post_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    cbh = (cbh_post - cbh_pre) * (t_epoch - t_pre) + cbh_pre

    if isinstance(tcc, float) != True:
        print("WARNING: Corrupt ERA5 Data for parameter \"cbh\"!")
        print("Assuming simplified value for parameter \"cbh\".")
        cbh = 2000

    lcc_pre = np.ma.dot(w2, ERAlcc[t_pre_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    lcc_post = np.ma.dot(w2, ERAlcc[t_post_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    lcc = (lcc_post - lcc_pre) * (t_epoch - t_pre) + lcc_pre

    if isinstance(lcc, float) != True:
        print("WARNING: Corrupt ERA5 Data for parameter \"lcc\"!")
        print("Assuming simplified value for parameter \"lcc\".")
        lcc = cc/3

    mcc_pre = np.ma.dot(w2, ERAmcc[t_pre_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    mcc_post = np.ma.dot(w2, ERAmcc[t_post_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    mcc = (mcc_post - mcc_pre) * (t_epoch - t_pre) + mcc_pre

    if isinstance(mcc, float) != True:
        print("WARNING: Corrupt ERA5 Data for parameter \"mcc\"!")
        print("Assuming simplified value for parameter \"mcc\".")
        mcc = cc/3

    hcc_pre = np.ma.dot(w2, ERAhcc[t_pre_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    hcc_post = np.ma.dot(w2, ERAhcc[t_post_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    hcc = (hcc_post - hcc_pre) * (t_epoch - t_pre) + hcc_pre

    if isinstance(hcc, float) != True:
        print("WARNING: Corrupt ERA5 Data for parameter \"hcc\"!")
        print("Assuming simplified value for parameter \"hcc\".")
        hcc = cc/3

    ssr_pre = np.ma.dot(w2, ERAssr[t_pre_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    ssr_post = np.ma.dot(w2, ERAssr[t_post_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    ssr = ((ssr_post - ssr_pre) * (t_epoch - t_pre) + ssr_pre) / 3600

    strn_pre = np.ma.dot(w2, ERAstrn[t_pre_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    strn_post = np.ma.dot(w2, ERAstrn[t_post_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    strn = ((strn_post - strn_pre) * (t_epoch - t_pre) + strn_pre) / 3600

    if isinstance(strn, float) != True:
        print("WARNING: Corrupt ERA5 Data for parameter \"strn\"!")
        print("Assuming simplified value for parameter \"strn\".")
        strn = 0

    skt_pre = np.ma.dot(w2, ERAskt[t_pre_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    skt_post = np.ma.dot(w2, ERAskt[t_post_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    skt = ((skt_post - skt_pre) * (t_epoch - t_pre) + skt_pre)

    if isinstance(skt, float) != True:
        print("WARNING: Corrupt ERA5 Data for parameter \"skt\"!")
        print("Assuming simplified value for parameter \"skt\".")
        skt = T_ground

    strd_pre = np.ma.dot(w2, ERAstrd[t_pre_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    strd_post = np.ma.dot(w2, ERAstrd[t_post_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    strd = ((strd_post - strd_pre) * (t_epoch - t_pre) + strd_pre) / 3600

    if isinstance(strd, float) != True:
        print("WARNING: Corrupt ERA5 Data for parameter \"strd\"!")
        print("Assuming simplified value for parameter \"strd\".")
        strd = epsilon_ground * sigma * T_ground ** 4

    strdc_pre = np.ma.dot(w2, ERAstrdc[t_pre_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    strdc_post = np.ma.dot(w2, ERAstrdc[t_post_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    strdc = ((strdc_post - strdc_pre) * (t_epoch - t_pre) + strdc_pre) / 3600

    if isinstance(strdc, float) != True:
        print("WARNING: Corrupt ERA5 Data for parameter \"strdc\"!")
        print("Assuming simplified value for parameter \"strdc\".")
        strdc = epsilon_ground * sigma * T_ground ** 4

    ssrd_pre = np.ma.dot(w2, ERAssrd[t_pre_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    ssrd_post = np.ma.dot(w2, ERAssrd[t_post_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    ssrd = ((ssrd_post - ssrd_pre) * (t_epoch - t_pre) + ssrd_pre) / 3600

    if isinstance(ssrd, float) != True:
        print("WARNING: Corrupt ERA5 Data for parameter \"ssrd\"!")
        print("Assuming simplified value for parameter \"ssrd\".")
        ssrd = 1
        ssr = 1 - Albedo

    if isinstance(ssr, float) != True:
        print("WARNING: Corrupt ERA5 Data for parameter \"ssr\"!")
        print("Assuming simplified value for parameter \"ssr\".")
        ssrd = 1
        ssr = 1 - Albedo

    tsr_pre = np.ma.dot(w2, ERAtsr[t_pre_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    tsr_post = np.ma.dot(w2, ERAtsr[t_post_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    tsr = ((tsr_post - tsr_pre) * (t_epoch - t_pre) + tsr_pre) / 3600

    tisr_pre = np.ma.dot(w2, ERAtisr[t_pre_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    tisr_post = np.ma.dot(w2, ERAtisr[t_post_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    tisr = ((tisr_post - tisr_pre) * (t_epoch - t_pre) + tisr_pre) / 3600

    if isinstance(tisr, float) != True:
        print("WARNING: Corrupt ERA5 Data for parameter \"tisr\"!")
        print("Assuming simplified value for parameter \"tisr\".")
        AZ, ELV = sun_angles_analytical(lat, lon, h, utc)
        MA = (357.52911 + 0.98560028 * (utc.jd - 2451545)) % 360  # in degree, reference: see folder "literature"
        TA = MA + 2 * e * np.sin(np.deg2rad(MA)) + 5 / 4 * e ** 2 * np.sin(np.deg2rad(2 * MA))
        I_Sun = 1367.5 * ((1 + e * np.cos(np.deg2rad(TA))) / (1 - e ** 2)) ** 2
        tisr = I_Sun * np.sin(np.deg2rad(ELV))

    if isinstance(tsr, float) != True:
        print("WARNING: Corrupt ERA5 Data for parameter \"tsr\"!")
        print("Assuming simplified value for parameter \"tsr\".")
        tsr = (1 - Albedo) * tisr

    ttr_pre = np.ma.dot(w2, ERAttr[t_pre_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    ttr_post = np.ma.dot(w2, ERAttr[t_post_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    ttr = ((ttr_post - ttr_pre) * (t_epoch - t_pre) + ttr_pre) / 3600

    p0_pre = np.ma.dot(w2, ERAsp[t_pre_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    p0_post = np.ma.dot(w2, ERAsp[t_post_ind, lat_ind2, lon_ind2]) / np.sum(w2)
    p0 = (p0_post - p0_pre) * (t_epoch - t_pre) + p0_pre

    if isinstance(p0, float) != True:
        print("WARNING: Corrupt ERA5 Data for parameter \"sp\"!")
        print("Assuming simplified value for parameter \"sp\".")
        p0 = 101325.0

    if isinstance(ttr, float) != True:
        print("WARNING: Corrupt ERA5 Data for parameter \"ttr\"!")
        print("Assuming simplified value for parameter \"ttr\".")
        AZ, ELV = sun_angles_analytical(lat, lon, h, utc)
        tau_atm, tau_atmIR = tau(ELV, h, p_air, p0)
        HalfConeAngle = np.arcsin(R_E / (R_E + h))
        ViewFactor = (1 - np.cos(HalfConeAngle)) / 2
        ttr = epsilon_ground * sigma * T_ground ** 4 * tau_atmIR * ViewFactor * 2

    if h > interp4d_height[0]:
        p_air = press_interp1d(interp4d_height[0])
        T_air = temp_interp1d(interp4d_height[0])
        u = vw_x_interp1d(interp4d_height[0])
        v = vw_y_interp1d(interp4d_height[0])
        w = -1 / grav(lat, h) * vw_z_interp1d(interp4d_height[0]) * R_air * T_air / p_air
    elif h < interp4d_height[-1]:
        p_air = press_interp1d(interp4d_height[-1])
        T_air = temp_interp1d(interp4d_height[-1])
        u = vw_x_interp1d(interp4d_height[-1])
        v = vw_y_interp1d(interp4d_height[-1])
        w = -1 / grav(lat, h) * vw_z_interp1d(interp4d_height[-1]) * R_air * T_air / p_air
    else:
        p_air = press_interp1d(h)
        T_air = temp_interp1d(h)
        u = vw_x_interp1d(h)
        v = vw_y_interp1d(h)
        w = -1 / grav(lat, h) * vw_z_interp1d(h) * R_air * T_air / p_air

    rho_air = p_air / (R_air * T_air)

    return p_air, p0, T_air, rho_air, u, v, w, cbh, tcc, lcc, mcc, hcc, ssr, strn, strd, strdc, ssrd, tsr, ttr, tisr, skt


deltaT_ERA = (Time(start_utc).jd - Time('1900-01-01 00:00:00.0').jd) * 24.000000
p_air0, p00, T_air0, rho_air0, u0, v0, w0, cbh0, tcc0, lcc0, mcc0, hcc0, ssr0, strn0, strd0, strdc0, ssrd0, tsr0, ttr0, tisr0, skt0 = ERA5Data(
    start_lon, start_lat, start_height, 0, deltaT_ERA)

y0 = [
    start_lon,     # start longitude [deg]
    start_lat,     # start latitude [deg]
    start_height,  # start altitude [m]
    0,             # initial v_x [m/s]
    0,             # initial v_y [m/s]
    0,             # initial v_z [m/s]
    T_air0,        # initial gas temperature [K] = initial air temperature [K]
    T_air0,        # initial film temperature [K] = initial air temperature [K]
    m_gas_init,    # initial lifting gas mass [kg]
    0,             # initial factor c2 [-]
    m_bal_init     # initial ballast mass [kg]
]


def model(t, y, m_pl, m_film, c_virt):
    utc = Time(start_utc) + t * unit.second
    lon = y[0]     # 1
    lat = y[1]     # 2
    h = y[2]       # 3
    v_x = y[3]     # 4
    v_y = y[4]     # 5
    v_z = y[5]     # 6
    T_gas = y[6]   # 7
    T_film = y[7]  # 8
    m_gas = y[8]   # 9
    c2 = y[9]      # 10
    m_bal = y[10]  # 11

    if (lon % 360) > 180:        # convert longitude to value in standard interval [-180, 180]
        lon = (lon % 360) - 360
    else:
        lon = (lon % 360)

    if lat > 90:                 # set plausible limits for latitudes
        lat = 90
    elif lat < -90:
        lat = -90
    else:
        lat = lat

    r_lon, r_lat = radii(lat, h)  # calculate radii for velocity conversion between cartesian and Earth reference frame

    deltaT_ERA = (Time(start_utc).jd - Time('1900-01-01 00:00:00.0').jd) * 24.000000  # conversion to ERA5 time format

    try:
        p_air, p0, T_air, rho_air, u, v, w, cbh, tcc, lcc, mcc, hcc, ssr, strn, strd, strdc, ssrd, tsr, ttr, tisr, skt = ERA5Data(lon, lat, h, t, deltaT_ERA)
    except:
        # in case of solver (temporarily) exceeding interpolation area (with subsequent correction)
        # or permanent drift out of interpolation area
        if h >= 30000 or (np.abs(lat - start_lat) <= 10.0 and np.abs(lon - start_lon) <= 10.0):
            print("solver exceeds definition area")
            p_air, p0, T_air, rho_air, u, v, w, cbh, tcc, lcc, mcc, hcc, ssr, strn, strd, strdc, ssrd, tsr, ttr, tisr, skt = 0, 101325, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
        else:
            print("instable trajectory!")
            p_air, p0, T_air, rho_air, u, v, w, cbh, tcc, lcc, mcc, hcc, ssr, strn, strd, strdc, ssrd, tsr, ttr, tisr, skt = 0, 101325, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0

    p_gas = p_air

    h_valve = 1.034 * V_design ** (1 / 3)
    h_duct = 0.47 * h_valve

    v_relx = u - v_x  # relative wind velocity x-dir (balloon frame)
    v_rely = v - v_y  # relative wind velocity y-dir (balloon frame)
    v_relz = w - v_z  # relative wind velocity z-dir (balloon frame)

    v_rel = np.sqrt(v_relx ** 2 + v_rely ** 2 + v_relz ** 2)  # total relative wind velocity (balloon frame)

    alpha = np.arcsin(v_relz / v_rel)  # "angle of attack": angle between longitudinal axis and rel. wind (in [rad])

    rho_gas = p_gas / (R_gas * T_gas)  # calculate gas density through *ideal* gas equation

    dP_valve = grav(lat, h) * (rho_air - rho_gas) * h_valve
    dP_duct = grav(lat, h) * (rho_air - rho_gas) * h_duct

    if m_gas > m_gas_init:  # limit gas mass to plausible value
        m_gas = m_gas_init  #
    elif m_gas < 0:
        m_gas = 0
    else:
        pass

    V_b = m_gas / rho_gas  # calculate balloon volume from current gas mass and gas density

    if V_b > V_design:
        c_duct = c_ducts
    elif V_b < 0:
        c_duct = 0
        V_b = 0
    else:
        c_duct = 0

    if ballasting(utc) == True:
        if m_bal >= 0:
            mdot = m_baldot
        else:
            mdot = 0
    else:
        mdot = 0

    if valving(utc) == True:  # opening valve process
        if c2 == 0:
            c2 = 1.0
            c2dot = 0
        elif c_valve < c2 <= 1.0:
            c2dot = (c_valve - 1.0) / t_open
        else:
            c2dot = 0
            c2 = c_valve

    if valving(utc) == False:  # closing valve process
        if c2 == 0:
            c2dot = 0
        elif c_valve <= c2 < 1.0:
            c2dot = (1.0 - c_valve) / t_close
        else:
            c2dot = 0
            c2 = 0

    m_gross = m_pl + m_film + m_bal
    m_tot = m_pl + m_film + m_gas
    m_virt = m_tot + c_virt * rho_air * V_b

    d_b = 1.383 * V_b ** (1 / 3)  # calculate diameter of balloon from its volume
    L_goreB = 1.914 * V_b ** (1 / 3)
    A_surf = 4.94 * V_b ** (2 / 3)
    A_surf1 = 4.94 * V_design ** (2 / 3) * (1 - np.cos(np.pi * L_goreB / L_goreDesign))
    A_eff = 0.65 * A_surf + 0.35 * A_surf1
    A_top = np.pi / 4 * d_b ** 2

    AZ, ELV = sun_angles_analytical(lat, lon, h, utc)

    A_proj = A_top * (0.9125 + 0.0875 * np.cos(np.pi - 2 * np.deg2rad(ELV)))  # projected area for sun radiation
    A_drag = A_top * (0.9125 + 0.0875 * np.cos(np.pi - 2 * alpha))            # projected area for drag

    # CALCULATIONS FOR THERMAL MODEL

    tau_atm, tau_atmIR = tau(ELV, h, p_air, p0)

    tau_atm0, tau_atmIR0 = tau(ELV, 0, p0, p0)

    MA = (357.52911 + 0.98560028 * (utc.jd - 2451545)) % 360  # in degree, reference: see folder "literature"
    TA = MA + 2 * e * np.sin(np.deg2rad(MA)) + 5 / 4 * e ** 2 * np.sin(np.deg2rad(2 * MA))
    I_Sun = 1367.5 * ((1 + e * np.cos(np.deg2rad(TA))) / (1 - e ** 2)) ** 2
    I_SunZ = I_Sun * tau_atm
    I_Sun0 = I_Sun * tau_atm0

    HalfConeAngle = np.arcsin(R_E / (R_E + h))
    ViewFactor = (1 - np.cos(HalfConeAngle)) / 2

    if simple == True:

        q_IREarth = epsilon_ground * sigma * T_ground ** 4 * tau_atmIR

        if ELV <= 0:
            q_Albedo = 0
        else:
            q_Albedo = Albedo * I_Sun * np.sin(np.deg2rad(ELV))

        Q_Albedo = alpha_VIS * A_surf * q_Albedo * ViewFactor * (1 + tau_VIS / (1 - r_VIS))
        Q_IREarth = alpha_IR * A_surf * q_IREarth * ViewFactor * (1 + tau_IR / (1 - r_IR))

        q_sun = I_SunZ

    else:

        if tcc <= 0.01:

            q_IREarth1 = alpha_IR * A_surf * (strd - strn) * tau_atmIR * ViewFactor * (1 + tau_VIS / (1 - r_VIS))
            q_IREarth2 = alpha_IR * A_surf * np.abs(ttr) * 0.5 * (1 + tau_IR / (1 - r_IR))
            # q_IREarth2 = alpha_IR * A_surf * (0.04321906 * np.abs(ttr) + 84.67820281) * (1 + tau_IR / (1 - r_IR))

            if h > 40000:
                Q_IREarth = q_IREarth2
            else:
                Q_IREarth = (q_IREarth2 - q_IREarth1) / 40000 * h + q_IREarth1

            if ELV <= 0:
                q_Albedo1 = 0
                q_Albedo2 = 0
            else:
                if ssrd == 0:
                    q_Albedo1 = 0
                else:
                    q_Albedo1 = alpha_VIS * A_surf * (1 - ssr / ssrd) * I_Sun0 * tau_atmIR * np.sin(
                        np.deg2rad(ELV)) * ViewFactor * (1 + tau_VIS / (1 - r_VIS))  # !
                if tisr == 0:
                    q_Albedo2 = 0
                else:
                    q_Albedo2 = alpha_VIS * A_surf * (1 - tsr / tisr) * I_Sun * np.sin(np.deg2rad(ELV)) * 0.5 * (
                            1 + tau_VIS / (1 - r_VIS))  # !
            if h > 40000:
                Q_Albedo = q_Albedo2
            else:
                Q_Albedo = (q_Albedo2 - q_Albedo1) / 40000 * h + q_Albedo1

            q_sun = I_SunZ

        else:
            q_IRground_bc = (strd - strn) + (strd - strdc)
            q_IREarth_bc = alpha_IR * A_surf * q_IRground_bc * tau_atmIR * ViewFactor * (1 + tau_VIS / (1 - r_VIS))
            q_sun_bc = I_SunZ * (1 - tcc)
            q_Albedo_bc = alpha_VIS * A_surf * (1 - ssr / ssrd) * I_Sun0 * tau_atmIR * np.sin(
                np.deg2rad(ELV)) * ViewFactor * (1 + tau_VIS / (1 - r_VIS))

            # q_IREarth_ac = alpha_IR * A_surf * (0.04321906 * np.abs(ttr) + 84.67820281) * (1 + tau_IR / (1 - r_IR))
            q_IREarth_ac = alpha_IR * A_surf * np.abs(ttr) * 0.5 * (1 + tau_VIS / (1 - r_VIS))
            q_sun_ac = I_SunZ
            q_Albedo_ac = alpha_VIS * A_surf * (1 - tsr / tisr) * I_Sun * np.sin(np.deg2rad(ELV)) * ViewFactor * (
                    1 + tau_VIS / (1 - r_VIS))

            if h <= cbh:                    # "below clouds"
                Q_IREarth = q_IREarth_bc
                q_sun = q_sun_bc
                if ELV <= 0:
                    Q_Albedo = 0
                else:
                    Q_Albedo = q_Albedo_bc
            elif h >= 12000:                # "above clouds"
                Q_IREarth = q_IREarth_ac
                q_sun = q_sun_ac
                if ELV <= 0:
                    Q_Albedo = 0
                else:
                    Q_Albedo = q_Albedo_ac
            elif h >= 6000:
                if hcc >= 0.01:
                    Q_IREarth = ((h - 6000) / 6000 * hcc + mcc + lcc) / (hcc + mcc + lcc) * (
                            q_IREarth_ac - q_IREarth_bc) + q_IREarth_bc
                    q_sun = ((h - 6000) / 6000 * hcc + mcc + lcc) / (hcc + mcc + lcc) * (q_sun_ac - q_sun_bc) + q_sun_bc
                    if ELV <= 0:
                        Q_Albedo = 0
                    else:
                        Q_Albedo = ((h - 6000) / 6000 * hcc + mcc + lcc) / (hcc + mcc + lcc) * (
                                q_Albedo_ac - q_Albedo_bc) + q_Albedo_bc
                else:
                    Q_IREarth = q_IREarth_ac
                    q_sun = q_sun_ac
                    if ELV <= 0:
                        Q_Albedo = 0
                    else:
                        Q_Albedo = q_Albedo_ac
            elif h >= 2000:
                if mcc > 0.01 or hcc > 0.01:
                    Q_IREarth = ((h - 2000) / 4000 * mcc + lcc) / (hcc + mcc + lcc) * (
                            q_IREarth_ac - q_IREarth_bc) + q_IREarth_bc
                    q_sun = ((h - 2000) / 4000 * mcc + lcc) / (hcc + mcc + lcc) * (q_sun_ac - q_sun_bc) + q_sun_bc
                    if ELV <= 0:
                        Q_Albedo = 0
                    else:
                        Q_Albedo = ((h - 2000) / 4000 * mcc + lcc) / (hcc + mcc + lcc) * (
                                q_Albedo_ac - q_Albedo_bc) + q_Albedo_bc
                else:
                    Q_IREarth = q_IREarth_ac
                    q_sun = q_sun_ac
                    if ELV <= 0:
                        Q_Albedo = 0
                    else:
                        Q_Albedo = q_Albedo_ac
            else:
                Q_IREarth = (h / 2000 * lcc) / (hcc + mcc + lcc) * (q_IREarth_ac - q_IREarth_bc) + q_IREarth_bc
                q_sun = (h / 2000 * lcc) / (hcc + mcc + lcc) * (q_sun_ac - q_sun_bc) + q_sun_bc

                if ELV <= 0:
                    Q_Albedo = 0
                else:
                    Q_Albedo = (h / 2000 * lcc) / (hcc + mcc + lcc) * (q_Albedo_ac - q_Albedo_bc) + q_Albedo_bc

    my_air = (1.458 * 10 ** -6 * T_air ** 1.5) / (T_air + 110.4)
    k_air = 0.0241 * (T_air / 273.15) ** 0.9
    k_gas = 0.144 * (T_gas / 273.15) ** 0.7
    Pr_air = 0.804 - 3.25 * 10 ** (-4) * T_air
    Pr_gas = 0.729 - 1.6 * 10 ** (-4) * T_gas
    Gr_air = (rho_air ** 2 * grav(lat, h) * np.abs(T_film - T_air) * d_b ** 3) / (T_air * my_air ** 2)
    Nu_air = 2 + 0.45 * (Gr_air * Pr_air) ** 0.25
    HC_free = Nu_air * k_air / d_b
    Re = np.abs(v_rel) * d_b * rho_air / my_air
    Fr = np.abs(v_rel) / np.sqrt(grav(lat, h) * d_b)

    HC_forced = k_air / d_b * (2 + 0.41 * Re ** 0.55)
    HC_internal = 0.13 * k_gas * (
            (rho_gas ** 2 * grav(lat, h) * np.abs(T_film - T_gas) * Pr_gas) / (T_gas * my_air ** 2)) ** (
                          1 / 3)
    HC_external = np.maximum(HC_free, HC_forced)

    Q_Sun = alpha_VIS * A_proj * q_sun * (1 + tau_VIS / (1 - r_VIS))
    Q_IRFilm = sigma * epsilon * alpha_IR * A_surf * T_film ** 4 * 1 / (1 - r_IR)
    Q_IRout = sigma * epsilon * A_surf * T_film ** 4 * (1 + tau_IR / (1 - r_IR))
    Q_ConvExt = HC_external * A_eff * (T_air - T_film)
    Q_ConvInt = HC_internal * A_eff * (T_film - T_gas)

    if simple == True:
        c_d = 0.47   # sphere
    else:
        c_d = cd_PalumboHigh(Fr, Re, A_top)  # Ref. xx

    #c_d = 0.47  # cd_PalumboHigh(Fr, Re, A_top)  #cd_sphere(Re)  #0.47  # cd_Palumbo(Fr, Re, A_top)  # 0.8  #cd_sphere(Re)  # cd_palumbo(Fr, Re, A_top)  # cd_sphere(Re) / 0.8 / 0.47

    D = drag(c_d, rho_air, A_drag, v_rel)  # calculate drag force

    if v_rel == 0:
        Drag_x, Drag_y, Drag_z = 0, 0, 0
    else:
        Drag_x, Drag_y, Drag_z = D * v_relx / v_rel, D * v_rely / v_rel, D * v_relz / v_rel

    F = grav(lat, h) * V_b * (rho_air - rho_gas) - grav(lat, h) * m_gross + Drag_z  # gross inflation - weight + drag

    a_x, a_y, a_z = Drag_x / m_virt, Drag_y / m_virt, F / m_virt

    eqn1 = np.rad2deg(y[3] / r_lon)
    eqn2 = np.rad2deg(y[4] / r_lat)
    eqn3 = v_z
    eqn4 = a_x
    eqn5 = a_y
    eqn6 = a_z
    eqn7 = Q_ConvInt / (gamma * c_v * m_gas) - (gamma - 1) / gamma * (rho_air * grav(lat, h)) / (rho_gas * R_gas) * v_z
    eqn8 = (Q_Sun + Q_Albedo + Q_IREarth + Q_IRFilm + Q_ConvExt - Q_ConvInt - Q_IRout) / (c_f * m_film)
    eqn9 = -(A_ducts * c_duct * np.sqrt(np.abs(2 * dP_duct * rho_gas))) - (A_valve * c2 * np.sqrt(np.abs(2 * dP_valve * rho_gas)))
    eqn10 = c2dot
    if m_bal > 0:
        eqn11 = -mdot
    else:
        eqn11 = 0

    return [eqn1, eqn2, eqn3, eqn4, eqn5, eqn6, eqn7, eqn8, eqn9, eqn10, eqn11]


# DEFINITION OF EVENTS FOR SOLVER

def at_ground(t, y, m_pl, m_film, c_virt):
    return y[2]


def above_float(t, y, m_pl, m_film, c_virt):
    return 45000 - y[2]


def below_float(t, y, m_pl, m_film, c_virt):
    return y[2] - 30050


hit_ground = lambda t, x: at_ground(t, x, m_pl, m_film, c_virt)
hit_ground.terminal = True
hit_ground.direction = -1
excess_ascend = lambda t, x: above_float(t, x, m_pl, m_film, c_virt)
excess_ascend.terminal = True
excess_ascend.direction = -1
instable = lambda t, x: below_float(t, x, m_pl, m_film, c_virt)
instable.terminal = True
instable.direction = -1

t0 = 2.0
tf = 30000

print("Beginning simulation")

sol = solve_ivp(fun=lambda t, x: model(t, x, m_pl, m_film, c_virt), t_span=[t0, tf], y0=y0, method='BDF', events=[hit_ground, excess_ascend, instable])

# tnew = np.linspace(0, sol.t[-1], len(Vol_list))

print(sol.message)

print(datetime.now() - starttime)

plt.plot(sol.t, sol.y[2, :], 'k--', label='Simulation')
plt.plot(comp_time, comp_height, 'r-', label='PoGo Flight Test')
plt.legend()
plt.title('high factor')
plt.xlabel('time in s')
plt.ylabel('Balloon Altitude in m')
plt.show()

plt.clf()
# ax = plt.axes(projection=ccrs.Mollweide())
ax = plt.axes(projection=ccrs.AzimuthalEquidistant(central_latitude=90))
ax.coastlines()
ax.gridlines()
ax.stock_img()
ax.set_extent([-120, 30, 60, 80], crs=ccrs.PlateCarree())

plt.plot(start_lon, start_lat, 'rx', transform=ccrs.Geodetic())
plt.plot(sol.y[0, :], sol.y[1, :], 'k--', transform=ccrs.Geodetic())
plt.plot(comp_lon, comp_lat, 'r-', transform=ccrs.Geodetic())
# plt.xticks()
# figname = "LatLon_%.2f_%.2f.png" % (Albedo, epsilon_ground)
# plt.savefig(os.path.join(rootdir, figname))
plt.show()