visualization/z_time.py at main · BottomRedoxModel/visualization · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import ticker
import matplotlib.dates as mdates
import config as cfg
import grid_maker as gm

sed = cfg.sed
sed2 = cfg.sed2
t1 = cfg.t1_ztime
t2 = cfg.t2_ztime
bbl_h = cfg.bbl_h

yspace = cfg.yspace


def plot_param(ds, name, x, y, y_sed, axis,axis_cb,axis_sed,axis_cb_sed):

    var = ds[name].values
    print(var.shape)
    levels = np.linspace(np.min(var[:,:sed2]), np.max(var[:,:sed2]), 30)
    sed_levels = np.linspace(np.min(var[:,sed2:]), np.max(var[:,sed2:]), 30)
    if all(lev == levels[0] for lev in levels):
        levels = np.linspace(levels[0], levels[0]+0.1, 30)
    if all(lev == sed_levels[0] for lev in sed_levels):
        sed_levels = np.linspace(sed_levels[0], sed_levels[0]+0.1, 30)

    X,Y = np.meshgrid(x,y[:sed2])
    X_sed,Y_sed = np.meshgrid(x,y_sed[sed2:])
    if name in cfg.cmap_dict.keys():
        cmap = cfg.cmap_dict[name]
    else:
        cmap = 'turbo'

    if var[:, sed2:].shape[1] == len(Y_sed):
        CS_1_sed = axis_sed.contourf(X_sed,Y_sed, var[:,sed2:].T, levels = sed_levels, cmap = cmap)
    else:
        CS_1_sed = axis_sed.contourf(X_sed, Y_sed, var[:, sed2+1:].T, levels=sed_levels, cmap=cmap)

    # TODO: fix it
    if name == 'Oxy':
        cmap = 'plasma'
    CS_1 = axis.contourf(X, Y, var[:, :sed2].T, levels=levels, cmap=cmap)

    locw = ticker.MaxNLocator(nbins=2, steps=[2, 3, 5, 10])

    cb = plt.colorbar(CS_1,cax = axis_cb)
    cb.ax.yaxis.set_major_locator(locw)

    cb_sed = plt.colorbar(CS_1_sed,cax = axis_cb_sed)
    locs = ticker.MaxNLocator(nbins=2, steps=[2, 3, 5, 10])
    cb_sed.ax.yaxis.set_major_locator(locs)

    axis.set_ylim(np.max(y[:sed2]),0)
    axis_sed.set_ylim(bbl_h,-bbl_h)

    axis_sed.axhline(0,linestyle = '--',linewidth = 0.5,color = 'w')

    axis.tick_params(axis='y', pad = 0.01)
    axis_sed.tick_params(axis='y', pad = 1)

    years = mdates.YearLocator(yspace)  # every Xth year
    years_fmt = mdates.DateFormatter('%Y')

    axis.xaxis.set_major_locator(years)
    axis.xaxis.set_major_formatter(years_fmt)
    axis_sed.xaxis.set_major_locator(years)
    axis_sed.xaxis.set_major_formatter(years_fmt)

    axis.format_xdata = mdates.DateFormatter('%Y-%m-%d')
    axis.set_xticklabels([])

    title = '%s, $\mu M$' % name

    # TODO: check how to simplify this
    for unit, vnames in cfg.units_dict.items():
        if name in vnames:
            title = name + ', ' + unit
            break

    axis.set_title(title)

def fig_ztime(ds, picname, varnames, icol, nrows, ncols):
    global gs

    ds = ds.sel(time=slice(t1, t2)).isel(i=icol)
    xs = ds['time'].values
    ys = ds['z'].values
    y2s = ds['z2'].values
    y_sed = ((ys - ys[sed]) * 100)

    nv = len(varnames)

    fig, (axes, axes_cb, axes_sed, axes_sed_cb) = gm.get_fig_axes(nrows, ncols, nv)

    for i in np.arange(len(varnames)):
        print(i,varnames[i])
        if varnames[i] != 'Kz':
            plot_param(ds, varnames[i], xs, ys, y_sed, axes[i], axes_cb[i], axes_sed[i], axes_sed_cb[i])
        else:
            plot_param(ds, varnames[i], xs, y2s, y_sed, axes[i], axes_cb[i], axes_sed[i], axes_sed_cb[i])

    plt.savefig(picname + '.png', bbox_inches='tight', dpi=300)