Pixel Densities

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 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

The autoreload extension is already loaded. To reload it, use:
  %reload_ext autoreload

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

!ls /gpfswork/rech/yrf/commun/data_challenges/dc20a_osse/staging/results/DUACS

# !cat configs/postprocess.yaml

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

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

from oceanbench._src.geoprocessing.gridding import (
    grid_to_regular_grid,
    coord_based_to_grid,
)
from oceanbench._src.geoprocessing import geostrophic as geocalc
from metpy.units import units

def calculate_physical_quantities(da):
    da["ssh"] *= units.meters
    da["lon"] = da.lon * units.degrees
    da["lat"] = da.lat * units.degrees
    da = geocalc.streamfunction(da, "ssh")
    da = geocalc.geostrophic_velocities(da, variable="psi")
    da = geocalc.kinetic_energy(da, variables=["u", "v"])
    da = geocalc.divergence(da, variables=["u", "v"])
    da = geocalc.coriolis_normalized(da, "div")
    da = geocalc.relative_vorticity(da, variables=["u", "v"])
    da = geocalc.coriolis_normalized(da, "vort_r")
    da = geocalc.strain_magnitude(da, variables=["u", "v"])
    da = geocalc.coriolis_normalized(da, variable="strain")
    return da

def correct_labels(ds):
    ds["lon"].attrs["units"] = "degrees"
    ds["lat"].attrs["units"] = "degrees"
    ds["ssh"].attrs["units"] = "m"
    ds["ssh"].attrs["standard_name"] = "sea_surface_height"
    ds["ssh"].attrs["long_name"] = "Sea Surface Height"
    ds["lon"].attrs["standard_name"] = "longitude"
    ds["lat"].attrs["standard_name"] = "latitude"
    ds["lat"].attrs["long_name"] = "Latitude"
    ds["lon"].attrs["long_name"] = "Longitude"

    return 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

CPU times: user 13.1 s, sys: 2.93 s, total: 16 s
Wall time: 45.1 s

<xarray.Dataset>
Dimensions:  (time: 42, lat: 600, lon: 600)
Coordinates:
  * lon      (lon) float64 -64.98 -64.97 -64.95 -64.93 ... -55.03 -55.02 -55.0
  * lat      (lat) float64 33.02 33.03 33.05 33.07 ... 42.95 42.97 42.98 43.0
  * time     (time) datetime64[ns] 2012-10-22 2012-10-23 ... 2012-12-02
Data variables:
    ssh      (time, lat, lon) float32 0.6549 0.6571 0.6593 ... -0.2152 -0.2174
Attributes:
    Info:     Horizontal grid read in regulargrid_NATL60.nc / Source field re...
    About:    Created by SOSIE interpolation environement => https://github.c...

xarray.Dataset

• Dimensions:
  • time: 42
  • lat: 600
  • lon: 600
• Coordinates: (3)
  • lon
    (lon)
    float64
    -64.98 -64.97 ... -55.02 -55.0

    array([-64.983333, -64.966667, -64.95    , ..., -55.033333, -55.016667,
           -55.      ])

  • lat
    (lat)
    float64
    33.02 33.03 33.05 ... 42.98 43.0

    units :

    degrees_west

    standard_name :

    latitude

    long_name :

    Latitude

    array([33.016667, 33.033333, 33.05    , ..., 42.966667, 42.983333, 43.      ])

  • time
    (time)
    datetime64[ns]
    2012-10-22 ... 2012-12-02

    array(['2012-10-22T00:00:00.000000000', '2012-10-23T00:00:00.000000000',
           '2012-10-24T00:00:00.000000000', '2012-10-25T00:00:00.000000000',
           '2012-10-26T00:00:00.000000000', '2012-10-27T00:00:00.000000000',
           '2012-10-28T00:00:00.000000000', '2012-10-29T00:00:00.000000000',
           '2012-10-30T00:00:00.000000000', '2012-10-31T00:00:00.000000000',
           '2012-11-01T00:00:00.000000000', '2012-11-02T00:00:00.000000000',
           '2012-11-03T00:00:00.000000000', '2012-11-04T00:00:00.000000000',
           '2012-11-05T00:00:00.000000000', '2012-11-06T00:00:00.000000000',
           '2012-11-07T00:00:00.000000000', '2012-11-08T00:00:00.000000000',
           '2012-11-09T00:00:00.000000000', '2012-11-10T00:00:00.000000000',
           '2012-11-11T00:00:00.000000000', '2012-11-12T00:00:00.000000000',
           '2012-11-13T00:00:00.000000000', '2012-11-14T00:00:00.000000000',
           '2012-11-15T00:00:00.000000000', '2012-11-16T00:00:00.000000000',
           '2012-11-17T00:00:00.000000000', '2012-11-18T00:00:00.000000000',
           '2012-11-19T00:00:00.000000000', '2012-11-20T00:00:00.000000000',
           '2012-11-21T00:00:00.000000000', '2012-11-22T00:00:00.000000000',
           '2012-11-23T00:00:00.000000000', '2012-11-24T00:00:00.000000000',
           '2012-11-25T00:00:00.000000000', '2012-11-26T00:00:00.000000000',
           '2012-11-27T00:00:00.000000000', '2012-11-28T00:00:00.000000000',
           '2012-11-29T00:00:00.000000000', '2012-11-30T00:00:00.000000000',
           '2012-12-01T00:00:00.000000000', '2012-12-02T00:00:00.000000000'],
          dtype='datetime64[ns]')

• Data variables: (1)
  • ssh
    (time, lat, lon)
    float32
    0.6549 0.6571 ... -0.2152 -0.2174

    standard_name :

    sea_surface_height_above_geoid

    long_name :

    seasurfaceheight

    units :

    m

    online_operation :

    average

    interval_operation :

    40 s

    interval_write :

    1 h

    cell_methods :

    time: mean (interval: 40 s)

    actual_range :

    [-0.7419584 1.1676793]

    array([[[ 0.6548789 ,  0.6571335 ,  0.6592872 , ...,  0.5913065 ,
              0.5900717 ,  0.58882123],
            [ 0.6497373 ,  0.65211684,  0.65439916, ...,  0.59274167,
              0.5915033 ,  0.59024507],
            [ 0.6443701 ,  0.6468519 ,  0.6492449 , ...,  0.59413385,
              0.5928928 ,  0.5916273 ],
            ...,
            [-0.05823348, -0.05635568, -0.05449666, ..., -0.02695263,
             -0.02947656, -0.03209217],
            [-0.06014866, -0.05832849, -0.05648171, ..., -0.02862907,
             -0.03114507, -0.03376382],
            [-0.06201855, -0.06033655, -0.0585744 , ..., -0.03032789,
             -0.03282989, -0.03544487]],
    
           [[ 0.6614337 ,  0.6636907 ,  0.6658804 , ...,  0.5849471 ,
              0.5837169 ,  0.5824688 ],
            [ 0.6560296 ,  0.658439  ,  0.6607885 , ...,  0.5864821 ,
              0.58524257,  0.58398175],
            [ 0.6503755 ,  0.652916  ,  0.65540755, ...,  0.58799535,
              0.5867505 ,  0.5854811 ],
    ...
            [ 0.01582234,  0.01615571,  0.01645508, ..., -0.20990212,
             -0.2122489 , -0.21467723],
            [ 0.01512165,  0.01543645,  0.01570038, ..., -0.21022458,
             -0.21266674, -0.21517338],
            [ 0.01385609,  0.0141472 ,  0.01437084, ..., -0.21028937,
             -0.21278173, -0.21531929]],
    
           [[ 0.45650592,  0.4500986 ,  0.44400105, ...,  0.63276714,
              0.6335951 ,  0.6342172 ],
            [ 0.45620465,  0.44972265,  0.44354692, ...,  0.6348397 ,
              0.63553923,  0.63605505],
            [ 0.45608178,  0.44951597,  0.44325253, ...,  0.6368277 ,
              0.63741523,  0.6378475 ],
            ...,
            [ 0.0345847 ,  0.03474202,  0.03505557, ..., -0.21224682,
             -0.21446429, -0.21657763],
            [ 0.03450282,  0.03470787,  0.03492999, ..., -0.21276844,
             -0.21503675, -0.21722902],
            [ 0.03375479,  0.03404886,  0.0342739 , ..., -0.21291387,
             -0.21521048, -0.21744703]]], dtype=float32)

• Indexes: (3)
  • lon
    PandasIndex

    PandasIndex(Float64Index([ -64.98333333333419,  -64.96666666666752,  -64.95000000000086,
                   -64.93333333333419,  -64.91666666666752,  -64.90000000000086,
                   -64.88333333333419,  -64.86666666666753,  -64.85000000000086,
                    -64.8333333333342,
                  ...
                   -55.15000000000141,  -55.13333333333475,  -55.11666666666808,
                  -55.100000000001415,  -55.08333333333475, -55.066666666668084,
                   -55.05000000000142,  -55.03333333333475,  -55.01666666666809,
                   -55.00000000000142],
                 dtype='float64', name='lon', length=600))

  • lat
    PandasIndex

    PandasIndex(Float64Index([33.016666666666325,  33.03333333333299, 33.049999999999656,
                   33.06666666666632,  33.08333333333299,  33.09999999999965,
                   33.11666666666632, 33.133333333332985,  33.14999999999965,
                  33.166666666666316,
                  ...
                   42.84999999999911, 42.866666666665765, 42.883333333332416,
                  42.899999999999096, 42.916666666665776,  42.93333333333243,
                   42.94999999999908,  42.96666666666576,  42.98333333333244,
                   42.99999999999909],
                 dtype='float64', name='lat', length=600))

  • time
    PandasIndex

    PandasIndex(DatetimeIndex(['2012-10-22', '2012-10-23', '2012-10-24', '2012-10-25',
                   '2012-10-26', '2012-10-27', '2012-10-28', '2012-10-29',
                   '2012-10-30', '2012-10-31', '2012-11-01', '2012-11-02',
                   '2012-11-03', '2012-11-04', '2012-11-05', '2012-11-06',
                   '2012-11-07', '2012-11-08', '2012-11-09', '2012-11-10',
                   '2012-11-11', '2012-11-12', '2012-11-13', '2012-11-14',
                   '2012-11-15', '2012-11-16', '2012-11-17', '2012-11-18',
                   '2012-11-19', '2012-11-20', '2012-11-21', '2012-11-22',
                   '2012-11-23', '2012-11-24', '2012-11-25', '2012-11-26',
                   '2012-11-27', '2012-11-28', '2012-11-29', '2012-11-30',
                   '2012-12-01', '2012-12-02'],
                  dtype='datetime64[ns]', name='time', freq=None))

• Attributes: (2)

  Info :

  Horizontal grid read in regulargrid_NATL60.nc / Source field read in /home/cubelmann/Data/Data_natl60_hourly/NATL60-CJM165_y2012m10d01.1h_SSH.nc / Interpolation method: bilin

  About :

  Created by SOSIE interpolation environement => https://github.com/brodeau/sosie/

nadir4_config = OmegaConf.load(f"./configs/natl60_obs.yaml")
ds_nadir4 = hydra.utils.instantiate(nadir4_config.ALONGTRACK_NADIR4.data).compute()
ds_swot1nadir5 = hydra.utils.instantiate(
    nadir4_config.ALONGTRACK_SWOT1NADIR5.data
).compute()
ds_swot1nadir5

<xarray.Dataset>
Dimensions:  (time: 1003548)
Coordinates:
    lon      (time) float64 -55.03 -55.06 -55.1 -55.13 ... -59.01 -59.03 -59.05
    lat      (time) float64 39.58 39.53 39.47 39.42 ... 42.81 42.87 42.93 42.98
  * time     (time) datetime64[ns] 2012-10-22T11:16:43.687588 ... 2012-12-02T...
Data variables:
    ssh      (time) float64 0.9958 1.014 1.027 1.032 ... -0.112 -0.1122 -0.1118
Attributes: (12/26)
    description:               SWOT fixed grid
    corresponding_grid:        
    title:                     Altimeter like data simulated by SWOT simulator
    keywords:                  check keywords
    Conventions:               CF-1.6
    summary:                   SWOT grid data produced
    ...                        ...
    geospatial_lon_units:      degrees_east
    project:                   SWOT
    date_created:              2018-10-12T12:39:50Z
    date_modified:             2018-10-12T12:39:50Z
    keywords_vocabulary:       NASA
    references:                Gaultier, L., C. Ubelmann, and L.-L. Fu, 2016:...

xarray.Dataset

• Dimensions:
  • time: 1003548
• Coordinates: (3)
  • lon
    (time)
    float64
    -55.03 -55.06 ... -59.03 -59.05

    array([-55.026272, -55.060736, -55.095136, ..., -59.005632, -59.025856,
           -59.046112])

  • lat
    (time)
    float64
    39.58 39.53 39.47 ... 42.93 42.98

    units :

    degrees_west

    standard_name :

    latitude

    long_name :

    Latitude

    array([39.58196 , 39.528312, 39.474664, ..., 42.86994 , 42.927328,
           42.984708])

  • time
    (time)
    datetime64[ns]
    2012-10-22T11:16:43.687588 ... 2...

    units :

    seconds since 2012-10-01

    array(['2012-10-22T11:16:43.687588000', '2012-10-22T11:16:44.676358000',
           '2012-10-22T11:16:45.665127000', ..., '2012-12-02T17:07:40.576272000',
           '2012-12-02T17:07:41.894631000', '2012-12-02T17:07:42.553811000'],
          dtype='datetime64[ns]')

• Data variables: (1)
  • ssh
    (time)
    float64
    0.9958 1.014 ... -0.1122 -0.1118

    valid_min :

    -1500000

    valid_max :

    1500000

    units :

    m

    long_name :

    SSH interpolated from model

    array([ 0.9958,  1.0138,  1.0272, ..., -0.112 , -0.1122, -0.1118])

• Indexes: (1)
  • time
    PandasIndex

    PandasIndex(DatetimeIndex(['2012-10-22 11:16:43.687588', '2012-10-22 11:16:44.676358',
                   '2012-10-22 11:16:45.665127', '2012-10-22 11:16:46.653897',
                   '2012-10-22 11:16:47.642666', '2012-10-22 11:16:48.631436',
                   '2012-10-22 11:16:49.620206', '2012-10-22 11:16:50.608975',
                   '2012-10-22 11:16:51.597745', '2012-10-22 11:16:52.586514',
                   ...
                   '2012-12-02 17:07:33.984475', '2012-12-02 17:07:35.302834',
                   '2012-12-02 17:07:35.962014', '2012-12-02 17:07:36.621193',
                   '2012-12-02 17:07:37.939553', '2012-12-02 17:07:38.598732',
                   '2012-12-02 17:07:39.917092', '2012-12-02 17:07:40.576272',
                   '2012-12-02 17:07:41.894631', '2012-12-02 17:07:42.553811'],
                  dtype='datetime64[ns]', name='time', length=1003548, freq=None))

• Attributes: (26)

  description :

  SWOT fixed grid

  corresponding_grid :

  title :

  Altimeter like data simulated by SWOT simulator

  keywords :

  check keywords

  Conventions :

  CF-1.6

  summary :

  SWOT grid data produced

  Metadata_Conventions :

  Unidata Dataset Discovery v1.0

  history :

  Grid File created by swotsimulator version 2.21

  processing_level :

  L2

  standard_name_vocabulary :

  CF-1.6

  creator_name :

  Lucile Gaultier and Clement Ubelmann

  creator_email :

  lucile.gaultier@gmail.com

  publisher_url :

  github/SWOTSimulator/

  time_coverage_start :

  0.3383046

  time_coverage_end :

  9.895123

  geospatial_lat_min :

  26.46

  geospatial_lat_max :

  66.15

  geospatial_lat_units :

  degrees_north

  geospatial_lon_max :

  360.00

  geospatial_lon_min :

  0.00

  geospatial_lon_units :

  degrees_east

  project :

  SWOT

  date_created :

  2018-10-12T12:39:50Z

  date_modified :

  2018-10-12T12:39:50Z

  keywords_vocabulary :

  NASA

  references :

  Gaultier, L., C. Ubelmann, and L.-L. Fu, 2016: The Challenge of Using Future SWOT Data for Oceanic Field Reconstruction. J. Atmos. Oceanic Technol., 33, 119-126, doi:10.1175/jtech-d-15-0160.1. http://dx.doi.org/10.1175/JTECH-D-15-0160.1.

Pixel Density - NATL60 | NADIR | SWOT

kwargs = {
    "stat": "density",
    "cumulative": False,
    "common_norm": True,
    "fill": False,
    "element": "poly",
    "linewidth": 3,
    "alpha": 0.9,
}
fig, ax = plt.subplots(figsize=(10, 6))
sns.histplot(
    data=ds_natl60.ssh.values.flatten(),
    ax=ax,
    label="NATL60",
    color="black",
    zorder=1,
    **kwargs,
)
sns.histplot(
    data=ds_nadir4.ssh.values.flatten(),
    ax=ax,
    label="NADIR4",
    color="black",
    linestyle=":",
    **kwargs
)
sns.histplot(
    data=ds_swot1nadir5.ssh.values.flatten(),
    ax=ax,
    label="SWOT1NADIR5",
    color="black",
    linestyle="--",
    **kwargs,
)


ax.set_xlabel("Sea Surface Height [m]")
ax.set_ylabel("Density")
plt.tight_layout()
plt.legend()
fig.savefig(f"./figures/dc20a/densities/natl60_obs_density.png")
plt.show()

Cumulative Density

kwargs["cumulative"] = True

fig, ax = plt.subplots(figsize=(10, 6))
sns.histplot(
    data=ds_natl60.ssh.values.flatten(),
    ax=ax,
    label="NATL60",
    color="black",
    zorder=1,
    **kwargs,
)
sns.histplot(
    data=ds_nadir4.ssh.values.flatten(),
    ax=ax,
    label="NADIR4",
    color="black",
    linestyle=":",
    **kwargs
)
sns.histplot(
    data=ds_swot1nadir5.ssh.values.flatten(),
    ax=ax,
    label="SWOT1NADIR5",
    color="black",
    linestyle="--",
    **kwargs,
)

ax.set_xlabel("Sea Surface Height [m]")
# ax.set_xlabel(r"Log Kinetic Energy [m$^2$s$^{-2}$]")
# ax.set_xlabel(r"Log Enstropy [s$^{-1}$]")
ax.set_ylabel("Cumulative Density")
plt.tight_layout()
plt.legend()
fig.savefig(f"./figures/dc20a/densities/natl60_obs_cdensity.png")
plt.show()

Regrdding: AlongTrack -> Uniform Grid

AlongTrack -> Uniform Grid

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

ds_natl60 = hydra.utils.instantiate(psd_config.fill_nans)(ds_natl60.pint.dequantify())

Coarsend Versions

ds_natl60 = ds_natl60.coarsen({"lon": 3, "lat": 3}).mean()
ds_natl60

<xarray.Dataset>
Dimensions:  (time: 42, lat: 200, lon: 200)
Coordinates:
  * lon      (lon) float64 -64.97 -64.92 -64.87 -64.82 ... -55.12 -55.07 -55.02
  * lat      (lat) float64 33.03 33.08 33.13 33.18 ... 42.83 42.88 42.93 42.98
  * time     (time) datetime64[ns] 2012-10-22 2012-10-23 ... 2012-12-02
Data variables:
    ssh      (time, lat, lon) float32 0.652 0.6585 0.6642 ... -0.2079 -0.2149
Attributes:
    Info:     Horizontal grid read in regulargrid_NATL60.nc / Source field re...
    About:    Created by SOSIE interpolation environement => https://github.c...

xarray.Dataset

• Dimensions:
  • time: 42
  • lat: 200
  • lon: 200
• Coordinates: (3)
  • lon
    (lon)
    float64
    -64.97 -64.92 ... -55.07 -55.02

    array([-64.966667, -64.916667, -64.866667, -64.816667, -64.766667, -64.716667,
           -64.666667, -64.616667, -64.566667, -64.516667, -64.466667, -64.416667,
           -64.366667, -64.316667, -64.266667, -64.216667, -64.166667, -64.116667,
           -64.066667, -64.016667, -63.966667, -63.916667, -63.866667, -63.816667,
           -63.766667, -63.716667, -63.666667, -63.616667, -63.566667, -63.516667,
           -63.466667, -63.416667, -63.366667, -63.316667, -63.266667, -63.216667,
           -63.166667, -63.116667, -63.066667, -63.016667, -62.966667, -62.916667,
           -62.866667, -62.816667, -62.766667, -62.716667, -62.666667, -62.616667,
           -62.566667, -62.516667, -62.466667, -62.416667, -62.366667, -62.316667,
           -62.266667, -62.216667, -62.166667, -62.116667, -62.066667, -62.016667,
           -61.966667, -61.916667, -61.866667, -61.816667, -61.766667, -61.716667,
           -61.666667, -61.616667, -61.566667, -61.516667, -61.466667, -61.416667,
           -61.366667, -61.316667, -61.266667, -61.216667, -61.166667, -61.116667,
           -61.066667, -61.016667, -60.966667, -60.916667, -60.866667, -60.816667,
           -60.766667, -60.716667, -60.666667, -60.616667, -60.566667, -60.516667,
           -60.466667, -60.416667, -60.366667, -60.316667, -60.266667, -60.216667,
           -60.166667, -60.116667, -60.066667, -60.016667, -59.966667, -59.916667,
           -59.866667, -59.816667, -59.766667, -59.716667, -59.666667, -59.616667,
           -59.566667, -59.516667, -59.466667, -59.416667, -59.366667, -59.316667,
           -59.266667, -59.216667, -59.166667, -59.116667, -59.066667, -59.016667,
           -58.966667, -58.916667, -58.866667, -58.816667, -58.766667, -58.716667,
           -58.666667, -58.616667, -58.566667, -58.516667, -58.466667, -58.416667,
           -58.366667, -58.316667, -58.266667, -58.216667, -58.166667, -58.116667,
           -58.066667, -58.016667, -57.966667, -57.916667, -57.866667, -57.816667,
           -57.766667, -57.716667, -57.666667, -57.616667, -57.566667, -57.516667,
           -57.466667, -57.416667, -57.366667, -57.316667, -57.266667, -57.216667,
           -57.166667, -57.116667, -57.066667, -57.016667, -56.966667, -56.916667,
           -56.866667, -56.816667, -56.766667, -56.716667, -56.666667, -56.616667,
           -56.566667, -56.516667, -56.466667, -56.416667, -56.366667, -56.316667,
           -56.266667, -56.216667, -56.166667, -56.116667, -56.066667, -56.016667,
           -55.966667, -55.916667, -55.866667, -55.816667, -55.766667, -55.716667,
           -55.666667, -55.616667, -55.566667, -55.516667, -55.466667, -55.416667,
           -55.366667, -55.316667, -55.266667, -55.216667, -55.166667, -55.116667,
           -55.066667, -55.016667])

  • lat
    (lat)
    float64
    33.03 33.08 33.13 ... 42.93 42.98

    standard_name :

    latitude

    long_name :

    Latitude

    units :

    degrees_west

    array([33.033333, 33.083333, 33.133333, 33.183333, 33.233333, 33.283333,
           33.333333, 33.383333, 33.433333, 33.483333, 33.533333, 33.583333,
           33.633333, 33.683333, 33.733333, 33.783333, 33.833333, 33.883333,
           33.933333, 33.983333, 34.033333, 34.083333, 34.133333, 34.183333,
           34.233333, 34.283333, 34.333333, 34.383333, 34.433333, 34.483333,
           34.533333, 34.583333, 34.633333, 34.683333, 34.733333, 34.783333,
           34.833333, 34.883333, 34.933333, 34.983333, 35.033333, 35.083333,
           35.133333, 35.183333, 35.233333, 35.283333, 35.333333, 35.383333,
           35.433333, 35.483333, 35.533333, 35.583333, 35.633333, 35.683333,
           35.733333, 35.783333, 35.833333, 35.883333, 35.933333, 35.983333,
           36.033333, 36.083333, 36.133333, 36.183333, 36.233333, 36.283333,
           36.333333, 36.383333, 36.433333, 36.483333, 36.533333, 36.583333,
           36.633333, 36.683333, 36.733333, 36.783333, 36.833333, 36.883333,
           36.933333, 36.983333, 37.033333, 37.083333, 37.133333, 37.183333,
           37.233333, 37.283333, 37.333333, 37.383333, 37.433333, 37.483333,
           37.533333, 37.583333, 37.633333, 37.683333, 37.733333, 37.783333,
           37.833333, 37.883333, 37.933333, 37.983333, 38.033333, 38.083333,
           38.133333, 38.183333, 38.233333, 38.283333, 38.333333, 38.383333,
           38.433333, 38.483333, 38.533333, 38.583333, 38.633333, 38.683333,
           38.733333, 38.783333, 38.833333, 38.883333, 38.933333, 38.983333,
           39.033333, 39.083333, 39.133333, 39.183333, 39.233333, 39.283333,
           39.333333, 39.383333, 39.433333, 39.483333, 39.533333, 39.583333,
           39.633333, 39.683333, 39.733333, 39.783333, 39.833333, 39.883333,
           39.933333, 39.983333, 40.033333, 40.083333, 40.133333, 40.183333,
           40.233333, 40.283333, 40.333333, 40.383333, 40.433333, 40.483333,
           40.533333, 40.583333, 40.633333, 40.683333, 40.733333, 40.783333,
           40.833333, 40.883333, 40.933333, 40.983333, 41.033333, 41.083333,
           41.133333, 41.183333, 41.233333, 41.283333, 41.333333, 41.383333,
           41.433333, 41.483333, 41.533333, 41.583333, 41.633333, 41.683333,
           41.733333, 41.783333, 41.833333, 41.883333, 41.933333, 41.983333,
           42.033333, 42.083333, 42.133333, 42.183333, 42.233333, 42.283333,
           42.333333, 42.383333, 42.433333, 42.483333, 42.533333, 42.583333,
           42.633333, 42.683333, 42.733333, 42.783333, 42.833333, 42.883333,
           42.933333, 42.983333])

  • time
    (time)
    datetime64[ns]
    2012-10-22 ... 2012-12-02

    array(['2012-10-22T00:00:00.000000000', '2012-10-23T00:00:00.000000000',
           '2012-10-24T00:00:00.000000000', '2012-10-25T00:00:00.000000000',
           '2012-10-26T00:00:00.000000000', '2012-10-27T00:00:00.000000000',
           '2012-10-28T00:00:00.000000000', '2012-10-29T00:00:00.000000000',
           '2012-10-30T00:00:00.000000000', '2012-10-31T00:00:00.000000000',
           '2012-11-01T00:00:00.000000000', '2012-11-02T00:00:00.000000000',
           '2012-11-03T00:00:00.000000000', '2012-11-04T00:00:00.000000000',
           '2012-11-05T00:00:00.000000000', '2012-11-06T00:00:00.000000000',
           '2012-11-07T00:00:00.000000000', '2012-11-08T00:00:00.000000000',
           '2012-11-09T00:00:00.000000000', '2012-11-10T00:00:00.000000000',
           '2012-11-11T00:00:00.000000000', '2012-11-12T00:00:00.000000000',
           '2012-11-13T00:00:00.000000000', '2012-11-14T00:00:00.000000000',
           '2012-11-15T00:00:00.000000000', '2012-11-16T00:00:00.000000000',
           '2012-11-17T00:00:00.000000000', '2012-11-18T00:00:00.000000000',
           '2012-11-19T00:00:00.000000000', '2012-11-20T00:00:00.000000000',
           '2012-11-21T00:00:00.000000000', '2012-11-22T00:00:00.000000000',
           '2012-11-23T00:00:00.000000000', '2012-11-24T00:00:00.000000000',
           '2012-11-25T00:00:00.000000000', '2012-11-26T00:00:00.000000000',
           '2012-11-27T00:00:00.000000000', '2012-11-28T00:00:00.000000000',
           '2012-11-29T00:00:00.000000000', '2012-11-30T00:00:00.000000000',
           '2012-12-01T00:00:00.000000000', '2012-12-02T00:00:00.000000000'],
          dtype='datetime64[ns]')

• Data variables: (1)
  • ssh
    (time, lat, lon)
    float32
    0.652 0.6585 ... -0.2079 -0.2149

    standard_name :

    sea_surface_height_above_geoid

    long_name :

    seasurfaceheight

    online_operation :

    average

    interval_operation :

    40 s

    interval_write :

    1 h

    cell_methods :

    time: mean (interval: 40 s)

    actual_range :

    [-0.7419584 1.1676793]

    units :

    m

    array([[[ 0.6520022 ,  0.6585398 ,  0.6642114 , ...,  0.59855586,
              0.59513116,  0.5914826 ],
            [ 0.63575274,  0.64306694,  0.6496819 , ...,  0.6025136 ,
              0.5991565 ,  0.5955087 ],
            [ 0.6190673 ,  0.62632686,  0.633281  , ...,  0.60601234,
              0.60273045,  0.5990826 ],
            ...,
            [-0.04787888, -0.04444524, -0.04109241, ..., -0.007847  ,
             -0.01461816, -0.02203192],
            [-0.05263177, -0.0483911 , -0.04548378, ..., -0.0121701 ,
             -0.01889833, -0.02627038],
            [-0.05833047, -0.05322452, -0.04955636, ..., -0.01737751,
             -0.02393646, -0.03118466]],
    
           [[ 0.658329  ,  0.66514313,  0.67125154, ...,  0.59218127,
              0.58885807,  0.5852296 ],
            [ 0.6411973 ,  0.6490578 ,  0.65647966, ...,  0.5965517 ,
              0.5933058 ,  0.5896903 ],
            [ 0.623477  ,  0.63121647,  0.6391002 , ...,  0.6004862 ,
              0.5973787 ,  0.5938805 ],
    ...
            [ 0.01493218,  0.01390165,  0.01254615, ..., -0.19566639,
             -0.20232889, -0.20773965],
            [ 0.01641853,  0.01650789,  0.01353683, ..., -0.19861259,
             -0.20470469, -0.21073315],
            [ 0.01522953,  0.01566913,  0.01434487, ..., -0.19968782,
             -0.20571834, -0.21258704]],
    
           [[ 0.44988114,  0.43149403,  0.41169986, ...,  0.6280331 ,
              0.6326583 ,  0.63545597],
            [ 0.44985878,  0.4309694 ,  0.41143242, ...,  0.63547236,
              0.63921857,  0.6410519 ],
            [ 0.4518173 ,  0.43222544,  0.41316524, ...,  0.64268017,
              0.6453408 ,  0.6468438 ],
            ...,
            [ 0.03271373,  0.02987434,  0.02491577, ..., -0.19762646,
             -0.20238733, -0.20735066],
            [ 0.03445117,  0.03456031,  0.03055872, ..., -0.19978192,
             -0.20572636, -0.21213359],
            [ 0.03451117,  0.03512124,  0.03456491, ..., -0.20084554,
             -0.20787542, -0.21487716]]], dtype=float32)

• Indexes: (3)
  • lon
    PandasIndex

    PandasIndex(Float64Index([  -64.9666666666675,  -64.91666666666752,  -64.86666666666753,
                   -64.81666666666753,  -64.76666666666752,  -64.71666666666754,
                   -64.66666666666754,  -64.61666666666754,  -64.56666666666753,
                   -64.51666666666755,
                  ...
                   -55.46666666666806,  -55.41666666666807,  -55.36666666666807,
                   -55.31666666666806,  -55.26666666666807,  -55.21666666666808,
                   -55.16666666666808, -55.116666666668074, -55.066666666668084,
                  -55.016666666668094],
                 dtype='float64', name='lon', length=200))

  • lat
    PandasIndex

    PandasIndex(Float64Index([ 33.03333333333299,  33.08333333333299, 33.133333333332985,
                   33.18333333333298,  33.23333333333298, 33.283333333332976,
                   33.33333333333297,  33.38333333333297,  33.43333333333297,
                  33.483333333332965,
                  ...
                   42.53333333333245, 42.583333333332455, 42.633333333332445,
                  42.683333333332435,  42.73333333333244,  42.78333333333244,
                   42.83333333333243,  42.88333333333242,  42.93333333333243,
                   42.98333333333243],
                 dtype='float64', name='lat', length=200))

  • time
    PandasIndex

    PandasIndex(DatetimeIndex(['2012-10-22', '2012-10-23', '2012-10-24', '2012-10-25',
                   '2012-10-26', '2012-10-27', '2012-10-28', '2012-10-29',
                   '2012-10-30', '2012-10-31', '2012-11-01', '2012-11-02',
                   '2012-11-03', '2012-11-04', '2012-11-05', '2012-11-06',
                   '2012-11-07', '2012-11-08', '2012-11-09', '2012-11-10',
                   '2012-11-11', '2012-11-12', '2012-11-13', '2012-11-14',
                   '2012-11-15', '2012-11-16', '2012-11-17', '2012-11-18',
                   '2012-11-19', '2012-11-20', '2012-11-21', '2012-11-22',
                   '2012-11-23', '2012-11-24', '2012-11-25', '2012-11-26',
                   '2012-11-27', '2012-11-28', '2012-11-29', '2012-11-30',
                   '2012-12-01', '2012-12-02'],
                  dtype='datetime64[ns]', name='time', freq=None))

• Attributes: (2)

  Info :

  Horizontal grid read in regulargrid_NATL60.nc / Source field read in /home/cubelmann/Data/Data_natl60_hourly/NATL60-CJM165_y2012m10d01.1h_SSH.nc / Interpolation method: bilin

  About :

  Created by SOSIE interpolation environement => https://github.com/brodeau/sosie/

Prediction Datasets

%%time

# load config

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()

CPU times: user 184 ms, sys: 59.8 ms, total: 243 ms
Wall time: 514 ms

!ls /gpfswork/rech/yrf/commun/data_challenges/dc20a_osse/staging/ml_ready/

nadir1.nc  nadir4.nc  nadir5.nc  swot1nadir5.nc  swot1.nc  swot.nc

Regrdding

Uniform Grid –> Uniform 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,
)

CPU times: user 12.2 s, sys: 77.7 ms, total: 12.2 s
Wall time: 12.3 s

Preprocess Chain

%%time

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

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

CPU times: user 1.05 s, sys: 2.01 ms, total: 1.05 s
Wall time: 1.06 s

Geophysical Variables

def calculate_physical_quantities(da):
    da["ssh"] *= units.meters
    da["lon"] = da.lon * units.degrees
    da["lat"] = da.lat * units.degrees
    da = geocalc.streamfunction(da, "ssh")
    da = geocalc.geostrophic_velocities(da, variable="psi")
    da = geocalc.kinetic_energy(da, variables=["u", "v"])
    da = geocalc.divergence(da, variables=["u", "v"])
    da = geocalc.coriolis_normalized(da, "div")
    da = geocalc.relative_vorticity(da, variables=["u", "v"])
    da = geocalc.coriolis_normalized(da, "vort_r")
    da = geocalc.strain_magnitude(da, variables=["u", "v"])
    da = geocalc.coriolis_normalized(da, variable="strain")
    return da

%%time

ds_natl60 = calculate_physical_quantities(correct_labels(ds_natl60).pint.dequantify())
ds_duacs = calculate_physical_quantities(correct_labels(ds_duacs).pint.dequantify())
ds_miost = calculate_physical_quantities(correct_labels(ds_miost).pint.dequantify())
ds_nerf_siren = calculate_physical_quantities(
    correct_labels(ds_nerf_siren).pint.dequantify()
)
ds_nerf_ffn = calculate_physical_quantities(
    correct_labels(ds_nerf_ffn).pint.dequantify()
)
ds_nerf_mlp = calculate_physical_quantities(
    correct_labels(ds_nerf_mlp).pint.dequantify()
)

CPU times: user 3.89 s, sys: 253 ms, total: 4.14 s
Wall time: 4.16 s

Sea Surface Height

kwargs = {
    "stat": "density",
    "cumulative": False,
    "common_norm": True,
    "fill": False,
    "element": "poly",
    "linewidth": 3,
    "alpha": 0.75,
}

fig, ax = plt.subplots(figsize=(10, 6))
sns.histplot(
    data=ds_natl60.ssh.values.flatten(),
    ax=ax,
    label="NATL60",
    color="black",
    zorder=1,
    **kwargs,
)
sns.histplot(
    data=ds_duacs.ssh.values.flatten(),
    ax=ax,
    label="DUACS",
    color="tab:green",
    linestyle="-",
    **kwargs
)
sns.histplot(
    data=ds_miost.ssh.values.flatten(),
    ax=ax,
    label="MIOST",
    color="tab:red",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_mlp.ssh.values.flatten(),
    ax=ax,
    label="NERF (MLP)",
    color="tab:olive",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_ffn.ssh.values.flatten(),
    ax=ax,
    label="NERF (FFN)",
    color="tab:blue",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_siren.ssh.values.flatten(),
    ax=ax,
    label="NERF (SIREN)",
    color="tab:cyan",
    linestyle="-",
    **kwargs
)


ax.set_xlabel("Sea Surface Height [m]")
ax.set_ylabel("Density")
plt.tight_layout()
plt.legend(fontsize=14, framealpha=0.5)
fig.savefig(f"./figures/dc20a/densities/dc20a_density_ssh_{experiment}.png")
plt.close()

kwargs["cumulative"] = True

fig, ax = plt.subplots(figsize=(10, 6))
sns.histplot(
    data=ds_natl60.ssh.values.flatten(),
    ax=ax,
    label="NATL60",
    color="black",
    zorder=1,
    **kwargs,
)
sns.histplot(
    data=ds_duacs.ssh.values.flatten(),
    ax=ax,
    label="DUACS",
    color="tab:green",
    linestyle="-",
    **kwargs
)
sns.histplot(
    data=ds_miost.ssh.values.flatten(),
    ax=ax,
    label="MIOST",
    color="tab:red",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_mlp.ssh.values.flatten(),
    ax=ax,
    label="NERF (MLP)",
    color="tab:olive",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_ffn.ssh.values.flatten(),
    ax=ax,
    label="NERF (FFN)",
    color="tab:blue",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_siren.ssh.values.flatten(),
    ax=ax,
    label="NERF (SIREN)",
    color="tab:cyan",
    linestyle="-",
    **kwargs
)


ax.set_xlabel("Sea Surface Height [m]")
ax.set_ylabel("Cumulative Density")
plt.tight_layout()
plt.legend(fontsize=14, framealpha=0.5)
fig.savefig(f"./figures/dc20a/densities/dc20a_cdensity_ssh_{experiment}.png")
plt.close()

Kinetic Energy

kwargs = {
    "stat": "density",
    "cumulative": False,
    "common_norm": True,
    "fill": False,
    "element": "poly",
    "linewidth": 3,
    "alpha": 0.75,
    "log_scale": True,
}

fig, ax = plt.subplots(figsize=(10, 6))
sns.histplot(
    data=ds_natl60.ke.values.flatten(),
    ax=ax,
    label="NATL60",
    color="black",
    zorder=1,
    **kwargs,
)
sns.histplot(
    data=ds_duacs.ke.values.flatten(),
    ax=ax,
    label="DUACS",
    color="tab:green",
    linestyle="-",
    **kwargs
)
sns.histplot(
    data=ds_miost.ke.values.flatten(),
    ax=ax,
    label="MIOST",
    color="tab:red",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_mlp.ke.values.flatten(),
    ax=ax,
    label="NERF (MLP)",
    color="tab:olive",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_ffn.ke.values.flatten(),
    ax=ax,
    label="NERF (FFN)",
    color="tab:blue",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_siren.ke.values.flatten(),
    ax=ax,
    label="NERF (SIREN)",
    color="tab:cyan",
    linestyle="-",
    **kwargs
)


ax.set_xlabel("Kinetic Energy")
ax.set_ylabel("Density")
plt.tight_layout()
plt.legend(fontsize=14, framealpha=0.5)
fig.savefig(f"./figures/dc20a/densities/dc20a_density_ke_{experiment}.png")
plt.close()

kwargs["cumulative"] = True

fig, ax = plt.subplots(figsize=(10, 6))
sns.histplot(
    data=ds_natl60.ke.values.flatten(),
    ax=ax,
    label="NATL60",
    color="black",
    zorder=1,
    **kwargs,
)
sns.histplot(
    data=ds_duacs.ke.values.flatten(),
    ax=ax,
    label="DUACS",
    color="tab:green",
    linestyle="-",
    **kwargs
)
sns.histplot(
    data=ds_miost.ke.values.flatten(),
    ax=ax,
    label="MIOST",
    color="tab:red",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_mlp.ke.values.flatten(),
    ax=ax,
    label="NERF (MLP)",
    color="tab:olive",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_ffn.ke.values.flatten(),
    ax=ax,
    label="NERF (FFN)",
    color="tab:blue",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_siren.ke.values.flatten(),
    ax=ax,
    label="NERF (SIREN)",
    color="tab:cyan",
    linestyle="-",
    **kwargs
)


ax.set_xlabel("Kinetic Energy")
ax.set_ylabel("Cumulative Density")
plt.tight_layout()
plt.legend(fontsize=14, framealpha=0.5)
fig.savefig(f"./figures/dc20a/densities/dc20a_cdensity_cke_{experiment}.png")
plt.close()

Relative Vorticity

kwargs = {
    "stat": "density",
    "cumulative": False,
    "common_norm": True,
    "fill": False,
    "element": "poly",
    "linewidth": 8,
    "alpha": 0.8,
    "log_scale": False,
}

fig, ax = plt.subplots(figsize=(10, 6))
sns.histplot(
    data=ds_natl60.vort_r.values.flatten(),
    ax=ax,
    label="NATL60",
    color="black",
    zorder=1,
    **kwargs,
)


kwargs["alpha"] = 0.9
kwargs["linewidth"] = 3
sns.histplot(
    data=ds_duacs.vort_r.values.flatten(),
    ax=ax,
    label="DUACS",
    color="tab:green",
    linestyle="-",
    **kwargs
)
sns.histplot(
    data=ds_miost.vort_r.values.flatten(),
    ax=ax,
    label="MIOST",
    color="tab:red",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_mlp.vort_r.values.flatten(),
    ax=ax,
    label="NERF (MLP)",
    color="tab:olive",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_ffn.vort_r.values.flatten(),
    ax=ax,
    label="NERF (FFN)",
    color="tab:blue",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_siren.vort_r.values.flatten(),
    ax=ax,
    label="NERF (SIREN)",
    color="tab:cyan",
    linestyle="-",
    **kwargs
)


ax.set_xlabel("Normalized Relative Vorticity")
ax.set_ylabel("Density")
plt.tight_layout()
plt.legend(fontsize=14, framealpha=0.5)
fig.savefig(f"./figures/dc20a/densities/dc20a_density_vort_r_{experiment}.png")
plt.close()

kwargs = {
    "stat": "density",
    "cumulative": False,
    "common_norm": True,
    "fill": False,
    "element": "poly",
    "linewidth": 8,
    "alpha": 0.8,
    "log_scale": False,
}

fig, ax = plt.subplots(figsize=(10, 6))
sns.histplot(
    data=ds_natl60.vort_r.values.flatten(),
    ax=ax,
    label="NATL60",
    color="black",
    zorder=1,
    **kwargs,
)


kwargs["alpha"] = 0.9
kwargs["linewidth"] = 3
sns.histplot(
    data=ds_duacs.vort_r.values.flatten(),
    ax=ax,
    label="DUACS",
    color="tab:green",
    linestyle="-",
    **kwargs
)
sns.histplot(
    data=ds_miost.vort_r.values.flatten(),
    ax=ax,
    label="MIOST",
    color="tab:red",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_mlp.vort_r.values.flatten(),
    ax=ax,
    label="NERF (MLP)",
    color="tab:olive",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_ffn.vort_r.values.flatten(),
    ax=ax,
    label="NERF (FFN)",
    color="tab:blue",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_siren.vort_r.values.flatten(),
    ax=ax,
    label="NERF (SIREN)",
    color="tab:cyan",
    linestyle="-",
    **kwargs
)

ax.set(xscale="symlog")

ax.set_xlabel("Normalized Relative Vorticity")
ax.set_ylabel("Density")
plt.tight_layout()
plt.legend(fontsize=14, framealpha=0.5)
fig.savefig(f"./figures/dc20a/densities/dc20a_density_vort_r_symlog_{experiment}.png")
plt.close()

kwargs["cumulative"] = True

fig, ax = plt.subplots(figsize=(10, 6))
sns.histplot(
    data=ds_natl60.vort_r.values.flatten(),
    ax=ax,
    label="NATL60",
    color="black",
    zorder=1,
    **kwargs,
)
sns.histplot(
    data=ds_duacs.vort_r.values.flatten(),
    ax=ax,
    label="DUACS",
    color="tab:green",
    linestyle="-",
    **kwargs
)
sns.histplot(
    data=ds_miost.vort_r.values.flatten(),
    ax=ax,
    label="MIOST",
    color="tab:red",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_mlp.vort_r.values.flatten(),
    ax=ax,
    label="NERF (MLP)",
    color="tab:olive",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_ffn.vort_r.values.flatten(),
    ax=ax,
    label="NERF (FFN)",
    color="tab:blue",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_siren.vort_r.values.flatten(),
    ax=ax,
    label="NERF (SIREN)",
    color="tab:cyan",
    linestyle="-",
    **kwargs
)


ax.set_xlabel("Normalized Relative Vorticity")
ax.set_ylabel("Density")
plt.tight_layout()
plt.legend(fontsize=14, framealpha=0.5)
fig.savefig(f"./figures/dc20a/densities/dc20a_cdensity_vort_r_{experiment}.png")
plt.close()

kwargs["cumulative"] = True

fig, ax = plt.subplots(figsize=(10, 6))
sns.histplot(
    data=ds_natl60.vort_r.values.flatten(),
    ax=ax,
    label="NATL60",
    color="black",
    zorder=1,
    **kwargs,
)
sns.histplot(
    data=ds_duacs.vort_r.values.flatten(),
    ax=ax,
    label="DUACS",
    color="tab:green",
    linestyle="-",
    **kwargs
)
sns.histplot(
    data=ds_miost.vort_r.values.flatten(),
    ax=ax,
    label="MIOST",
    color="tab:red",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_mlp.vort_r.values.flatten(),
    ax=ax,
    label="NERF (MLP)",
    color="tab:olive",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_ffn.vort_r.values.flatten(),
    ax=ax,
    label="NERF (FFN)",
    color="tab:blue",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_siren.vort_r.values.flatten(),
    ax=ax,
    label="NERF (SIREN)",
    color="tab:cyan",
    linestyle="-",
    **kwargs
)

ax.set(xscale="symlog")

ax.set_xlabel("Normalized Relative Vorticity")
ax.set_ylabel("Density")
plt.tight_layout()
plt.legend(fontsize=14, framealpha=0.5)
fig.savefig(f"./figures/dc20a/densities/dc20a_cdensity_vort_r_symlog_{experiment}.png")
plt.close()

Strain

kwargs = {
    "stat": "density",
    "cumulative": False,
    "common_norm": True,
    "fill": False,
    "element": "poly",
    "linewidth": 5,
    "alpha": 0.8,
    "log_scale": True,
}

fig, ax = plt.subplots(figsize=(10, 6))
sns.histplot(
    data=ds_natl60.strain.values.flatten(),
    ax=ax,
    label="NATL60",
    color="black",
    zorder=1,
    **kwargs,
)


kwargs["alpha"] = 0.9
kwargs["linewidth"] = 3
sns.histplot(
    data=ds_duacs.strain.values.flatten(),
    ax=ax,
    label="DUACS",
    color="tab:green",
    linestyle="-",
    **kwargs
)
sns.histplot(
    data=ds_miost.strain.values.flatten(),
    ax=ax,
    label="MIOST",
    color="tab:red",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_mlp.strain.values.flatten(),
    ax=ax,
    label="NERF (MLP)",
    color="tab:olive",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_ffn.strain.values.flatten(),
    ax=ax,
    label="NERF (FFN)",
    color="tab:blue",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_siren.strain.values.flatten(),
    ax=ax,
    label="NERF (SIREN)",
    color="tab:cyan",
    linestyle="-",
    **kwargs
)


ax.set_xlabel("Strain")
ax.set_ylabel("Density")
plt.tight_layout()
plt.legend(fontsize=14, framealpha=0.5)
fig.savefig(f"./figures/dc20a/densities/dc20a_density_strain_{experiment}.png")
plt.close()

kwargs = {
    "stat": "density",
    "cumulative": True,
    "common_norm": True,
    "fill": False,
    "element": "poly",
    "linewidth": 5,
    "alpha": 0.8,
    "log_scale": True,
}

fig, ax = plt.subplots(figsize=(10, 6))
sns.histplot(
    data=ds_natl60.strain.values.flatten(),
    ax=ax,
    label="NATL60",
    color="black",
    zorder=1,
    **kwargs,
)


kwargs["alpha"] = 0.9
kwargs["linewidth"] = 3
sns.histplot(
    data=ds_duacs.strain.values.flatten(),
    ax=ax,
    label="DUACS",
    color="tab:green",
    linestyle="-",
    **kwargs
)
sns.histplot(
    data=ds_miost.strain.values.flatten(),
    ax=ax,
    label="MIOST",
    color="tab:red",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_mlp.strain.values.flatten(),
    ax=ax,
    label="NERF (MLP)",
    color="tab:olive",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_ffn.strain.values.flatten(),
    ax=ax,
    label="NERF (FFN)",
    color="tab:blue",
    linestyle="-",
    **kwargs
)

sns.histplot(
    data=ds_nerf_siren.strain.values.flatten(),
    ax=ax,
    label="NERF (SIREN)",
    color="tab:cyan",
    linestyle="-",
    **kwargs
)


ax.set_xlabel("Strain")
ax.set_ylabel("Density")
plt.tight_layout()
plt.legend(fontsize=14, framealpha=0.5)
fig.savefig(f"./figures/dc20a/densities/dc20a_cdensity_strain_{experiment}.png")
plt.close()

var_names = [
    "Sea Surface Height", "Kinetic Energy", "Relative Vorticity", "Strain"
]


def pixel_stack(ds):
    
    pixel_stack = np.vstack([
        ds.ssh.mean(dim="time").values.ravel(),
        np.log(ds.ke.mean(dim="time").values.ravel()),
        ds.vort_r.mean(dim="time").values.ravel(),
        np.sqrt(ds.strain.mean(dim="time").values.ravel())
    ]).T
    
    return pd.DataFrame(data=pixel_stack, columns=var_names)

NATL60

df = pixel_stack(ds_natl60)

# g = sns.PairGrid(df.sample(5_000))
g = sns.PairGrid(df.sample(25_000), corner=True)
# g.map_upper(sns.scatterplot, size=0.1)
g.map_lower(sns.kdeplot, levels=4, gridsize=50, color="black")
g.map_lower(sns.histplot, bins=50, pthresh=0.1, color="black")
# g.map_diag(sns.kdeplot, lw=3, legend=False, levels=4, gridsize=50)
g.map_diag(sns.histplot, bins=50, pthresh=0.1, color="black")

g.savefig(f"./figures/dc20a/densities/dc20a_natl60_jointdensity_{experiment}.png")
plt.close()

DUACS

df = pixel_stack(ds_duacs)

# g = sns.PairGrid(df.sample(5_000))
g = sns.PairGrid(df.sample(25_000), corner=True)
# g.map_upper(sns.scatterplot, size=0.1)
g.map_lower(sns.kdeplot, levels=4, gridsize=50, color="tab:green")
g.map_lower(sns.histplot, bins=50, pthresh=0.1, color="tab:green")
# g.map_diag(sns.kdeplot, lw=3, legend=False, levels=4, gridsize=50)
g.map_diag(sns.histplot, bins=50, pthresh=0.1, color="tab:green")

g.savefig(f"./figures/dc20a/densities/dc20a_duacs_jointdensity_{experiment}.png")
plt.close()

MIOST

df = pixel_stack(ds_miost)

# g = sns.PairGrid(df.sample(5_000))
g = sns.PairGrid(df.sample(25_000), corner=True)
# g.map_upper(sns.scatterplot, size=0.1)
g.map_lower(sns.kdeplot, levels=4, gridsize=50, color="tab:red")
g.map_lower(sns.histplot, bins=50, pthresh=0.1, color="tab:red")
# g.map_diag(sns.kdeplot, lw=3, legend=False, levels=4, gridsize=50)
g.map_diag(sns.histplot, bins=50, pthresh=0.1, color="tab:red")

g.savefig(f"./figures/dc20a/densities/dc20a_miost_jointdensity_{experiment}.png")
plt.close()

NERF (MLP)

df = pixel_stack(ds_nerf_mlp)

# g = sns.PairGrid(df.sample(5_000))
g = sns.PairGrid(df.sample(25_000), corner=True)
# g.map_upper(sns.scatterplot, size=0.1)
g.map_lower(sns.kdeplot, levels=4, gridsize=50, color="tab:olive")
g.map_lower(sns.histplot, bins=50, pthresh=.1, color="tab:olive")
# g.map_diag(sns.kdeplot, lw=3, legend=False, levels=4, gridsize=50)
g.map_diag(sns.histplot, bins=50, pthresh=.1, color="tab:olive")

g.savefig(f"./figures/dc20a/densities/dc20a_nerf_mlp_jointdensity_{experiment}.png")
plt.close()

NERF (FFN)

df = pixel_stack(ds_nerf_ffn)

# g = sns.PairGrid(df.sample(5_000))
g = sns.PairGrid(df.sample(25_000), corner=True)
# g.map_upper(sns.scatterplot, size=0.1)
g.map_lower(sns.kdeplot, levels=4, gridsize=50, color="tab:blue")
g.map_lower(sns.histplot, bins=50, pthresh=0.1, color="tab:blue")
# g.map_diag(sns.kdeplot, lw=3, legend=False, levels=4, gridsize=50)
g.map_diag(sns.histplot, bins=50, pthresh=0.1, color="tab:blue")

g.savefig(f"./figures/dc20a/densities/dc20a_nerf_ffn_jointdensity_{experiment}.png")
plt.close()

NERF (SIREN)

df = pixel_stack(ds_nerf_siren)

# g = sns.PairGrid(df.sample(5_000))
g = sns.PairGrid(df.sample(25_000), corner=True)
# g.map_upper(sns.scatterplot, size=0.1)
g.map_lower(sns.kdeplot, levels=4, gridsize=50, color="tab:cyan", )
g.map_lower(sns.histplot, bins=50, pthresh=.1, color="tab:cyan")
# g.map_diag(sns.kdeplot, lw=3, legend=False, levels=4, gridsize=50)
g.map_diag(sns.histplot, bins=50, pthresh=.1, color="tab:cyan")

g.savefig(f"./figures/dc20a/densities/dc20a_nerf_siren_jointdensity_{experiment}.png")
plt.close()