More visualisations

import os
import pandas as pd
import numpy as np
import xarray as xr
import dask as da
import zarr
import holoviews as hv
import hvplot.xarray
import matplotlib.pyplot as plt

from chaosmagpy.plot_utils import nio_colormap

TMPDIR = os.getcwd()
zarr_store = os.path.join(TMPDIR, "datacube_test.zarr")
print("Using:", zarr_store)

xr.set_options(
    display_expand_attrs=False,
    display_expand_data_vars=True
);

Using: /home/ash/code/geomagnetic_datacubes_dev/notebooks/datacube_test.zarr

ds = xr.open_dataset(
    zarr_store, engine="zarr",
)
# Dataset downsampled by 1/30 (i.e. 5-minute sampling)
_ds = ds.isel(Timestamp=slice(0, -1, 30))

def make_plots(labels=True, titles=True):
    plots = {}
    for NEC in "NEC":
        kwargs = dict()
        if not labels:
            kwargs["xlabel"] = ""
            kwargs["ylabel"] = ""
        plots[f"{NEC}_core"] = (
            _ds.drop("Timestamp").sel(NEC=NEC)
            .hvplot.scatter(
                x="Longitude", y="Latitude", c="B_NEC",
                rasterize=True,
                colorbar=False,
                hover=False,
                width=300, height=200,
                clim=(-50000, 50000), cmap=nio_colormap(),
                title=f"B_{NEC} (GEO)" if titles else "",
                **kwargs
            )
        )
        plots[f"{NEC}_iono"] = (
            _ds.drop("Timestamp").sel(NEC=NEC)
            .hvplot.scatter(
                x="Longitude", y="Latitude", c="B_NEC_res_CHAOS-full",
                rasterize=True,
                colorbar=False,
                hover=False,
                tools=[],
                width=300, height=200,
                clim=(-50, 50), cmap=nio_colormap(),
                title=f"B_{NEC}_res_CHAOS-full (GEO)" if titles else "",
                **kwargs
            )
        )
        plots[f"{NEC}_iono_qdmlt"] = (
            _ds.drop("Timestamp").sel(NEC=NEC)
            .hvplot.scatter(
                x="MLT", y="QDLat", c="B_NEC_res_CHAOS-full" if titles else "",
                rasterize=True,
                colorbar=False,
                hover=False,
                tools=[],
                width=300, height=200,
                clim=(-50, 50), cmap=nio_colormap(),
                title=f"B_{NEC}_res_CHAOS-full (QDMLT)",
                **kwargs
            )
        )
    return plots

plots = make_plots(labels=False, titles=True)

(
    plots["N_core"] + plots["N_iono"] + plots["N_iono_qdmlt"] +
    plots["E_core"] + plots["E_iono"] + plots["E_iono_qdmlt"] +
    plots["C_core"] + plots["C_iono"] + plots["C_iono_qdmlt"]
).cols(3)

# # can't set dpi with bokeh (matplotlib backend bugs too)
# # -> adjust the size and save that instead
# hv.save(
#     (
#         plots["N_core"] + plots["N_iono"] + plots["N_iono_qdmlt"] +
#         plots["E_core"] + plots["E_iono"] + plots["E_iono_qdmlt"] +
#         plots["C_core"] + plots["C_iono"] + plots["C_iono_qdmlt"]
#     ).cols(3).opts(width=1000, height=500),
#     f"datacube.png"
# )