SpaceTime PSD

import sys, os

# spyder up to find the root
oceanbench_root = "/gpfswork/rech/cli/uvo53rl/projects/oceanbench"

# append to path
sys.path.append(str(oceanbench_root))

import autoroot
import typing as tp
import jax
import jax.numpy as jnp
import jax.scipy as jsp
import jax.random as jrandom
import numpy as np
import numba as nb
import pandas as pd
import equinox as eqx
import kernex as kex
import finitediffx as fdx
import diffrax as dfx
import xarray as xr
import matplotlib.pyplot as plt
import seaborn as sns
from tqdm.notebook import tqdm, trange
from jaxtyping import Float, Array, PyTree, ArrayLike
import wandb
from omegaconf import OmegaConf
import hydra
import metpy
from sklearn.pipeline import Pipeline
from jejeqx._src.transforms.dataframe.spatial import Spherical2Cartesian
from jejeqx._src.transforms.dataframe.temporal import TimeDelta
from jejeqx._src.transforms.dataframe.scaling import MinMaxDF


sns.reset_defaults()
sns.set_context(context="poster", font_scale=0.7)
jax.config.update("jax_enable_x64", False)

%matplotlib inline
%load_ext autoreload
%autoreload 2

Processing Chain

Part I:

• Open Dataset

• Validate Coordinates + Variables

• Decode Time

• Select Region

• Sortby Time

Part II: Regrid

Part III:

• Interpolate Nans

• Add Units

• Spatial Rescale

• Time Rescale

Part IV: Metrics

Data

# !wget wget -nc https://s3.us-east-1.wasabisys.com/melody/osse_data/ref/NATL60-CJM165_GULFSTREAM_ssh_y2013.1y.nc

# !cat configs/postprocess.yaml

# # load config
# config_dm = OmegaConf.load('./configs/postprocess.yaml')

# # instantiate
# ds = hydra.utils.instantiate(config_dm.NATL60_GF_1Y1D)
# ds

Reference Dataset

For the reference dataset, we will look at the NEMO simulation of the Gulfstream.

%%time

# load config
config_dm = OmegaConf.load("./configs/postprocess.yaml")

# instantiate
ds_natl60 = hydra.utils.instantiate(config_dm.NATL60_GF_FULL).compute()
ds_natl60

Prediction Datasets - NADIR

%%time

experiment = "swot"  # "nadir" #
if experiment == "nadir":
    # load config
    results_config = OmegaConf.load(f"./configs/results_dc20a_nadir.yaml")

    # instantiate
    ds_duacs = hydra.utils.instantiate(results_config.DUACS_NADIR.data).compute()
    ds_miost = hydra.utils.instantiate(results_config.MIOST_NADIR.data).compute()
    ds_nerf_siren = hydra.utils.instantiate(
        results_config.NERF_SIREN_NADIR.data
    ).compute()
    ds_nerf_ffn = hydra.utils.instantiate(results_config.NERF_FFN_NADIR.data).compute()
    ds_nerf_mlp = hydra.utils.instantiate(results_config.NERF_MLP_NADIR.data).compute()
elif experiment == "swot":
    # load config
    results_config = OmegaConf.load(f"./configs/results_dc20a_swot.yaml")

    # instantiate
    ds_duacs = hydra.utils.instantiate(results_config.DUACS_SWOT.data).compute()
    ds_miost = hydra.utils.instantiate(results_config.MIOST_SWOT.data).compute()
    ds_nerf_siren = hydra.utils.instantiate(
        results_config.NERF_SIREN_SWOT.data
    ).compute()
    ds_nerf_ffn = hydra.utils.instantiate(results_config.NERF_FFN_SWOT.data).compute()
    ds_nerf_mlp = hydra.utils.instantiate(results_config.NERF_MLP_SWOT.data).compute()

Regrdding

from oceanbench._src.geoprocessing.gridding import grid_to_regular_grid

%%time

ds_duacs = grid_to_regular_grid(
    src_grid_ds=ds_duacs.pint.dequantify(),
    tgt_grid_ds=ds_natl60.pint.dequantify(),
    keep_attrs=False,
)
ds_miost = grid_to_regular_grid(
    src_grid_ds=ds_miost.pint.dequantify(),
    tgt_grid_ds=ds_natl60.pint.dequantify(),
    keep_attrs=False,
)
ds_nerf_siren = grid_to_regular_grid(
    src_grid_ds=ds_nerf_siren.pint.dequantify(),
    tgt_grid_ds=ds_natl60.pint.dequantify(),
    keep_attrs=False,
)
ds_nerf_ffn = grid_to_regular_grid(
    src_grid_ds=ds_nerf_ffn.pint.dequantify(),
    tgt_grid_ds=ds_natl60.pint.dequantify(),
    keep_attrs=False,
)
ds_nerf_mlp = grid_to_regular_grid(
    src_grid_ds=ds_nerf_mlp.pint.dequantify(),
    tgt_grid_ds=ds_natl60.pint.dequantify(),
    keep_attrs=False,
)

Preprocess Chain

%%time

# load config
psd_config = OmegaConf.load("./configs/metrics.yaml")

ds_natl60 = hydra.utils.instantiate(psd_config.psd_preprocess_chain)(
    ds_natl60.pint.dequantify()
)
ds_duacs = hydra.utils.instantiate(psd_config.psd_preprocess_chain)(
    ds_duacs.pint.dequantify()
)
ds_miost = hydra.utils.instantiate(psd_config.psd_preprocess_chain)(
    ds_miost.pint.dequantify()
)
ds_nerf_siren = hydra.utils.instantiate(psd_config.psd_preprocess_chain)(
    ds_nerf_siren.pint.dequantify()
)
ds_nerf_ffn = hydra.utils.instantiate(psd_config.psd_preprocess_chain)(
    ds_nerf_ffn.pint.dequantify()
)
ds_nerf_mlp = hydra.utils.instantiate(psd_config.psd_preprocess_chain)(
    ds_nerf_mlp.pint.dequantify()
)

Power Spectrum (Spacetime)

%%time

# load config
psd_config = OmegaConf.load("./configs/metrics.yaml")

ds_natl60_psd = hydra.utils.instantiate(psd_config.psd_spacetime_chain)(
    ds_natl60.pint.dequantify()
)
ds_duacs_psd = hydra.utils.instantiate(psd_config.psd_spacetime_chain)(
    ds_duacs.pint.dequantify()
)
ds_miost_psd = hydra.utils.instantiate(psd_config.psd_spacetime_chain)(
    ds_miost.pint.dequantify()
)
ds_nerf_siren_psd = hydra.utils.instantiate(psd_config.psd_spacetime_chain)(
    ds_nerf_siren.pint.dequantify()
)
ds_nerf_ffn_psd = hydra.utils.instantiate(psd_config.psd_spacetime_chain)(
    ds_nerf_ffn.pint.dequantify()
)
ds_nerf_mlp_psd = hydra.utils.instantiate(psd_config.psd_spacetime_chain)(
    ds_nerf_mlp.pint.dequantify()
)

from jejeqx._src.viz.xarray.psd import PlotPSDSpaceTime, PlotPSDSpaceTimeScore

fig, ax = plt.subplots(ncols=2, nrows=3, figsize=(15, 20))

# NATL60
psd_st_plot = PlotPSDSpaceTime()
psd_st_plot.init_fig(ax=ax[0, 0])
psd_st_plot.plot_wavelength(
    ds_natl60_psd.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
ax[0, 0].set(title="NATL60")

# DUACS
psd_st_plot = PlotPSDSpaceTime()
psd_st_plot.init_fig(ax=ax[0, 1])
psd_st_plot.plot_wavelength(
    ds_duacs_psd.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
ax[0, 1].set(title="DUACS")

# MIOST
psd_st_plot = PlotPSDSpaceTime()
psd_st_plot.init_fig(ax=ax[1, 0])
psd_st_plot.plot_wavelength(
    ds_miost_psd.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
ax[1, 0].set(title="MIOST")

# NERF - MLP
psd_st_plot = PlotPSDSpaceTime()
psd_st_plot.init_fig(ax=ax[1, 1])
psd_st_plot.plot_wavelength(
    ds_nerf_mlp_psd.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
ax[1, 1].set(title="NerF (MLP)")

# NERF - FFN
psd_st_plot = PlotPSDSpaceTime()
psd_st_plot.init_fig(ax=ax[2, 0])
psd_st_plot.plot_wavelength(
    ds_nerf_ffn_psd.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
ax[2, 0].set(title="NerF (FFN)")

# NERF - SIREN
psd_st_plot = PlotPSDSpaceTime()
psd_st_plot.init_fig(ax=ax[2, 1])
psd_st_plot.plot_wavelength(
    ds_nerf_siren_psd.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
ax[2, 1].set(title="NerF (SIREN)")

plt.tight_layout()
plt.gcf().savefig(f"./figures/dc20a/psd/spacetime/dc20a_psd_spacetime_{experiment}.png")
plt.show()

Power Spectrum Score (Spacetime)

%%time

# load config
psd_config = OmegaConf.load("./configs/metrics.yaml")

ds_psd_duacs_score = hydra.utils.instantiate(
    psd_config.psd_spacetime_score,
    da=ds_duacs.pint.dequantify(),
    da_ref=ds_natl60.pint.dequantify(),
)

ds_psd_miost_score = hydra.utils.instantiate(
    psd_config.psd_spacetime_score,
    da=ds_miost.pint.dequantify(),
    da_ref=ds_natl60.pint.dequantify(),
)

ds_psd_nerf_mlp_score = hydra.utils.instantiate(
    psd_config.psd_spacetime_score,
    da=ds_nerf_mlp.pint.dequantify(),
    da_ref=ds_natl60.pint.dequantify(),
)

ds_psd_nerf_ffn_score = hydra.utils.instantiate(
    psd_config.psd_spacetime_score,
    da=ds_nerf_ffn.pint.dequantify(),
    da_ref=ds_natl60.pint.dequantify(),
)

ds_psd_nerf_siren_score = hydra.utils.instantiate(
    psd_config.psd_spacetime_score,
    da=ds_nerf_siren.pint.dequantify(),
    da_ref=ds_natl60.pint.dequantify(),
)

NATL60

# NATL60
fig, ax = plt.subplots(figsize=(7, 5.5))

psd_st_plot = PlotPSDSpaceTime()
psd_st_plot.init_fig(ax=ax)
psd_st_plot.plot_wavelength(
    ds_natl60_psd.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
plt.tight_layout()
plt.gcf().savefig(
    f"./figures/dc20a/psd/spacetime/dc20a_psd_spacetime_ssh_natl60_{experiment}.png"
)
plt.close()

DUACS

# DUACS
fig, ax = plt.subplots(figsize=(7, 5.5))

psd_st_plot = PlotPSDSpaceTime()
psd_st_plot.init_fig(ax=ax)
psd_st_plot.plot_wavelength(
    ds_duacs_psd.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
plt.tight_layout()
plt.gcf().savefig(
    f"./figures/dc20a/psd/spacetime/dc20a_psd_spacetime_ssh_duacs_{experiment}.png"
)
plt.close()

MIOST

# DUACS
fig, ax = plt.subplots(figsize=(7, 5.5))

psd_st_plot = PlotPSDSpaceTime()
psd_st_plot.init_fig(ax=ax)
psd_st_plot.plot_wavelength(
    ds_miost_psd.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
plt.tight_layout()
plt.gcf().savefig(
    f"./figures/dc20a/psd/spacetime/dc20a_psd_spacetime_ssh_miost_{experiment}.png"
)
plt.close()

NERF (MLP)

# DUACS
fig, ax = plt.subplots(figsize=(7, 5.5))

psd_st_plot = PlotPSDSpaceTime()
psd_st_plot.init_fig(ax=ax)
psd_st_plot.plot_wavelength(
    ds_nerf_mlp_psd.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
plt.tight_layout()
plt.gcf().savefig(
    f"./figures/dc20a/psd/spacetime/dc20a_psd_spacetime_ssh_nerf_mlp_{experiment}.png"
)
plt.close()

NERF (FFN)

# DUACS
fig, ax = plt.subplots(figsize=(7, 5.5))

psd_st_plot = PlotPSDSpaceTime()
psd_st_plot.init_fig(ax=ax)
psd_st_plot.plot_wavelength(
    ds_nerf_ffn_psd.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
plt.tight_layout()
plt.gcf().savefig(
    f"./figures/dc20a/psd/spacetime/dc20a_psd_spacetime_ssh_nerf_ffn_{experiment}.png"
)
plt.close()

NERF (SIREN)

# DUACS
fig, ax = plt.subplots(figsize=(7, 5.5))

psd_st_plot = PlotPSDSpaceTime()
psd_st_plot.init_fig(ax=ax)
psd_st_plot.plot_wavelength(
    ds_nerf_siren_psd.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
plt.tight_layout()
plt.gcf().savefig(
    f"./figures/dc20a/psd/spacetime/dc20a_psd_spacetime_ssh_nerf_siren_{experiment}.png"
)
plt.close()

from jejeqx._src.viz.xarray.psd import PlotPSDSpaceTime, PlotPSDSpaceTimeScore

fig, ax = plt.subplots(ncols=2, nrows=3, figsize=(15, 20))

# NATL60
psd_st_plot = PlotPSDSpaceTime()
psd_st_plot.init_fig(ax=ax[0, 0])
psd_st_plot.plot_wavelength(
    ds_natl60_psd.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
ax[0, 0].set(title="NATL60")

# DUACS
psd_st_plot = PlotPSDSpaceTimeScore()
psd_st_plot.init_fig(ax=ax[0, 1])
psd_st_plot.plot_wavelength(
    ds_psd_duacs_score.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
ax[0, 1].set(title="DUACS")

# MIOST
psd_st_plot = PlotPSDSpaceTimeScore()
psd_st_plot.init_fig(ax=ax[1, 0])
psd_st_plot.plot_wavelength(
    ds_psd_miost_score.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
ax[1, 0].set(title="MIOST")

# NERF - MLP
psd_st_plot = PlotPSDSpaceTimeScore()
psd_st_plot.init_fig(ax=ax[1, 1])
psd_st_plot.plot_wavelength(
    ds_psd_nerf_mlp_score.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
ax[1, 1].set(title="NerF (MLP)")

# NERF - FFN
psd_st_plot = PlotPSDSpaceTimeScore()
psd_st_plot.init_fig(ax=ax[2, 0])
psd_st_plot.plot_wavelength(
    ds_psd_nerf_ffn_score.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
ax[2, 0].set(title="NerF (FFN)")

# NERF - SIREN
psd_st_plot = PlotPSDSpaceTimeScore()
psd_st_plot.init_fig(ax=ax[2, 1])
psd_st_plot.plot_wavelength(
    ds_psd_nerf_siren_score.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
ax[2, 1].set(title="NerF (SIREN)")

plt.tight_layout()
plt.gcf().savefig(
    f"./figures/dc20a/psd_score/spacetime/dc20a_psd_spacetime_score_{experiment}.png"
)
plt.show()

DUACS

# DUACS
fig, ax = plt.subplots(figsize=(7, 5.5))

psd_st_plot = PlotPSDSpaceTimeScore()
psd_st_plot.init_fig(ax=ax)
psd_st_plot.plot_wavelength(
    ds_psd_duacs_score.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
plt.tight_layout()
plt.gcf().savefig(
    f"./figures//dc20a/psd_score/spacetime/dc20a_psd_spacetime_score_ssh_duacs_{experiment}.png"
)
plt.close()

MIOST

# DUACS
fig, ax = plt.subplots(figsize=(7, 5.5))

psd_st_plot = PlotPSDSpaceTimeScore()
psd_st_plot.init_fig(ax=ax)
psd_st_plot.plot_wavelength(
    ds_psd_miost_score.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
plt.tight_layout()
plt.gcf().savefig(
    f"./figures/dc20a/psd_score/spacetime/dc20a_psd_spacetime_score_ssh_miost_{experiment}.png"
)
plt.close()

NERF (MLP)

# DUACS
fig, ax = plt.subplots(figsize=(7, 5.5))

psd_st_plot = PlotPSDSpaceTimeScore()
psd_st_plot.init_fig(ax=ax)
psd_st_plot.plot_wavelength(
    ds_psd_nerf_mlp_score.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
plt.tight_layout()
plt.gcf().savefig(
    f"./figures/dc20a/psd_score/spacetime/dc20a_psd_spacetime_score_ssh_nerf_mlp_{experiment}.png"
)
plt.close()

NERF (FFN)

# DUACS
fig, ax = plt.subplots(figsize=(7, 5.5))

psd_st_plot = PlotPSDSpaceTimeScore()
psd_st_plot.init_fig(ax=ax)
psd_st_plot.plot_wavelength(
    ds_psd_nerf_ffn_score.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
plt.tight_layout()
plt.gcf().savefig(
    f"./figures/dc20a/psd_score/spacetime/dc20a_psd_spacetime_score_ssh_nerf_ffn_{experiment}.png"
)
plt.close()

NERF (SIREN)

# DUACS
fig, ax = plt.subplots(figsize=(7, 5.5))

psd_st_plot = PlotPSDSpaceTimeScore()
psd_st_plot.init_fig(ax=ax)
psd_st_plot.plot_wavelength(
    ds_psd_nerf_siren_score.ssh,
    space_scale=1e3,
    space_units="km",
    time_units="days",
    psd_units="SSH",
)
plt.tight_layout()
plt.gcf().savefig(
    f"./figures/dc20a/psd_score/spacetime/dc20a_psd_spacetime_score_ssh_nerf_siren_{experiment}.png"
)
plt.close()