Plotting temperature contour by depth

[aiza-k]

This function will create a contour plot of temperature by soil depth. It will call load_trsc_dataframe from output_utils, so make sure to import this function.

Note these bugs:

1. there is an issue with cutting off the tail end of the data by depth when plotting.
2. there may be interpolation issues, as seen with the talik formation (to be debugged).
3. there is a bug making the colorbar symmetric - current implementation using vmin and vmax doesn't work

def soil_contourbydepth(data_path,output_var,res,px,py,output_folder, start_time=None, end_time=None, 
                  depth_start=None, depth_end=None, n=100):
    '''
    data_path: path to output data to plot
    output_var: variable to plot (technically, doesn't need to be temperature)
    res: time resolution of output data
    px, py: pixel location of data
    output_folder: name of output folder
    start_time, end_time: start and end year of data to plot
    depth_start, depth_end: starting and ending depth to be plotted (in meters)
    n: levels on the colobar
    '''
    
    os.chdir(data_path)
    df, meta = load_trsc_dataframe(var=output_var, timeres=res, px_y=py, px_x=px, fileprefix=output_folder) 
    df_depth, meta_depth = load_trsc_dataframe(var='LAYERDEPTH',  timeres=res, px_y=py, px_x=px, fileprefix=output_folder)
    df_dz, meta_dzh = load_trsc_dataframe(var='LAYERDZ',  timeres=res, px_y=py, px_x=px, fileprefix=output_folder)
    
    layers = df.columns.astype(float)
    times = pd.to_datetime(df.index)
    
    # Filter data based on start_time and end_time
    if start_time is not None and end_time is not None:
        mask = (times >= start_time) & (times <= end_time)
        df = df.loc[mask]
        times = times[mask]

    # Extract necessary data
    depths = df_depth.iloc[:, :-2].values
    dz = df_dz.iloc[:, :-2].values
    temperature = df.iloc[:, :-2].values
    xp = depths + dz / 2  # Center of each layer, x-coordinates of the data points for interp1d

    # Create a regular grid of depth values
    ii=np.unravel_index(np.argmax(depths), depths.shape)
    maxd=depths.max()+(dz[ii]/2)
    regular_depths = np.arange(0, maxd, 0.01)

    # Interpolate temperature onto the regular grid
    interp_temperature = np.empty((temperature.shape[0], regular_depths.shape[0]))
    for i in range(temperature.shape[0]):
        f = interp1d(xp[i], temperature[i], kind='linear', fill_value='extrapolate')
        interp_temperature[i] = f(regular_depths)

    # Create contour plot
    color_axes = max(np.max(temperature), np.abs(np.min(temperature)))
    plt.contourf(times, regular_depths, interp_temperature.T, cmap='seismic', vmin=-color_axes, vmax=color_axes, levels=n)

    # Add colorbar
    plt.colorbar(label='Temperature (C)')

    # Set labels and title
    plt.xlabel('Time')
    plt.ylabel('Depth (m)')

    # Show the plot
    plt.ylim(depth_start, depth_end)
    plt.gca().invert_yaxis()
    return

You can call the function like this:
To choose start and end times to plot, you need to specify Year-Month-Date, like start_time='1980-01-01', end_time='2100-12-31'

fig, ax = plt.subplots(1,1,figsize=(15, 5))
data_path='/data/workflows/MD_deep_soil/experiment_allout_monthly_11px/'
soil_contourbydepth(data_path,'TLAYER','monthly',1,1,'output', start_time=None, end_time=None, depth_start=0, 
                    depth_end=30, n=50)
plt.title('Temperature Contour Plot: Deep Soil branch MD Pixel 1,1, cmt 1')