plot_cross_section.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Aug 25 09:50:51 2022

@author: lunelt
"""

import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import numpy as np
import xarray as xr
import tools
# import metpy.calc as mcalc
import global_variables as gv
import matplotlib as mpl
from matplotlib import ticker
from shapely.geometry import Point, LineString

########## Independant parameters ###############

# model = '1_planier_0105/02_1km'
# model = '1_planier_0105/05_1km_noEDKF_3DTKE'
# model = '1_planier_0105/06_1km_noEDKF'
# model = '1_planier_0105/07_1km_3DTKE'
# model = '1_planier_0105/021_200m'
# model = '1_planier_0105/0221_40m'
# model = '1_planier_0105/0211_40m'
model = '2_planier_0122/04_1km_COARE'
# model = '2_planier_0122/0411_40m'
# model = '2_planier_0122/06_1km_UPDRAFTDIAG'
# model = '3_planier_0127/04_1km_COARE'
# model = '3_planier_0127/0411_40m'
# model = 'std_d1'

global_simu_folder=gv.global_simu_folder
output_type='diag'
# Datetime
wanted_date = '20230122-0700'

varname_colormap = 'TKET'
varname_contourmap = 'THTV'

# Surface variable to show below the section
subplot_type = 'surface_var'  #'surface_var', 'distance'
# only useful if subplot_type=='surface_var' (ex: WG2_ISBA, H, LE, etc)
surf_var = 'RI'   
surf_var_label = 'RI'
vmin_surf_var, vmax_surf_var = -0.025, 0.025
 
# Set type of wind representation: 'verti_proj' or 'horiz'
wind_visu = 'verti_proj'

# altitude ASL or height AGL: 'asl' or 'agl'
alti_type = 'asl'
# maximum level (height AGL) to plot
toplevel = 900

# where to place the cross section
nb_points_beyond = 3
site_start = 'carry'
site_end = 'planier'

# sites different from start and end, to project on the x-axis
# sites_to_project = ['elsplans', 'serra_tallat', 'coll_lilla']  
sites_to_project = [
    # 'carry',
    ]  

# Arrow/barbs esthetics:
skip_barbs_x = 2
skip_barbs_y = 10    #if 1: 1barb/10m, if 5: 1barb/50m, etc
arrow_size = 0.7  #works for arrow and barbs
barb_size_option = 'standard'  # 'weak_winds' or 'standard'

# Save the figure
figsize = (9,4)
save_plot = True
save_folder = f'./figures/cross_sections/{model}/section_{site_start}_{site_end}/{wind_visu}/{varname_colormap}_{varname_contourmap}/'

plt.rcParams.update({'font.size': 11})
###########################################

vmin = gv.minmax_dict[varname_colormap]['vmin']
vmax = gv.minmax_dict[varname_colormap]['vmax']
vstep = gv.minmax_dict[varname_colormap]['vstep']
colormap = gv.minmax_dict[varname_colormap]['colormap']

# -- modify by hand if necessary:
# vmax = 5

# -- Colorscale linear
# norm_cm=mpl.colors.Normalize(vmin=vmin, vmax=vmax)
# levels = np.linspace(vmin, vmax, 12+1)
levels = np.arange(vmin, vmax, vstep)
# levels = np.linspace(vmin, vmax, vmax-vmin+1)  # to have 1 unit per color variation
# -- Colorscale log
# norm_cm=mpl.colors.LogNorm(vmin=vmin+0.01, vmax=vmax)  # for TKE
# levels = 20
# -- Colorscale unknown
if None in [vmin, vmax]:
    levels = 10

vmin_contour = gv.minmax_dict[varname_contourmap]['vmin']
vmax_contour = gv.minmax_dict[varname_contourmap]['vmax']
vstep_contour = gv.minmax_dict[varname_contourmap]['vstep']
levels_contour = np.arange(vmin_contour, vmax_contour, vstep_contour)
#vmin_contour, vmax_contour = None, None

barb_size_increments = gv.barb_size_increments
barb_size_description = gv.barb_size_description

# Filename to open
filename = tools.get_simu_filepath(model, wanted_date, 
                                   output_type=output_type,
                                   global_simu_folder=global_simu_folder)
# OR:
#filename = '/home/lunelt/Data/temp_outputs/LIAIS.1.SEG03.009.nc'

ds = xr.open_dataset(filename,)

# Coordinates of starting and ending points of the cross-section:
end = (gv.whole[site_end]['lat'], gv.whole[site_end]['lon'])
start = (gv.whole[site_start]['lat'], gv.whole[site_start]['lon'])
# OR:
# end = (44,1)
# start = (43,1.5)

# Check order of start and end sites
if gv.whole[site_start]['lon'] > gv.whole[site_end]['lon']:
    raise ValueError("site_start must be west of site_end")


# pre-processing of data
ds_sub = tools.subset_ds(ds, 
                         lat_range = [start[0], end[0]], 
                         lon_range = [start[1], end[1]],
                         nb_indices_exterior=nb_points_beyond) 

# ds_subcen = tools.flux_pt_to_mass_pt(ds_sub, only_basic_vars=False)
ds_subcen = tools.flux_pt_to_mass_pt(
    ds_sub, only_basic_vars=True,
    basic_vars=['UT', 'VT', 'WT', 
                'UMME', 'VMME', 'WMME',
                'U2ME', 'V2ME', 'W2ME',
                'MF_FRAC_UP', 'MF_W_UP', 'MF_U_UP', 'MF_V_UP',
                ],
    )

#%% DIAG
# Computation of other diagnostic variable
# ds_subcen['DIV'] = mcalc.divergence(ds_subcen['UT'], ds_subcen['VT'])
ds_subcen['WS'], ds_subcen['WD'] = tools.calc_ws_wd(ds_subcen['UT'], ds_subcen['VT'])
try:
    ds_subcen['THTV'] = ds_subcen['THT']*(1 + 0.61*ds_subcen['RVT'] - (ds_subcen['MRR']+ds_subcen['MRC'])/1000)
except KeyError:
    ds_subcen['THTV'] = ds_subcen['THT']*(1 + 0.61*ds_subcen['RVT'])
# ds_subcen['MSLP3D'] = tools.calc_mslp(ds)
try:
    ds_subcen['RES_TKE'] = 0.5* (ds_subcen['U2ME'] + ds_subcen['V2ME'] + ds_subcen['W2ME'])
    ds_subcen['SBG_TKE'] = ds_subcen['TKEMME']
    ds_subcen['TOT_TKE'] = ds_subcen['SBG_TKE'] + ds_subcen['RES_TKE']
    ds_subcen['RES_TKE_FRAC'] = ds_subcen['RES_TKE']/ds_subcen['TOT_TKE']
    ds_subcen['TOT_TI'] = np.sqrt(1.5*ds_subcen['TOT_TKE'])/ds_subcen['WS']
except:
    pass

try:
    # ds_subcen['TKE_MF_UP'] = \
    #     ds_subcen['MF_FRAC_UP'] / (1 - ds_subcen['MF_FRAC_UP']) * \
    #     (ds_subcen['MF_W_UP'] - ds_subcen['WT'])**2
    ds_subcen['TKE_MF_UP'] = \
        0.5*(ds_subcen['MF_FRAC_UP'] / (1 - ds_subcen['MF_FRAC_UP'])) * \
        ((ds_subcen['MF_W_UP'] - ds_subcen['WT'])**2 + \
         (ds_subcen['MF_U_UP'] - ds_subcen['UT'])**2 + \
         (ds_subcen['MF_V_UP'] - ds_subcen['VT'])**2)
    ds_subcen['TKE_EQ'] = ds_subcen['TKET'] + ds_subcen['TKE_MF_UP']
except:
    pass

data_reduced = ds_subcen[[
    'UT', 'VT', 'WT', 
    'ZS', 
    # 'ALT',
#    'TEMP', 'PRES', 'HBLTOP',
#    'DENS', 'DIV', 'WS', 'WD',
    varname_colormap, varname_contourmap, surf_var]]


data = data_reduced
#data = data_redsub


#%% CREATE SECTION LINE

line = tools.get_line_coords(data, start, end, 
                             nb_indices_exterior=nb_points_beyond)
ni_range = line['ni_range']
nj_range = line['nj_range']
slope = line['slope']

if slope == 'vertical':
    angle = np.pi/2
else:
    angle = np.arctan(slope)
    
data['WPROJ'] = tools.windvec_verti_proj(data['UT'], data['VT'], 
                                         data.level, angle)

#%% INTERPOLATION

section = []
abscisse_coords = []
abscisse_sites = {}

#get total maximum height of relief on domain
max_ZS = data['ZS'].max()
if alti_type == 'asl':
    level_range = np.arange(10, toplevel+max_ZS, 10)
else:
    level_range = np.arange(10, toplevel, 10)
    
print('section interpolation on {0} points (~1sec/pt)'.format(len(ni_range)))
for i, ni in enumerate(ni_range):
    nj=nj_range[i]
    #interpolation of all variables on ni_range
    profile = data.interp(ni=ni, 
                          nj=nj, 
                          level=level_range).expand_dims({'i_sect':[i]})
    section.append(profile)
    
    #store values of lat-lon for the horiz axis
    lat = np.round(profile.latitude.values, decimals=3)
    lon = np.round(profile.longitude.values, decimals=3)
    latlon = str(lat) + '\n' + str(lon)
    abscisse_coords.append(latlon)
    
    #Store values of i and name of site in dict for horiz axis
    if slope == 'vertical':
        if nj == line['nj_start']:
            abscisse_sites[i] = site_start
        elif nj == line['nj_end']:
            abscisse_sites[i] = site_end
    else:
        if ni == line['ni_start']:
            abscisse_sites[i] = site_start
        elif ni == line['ni_end']:
            abscisse_sites[i] = site_end

#concatenation of all profile in order to create the 2D section dataset
section_ds = xr.concat(section, dim="i_sect")


#%% PLOT
# create figure
fig, ax = plt.subplots(2, figsize=figsize,
                       gridspec_kw={'height_ratios': [20, 1]})

## --- Subplot of section, i.e. the main plot ----
#get maximum height of relief in cross-section
max_ZS = section_ds['ZS'].max()

# remove top layers of troposphere
section_ds = section_ds.where(section_ds.level<(level_range.max()), drop=True)

## --- Adapt to alti_type ------
#create grid mesh (eq. to X)
X = np.meshgrid(section_ds.i_sect, section_ds.level)[0]
Xmesh = xr.DataArray(X, dims=['level', 'i_sect'])
#create alti mesh (eq. to Y)
if alti_type == 'asl': 
    #compute altitude ASL from height AGL, and transpose (eq. Y)
    # alti = section_ds['ALT']
    #old: (issue)
    alti = section_ds.ZS[:, 0] + section_ds.level
    alti = alti.T
    #for plot
    ylabel = 'altitude a.s.l. [m]'
elif alti_type == 'agl':  # TO CORRECT for strectched grid with var 'ALT'
    #create grid mesh (eq. Y)
    alti = np.meshgrid(section_ds.i_sect, section_ds.level)[1]
    alti = xr.DataArray(alti, dims=['level', 'i_sect'])
    #for plot
    ylabel = 'height a.g.l. [m]'
    

### 1.1. Color map (pcolor or contourf)
data1 = section_ds[varname_colormap]
# cm = ax[0].pcolormesh(Xmesh,
#                       alti,
#                       data1.T,
#                       cmap=colormap,
#                       norm=norm_cm,  # for logscale of colormap
#                       # vmin=vmin, vmax=vmax
#                       )

cm = ax[0].contourf(Xmesh,
                    alti,
                    data1.T, 
                    cmap=colormap,  # 'OrRd', 'coolwarm'
                    # norm=norm_cm,  # for logscale of colormap
                    # locator=ticker.LogLocator(),
                    levels=levels,
                    extend = 'both',  #highlights the min and max in different color
                    )
#manage colorbar
divider = make_axes_locatable(ax[0])
cax = divider.append_axes('right', size='2%', pad=0.05)
cbar = fig.colorbar(cm, cax=cax, orientation='vertical')
try:
    cbar.set_label(f'{data1.long_name} [{data1.units}]')
except:
    cbar.set_label(varname_colormap)

### 1.2. Contour map
if varname_contourmap in ['HBLTOP', 'HLOWJET', 'HLOWJET_WS']:  #1D
#    ax[0].plot(section_ds['HBLTOP'] + section_ds['ZS'],
#              linestyle='--', color='r')
    ax[0].plot(section_ds['HLOWJET_NOSE'] + section_ds['ZS'],
              linestyle='-.', color='g')
    ax[0].plot(section_ds['HLOWJET_MAX'] + section_ds['ZS'],
              linestyle='-.', color='y')
else:  
    data2 = section_ds[varname_contourmap]  # x1000 to get it in g/kg if RVT
    cont = ax[0].contour(Xmesh,
                         alti,
                         data2.T,
                         cmap='viridis_r',  #viridis_r, copper_r
#                         colors='r',
#                         levels=np.arange(vmin_contour, vmax_contour),
                         levels=levels_contour,
#                         vmin=vmin, vmax=vmax,  # for adaptative colormap
                         )
    ax[0].clabel(cont, cont.levels, inline=True, fontsize=13)

### 1.3. Winds
if wind_visu == 'horiz':            # 2.1 winds - flat direction and force
    ax[0].barbs(
            #Note that X & alti have dimensions reversed
            Xmesh[::skip_barbs_y, ::skip_barbs_x], 
            alti[::skip_barbs_y, ::skip_barbs_x], 
            #Here dimensions are in the proper order
            section_ds['UT'][::skip_barbs_x, ::skip_barbs_y].T, 
            section_ds['VT'][::skip_barbs_x, ::skip_barbs_y].T, 
            pivot='middle',
            length=5*arrow_size,     #length of barbs
            sizes={
    #              'spacing':1, 'height':1, 'width':1,
                    'emptybarb':0.01},
            barb_increments=barb_size_increments[barb_size_option] # [kts], 1.94kt = 1m/s
            )
    ax[0].annotate(barb_size_description[barb_size_option],
                   xy=(0.1, 0.9),
                   xycoords='subfigure fraction'
                   )
elif wind_visu == 'verti_proj':     # 2.2  winds - verti and projected wind
    Q = ax[0].quiver(
            # Note that X & alti have dimensions reversed
            Xmesh[::skip_barbs_y, ::skip_barbs_x], 
            alti[::skip_barbs_y, ::skip_barbs_x], 
            # Here dimensions are in the proper order
            section_ds['WPROJ'][::skip_barbs_x, ::skip_barbs_y].T, 
            section_ds['WT'][::skip_barbs_x, ::skip_barbs_y].T, 
            pivot='middle',
            scale=150/arrow_size,  # arrows scale, if higher, smaller arrows
            alpha=0.4,
            )
    # Add arrow scale in top-right corner
    u_max = abs(section_ds['WPROJ'][::skip_barbs_x, ::skip_barbs_y]).max()
    ax[0].quiverkey(Q, 0.8, 0.9, 
                    U=u_max, 
                    label=str((np.round(u_max, decimals=1)).data) + ' $m.s^{-1}$', 
                    labelpos='E',
                    coordinates='figure')

# add sites that are between the two defining sites 'start' and 'end'
for site_inter in sites_to_project:
    coords_site_inter = (gv.whole[site_inter]['lat'], gv.whole[site_inter]['lon'])
    
    point_site_inter = Point(coords_site_inter)
    line_cross_section = LineString([start, end])
    dist = line_cross_section.project(point_site_inter)
    coords_site_inter_proj = list(line_cross_section.interpolate(dist).coords)[0]
    
    fraction_lon_point_inter = (coords_site_inter_proj[1] - start[1]) / (end[1] - start[1])
    # in term of abscisse
    list_abscisses_sites = list(abscisse_sites.keys())
    diff_abscisses = list_abscisses_sites[1] - list_abscisses_sites[0]
    abscisse_inter = fraction_lon_point_inter * diff_abscisses + list_abscisses_sites[0]
    # add to the dict
    abscisse_sites[abscisse_inter] = site_inter 

# x-axis with sites names
ax[0].set_xticks(list(abscisse_sites.keys()))
ax[0].set_xticklabels(list(abscisse_sites.values()), 
                   rotation=0, fontsize=12)
#ax[0].set_xticklabels(['La Cendrosa', 'Els Plans'], 
#                       rotation=0, fontsize=12)
# x-axis with lat-lon values
#ax.set_xticks(data1.i_sect[::10])
#ax.set_xticklabels(abscisse_coords[::10], rotation=0, fontsize=9)

# set y limits (height ASL)
if alti_type == 'asl':
    min_ZS = section_ds['ZS'].min()
    ax[0].set_ylim([min_ZS, max_ZS + toplevel])
ax[0].set_ylabel(ylabel)


### 2. Subplot of surface characteristic ---

if subplot_type == 'surface_var':
    
    surfplot = ax[1].pcolor(
        section_ds.i_sect, 
        section_ds[surf_var][:,:2].level, 
        section_ds[surf_var][:,:2].transpose(), 
        cmap='seismic_r',
        vmin=vmin_surf_var, vmax=vmax_surf_var
        )
    # create colorbar dedicated to the subplot
    divider = make_axes_locatable(ax[1])
    cax = divider.append_axes('right', size='2%', pad=0.05)
    cbar2 = fig.colorbar(surfplot, cax=cax, orientation='vertical')
    cbar2.ax.get_yaxis().set_ticks([vmin_surf_var, vmax_surf_var])
    cbar2.set_label(surf_var_label)
    
    # labels_arr = np.arange(0, 100, 10)
    transect_length = (line['index_distance'] + 2*nb_points_beyond)*line['nij_step']/1000  # length in km
    labels_arr = np.arange(0, transect_length, 5)
    tick_pos = labels_arr/ (line['nij_step']/1000)
    ax[1].set_xticks(ticks = tick_pos,
                     labels = labels_arr
                     )
    ax[1].set_xlabel('distance [km]')
    ax[1].set_yticks([])
    ax[1].set_ylabel(surf_var)
    
if subplot_type == 'distance':
    # remove the surface subplot
    fig.delaxes(ax[1])
    
    # get index of torredembarra in abscisse_sites
    distance_ref_ind = list(abscisse_sites.values()).index('carry')
    # get corresponding abscisse for torredembarra
    distance_ref_xval = list(abscisse_sites.keys())[distance_ref_ind]
    
    def ftest(x):
        return -(x - distance_ref_xval) * (line['nij_step']/1000)
    def ftest_recip(x):
        return -(x/(line['nij_step']/1000) + distance_ref_xval)
    
    # add secondary axis
    secax = ax[0].secondary_xaxis(-0.1, functions=(ftest, ftest_recip))
    secax.set_xlabel('distance to the sea [km]')


### Global options
plot_title = 'Cross section on {0}-{1}-{2}'.format(
        wanted_date, model, wind_visu)
#plot_title = 'Cross section of ABL between irrigated and rainfed areas on July 22 at 12:00 - {0}'.format(
#        model)
#plot_title = 'Cross section on July 22 at 12:00 - {0}'.format(model)
# ax[0].set_title(plot_title)
fig.suptitle(plot_title)

if save_plot:
    tools.save_figure(plot_title, save_folder)