virtualship/docs/paper/figure1.py at e3bb6f910195e581055cb2cb176dde82bc0b3c87 · Parcels-code/virtualship · 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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
# %%

import cartopy.crs as ccrs
import matplotlib.gridspec as gridspec
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
from plotting_functions import (
    _ctd_distance_along_expedition,
    plot_adcp,
    plot_ctd,
    plot_drifters,
)

SAMPLE_DIR = "sample_expedition/"
CONFIG = "expedition.yaml"
EXPEDITION = "MY_EXPEDITION"

# %%

# VirtualShip output
ctd_ds = xr.open_dataset(f"{SAMPLE_DIR}{EXPEDITION}/results/ctd.zarr")
ctd_bgc_ds = xr.open_dataset(f"{SAMPLE_DIR}{EXPEDITION}/results/ctd_bgc.zarr")
drifter_ds = xr.open_dataset(f"{SAMPLE_DIR}{EXPEDITION}/results/drifter.zarr")
adcp_ds = xr.open_dataset(f"{SAMPLE_DIR}{EXPEDITION}/results/adcp.zarr")


# %%

# plot

PROJ = ccrs.PlateCarree()

waypoint_distances = np.unique(_ctd_distance_along_expedition(ctd_ds)["distance"])


def add_waypoint_markers(ax, distances, offset=15, marker_size=70):
    ax.scatter(
        (distances / 1000),
        np.zeros_like(distances) - offset,
        marker="v",
        color="black",
        edgecolors="white",
        s=marker_size,
        clip_on=False,
    )


def add_title(ax, title, fontsize=13, y0=1.03):
    """Add title."""
    ax.text(
        0,
        y0,
        title,
        ha="left",
        va="bottom",
        transform=ax.transAxes,
        fontsize=fontsize,
    )


# fig
fig = plt.figure(figsize=(9, 13), dpi=300)

# custom layout
gs = gridspec.GridSpec(3, 2, height_ratios=[1.5, 1, 1])
ax0 = fig.add_subplot(gs[0, :])
ax1 = fig.add_subplot(gs[1, 0])
ax2 = fig.add_subplot(gs[1, 1], projection=PROJ)
ax3 = fig.add_subplot(gs[2, 0])
ax4 = fig.add_subplot(gs[2, 1])

# overview image
add_title(ax0, r"$\bf{a}$" + ") MFP expedition overview", y0=1.01)
img = mpimg.imread(f"{SAMPLE_DIR}expedition_overview.png")
ax0.imshow(img)
ax0.axis("off")

# adcp
add_title(ax1, r"$\bf{b}$" + ") ADCP (flow speed)")
ax1.set_ylabel("Depth (m)")
ax1.set_xlabel("Distance (km)")
plot_adcp(adcp_ds, ax1)
add_waypoint_markers(ax1, waypoint_distances)

# drifters
add_title(ax2, r"$\bf{c}$" + ") Surface drifters")
plot_drifters(
    drifter_ds,
    ax2,
    vmin=drifter_ds.temperature.min(),
    vmax=drifter_ds.temperature.max(),
)

# CTD (temperature)
add_title(ax3, r"$\bf{d}$" + ") CTD (temperature)")
ax3.set_ylabel("Depth (m)")
ax3.set_xlabel("Distance (km)")
_, _distances_regular, _var_masked = plot_ctd(
    ctd_ds,
    ax3,
    plot_variable="temperature",
    vmin=ctd_ds.temperature.min(),
    vmax=ctd_ds.temperature.max(),
)

ctd_wp_distances = _distances_regular[np.nansum(_var_masked, axis=1) > 0]
add_waypoint_markers(ax3, ctd_wp_distances, offset=45, marker_size=60)

# CTD (oxygen)
add_title(ax4, r"$\bf{e}$" + ") CTD (oxygen)")
ax4.set_xlabel("Distance (km)")
plot_ctd(
    ctd_bgc_ds, ax4, plot_variable="oxygen", vmin=0, vmax=ctd_bgc_ds.o2.max()
)  # vmin tailored to mark red as approximate oxygen minimum zone (~ 45 mmol/m-3)
add_waypoint_markers(ax4, ctd_wp_distances, offset=45, marker_size=60)


plt.tight_layout()
plt.show()

fig.savefig("figure1.png", dpi=300, bbox_inches="tight")

# %%