Preserving NaN missing values with safeguards¶
In this example, we compare how three lossy compressors (ZFP, SZ3 [v3.2 and v3.3], and SPERR) handle missing values that are encoded as NaNs. We compress a single day (four 6h observations) of sea-only satellite observation data, where data is missing on land (static) and wherever is not currently covered by the satellites (dynamic). Next, we apply safeguards to guarantee that NaN values are preserved. We also compare the safeguards with the compressor configuration auto-tuner OptZConfig.
Since the SPERR codec does not support NaN values, we replace them with the data mean before encoding with SPERR.
Copied!
import ssl
ssl._create_default_https_context = ssl._create_stdlib_context
import ssl
ssl._create_default_https_context = ssl._create_stdlib_context
Copied!
import copy
from pathlib import Path
import earthkit.plots
import humanize
import matplotlib as mpl
import numcodecs
import numpy as np
import pandas as pd
import xarray as xr
from matplotlib import patheffects as PathEffects
from matplotlib import pyplot as plt
from numcodecs_combinators.stack import CodecStack
from numcodecs_replace import ReplaceFilterCodec
from numcodecs_safeguards import SafeguardedCodec
from numcodecs_wasm_pressio import Pressio
from numcodecs_wasm_sperr import Sperr
from numcodecs_wasm_sz3 import Sz3 as _Sz3_3
from numcodecs_wasm_zfp import Zfp
from numcodecs_wasm_zstd import Zstd
from numcodecs_zero import ZeroCodec
import copy
from pathlib import Path
import earthkit.plots
import humanize
import matplotlib as mpl
import numcodecs
import numpy as np
import pandas as pd
import xarray as xr
from matplotlib import patheffects as PathEffects
from matplotlib import pyplot as plt
from numcodecs_combinators.stack import CodecStack
from numcodecs_replace import ReplaceFilterCodec
from numcodecs_safeguards import SafeguardedCodec
from numcodecs_wasm_pressio import Pressio
from numcodecs_wasm_sperr import Sperr
from numcodecs_wasm_sz3 import Sz3 as _Sz3_3
from numcodecs_wasm_zfp import Zfp
from numcodecs_wasm_zstd import Zstd
from numcodecs_zero import ZeroCodec
Copied!
# we want to test both SZ3 v3.2 and SZ3 v3.3
# but pip doesn't allow installing two different versions of the same package
# so we instead import the SZ3 v3.2 package directly from its PyPi wheel
import shutil
import sys
import tempfile
from urllib.request import urlopen
from zipfile import ZipFile
numcodecs_wasm_sz3_2_url = (
"https://files.pythonhosted.org/packages/a0/c4/"
+ "8e0cdbda2ffd5a591bf31c15e274984e065dfbfff27a725acf481c2a2bde/"
+ "numcodecs_wasm_sz3-0.7.0-py3-none-any.whl"
)
with tempfile.TemporaryDirectory() as tmp:
numcodecs_wasm_sz3_2_wheel = Path(tmp) / Path(numcodecs_wasm_sz3_2_url).name
with urlopen(numcodecs_wasm_sz3_2_url) as response:
with numcodecs_wasm_sz3_2_wheel.open("wb") as wheel:
shutil.copyfileobj(response, wheel)
with ZipFile(numcodecs_wasm_sz3_2_wheel, "r") as wheel:
numcodecs_wasm_sz3_2_source = wheel.extractall(tmp)
(Path(tmp) / "numcodecs_wasm_sz3").rename(
Path(tmp) / "numcodecs_wasm_sz3_2"
)
sys.path.append(tmp)
from numcodecs_wasm_sz3_2 import Sz3 as _Sz3_2
sys.path.remove(tmp)
# we want to test both SZ3 v3.2 and SZ3 v3.3
# but pip doesn't allow installing two different versions of the same package
# so we instead import the SZ3 v3.2 package directly from its PyPi wheel
import shutil
import sys
import tempfile
from urllib.request import urlopen
from zipfile import ZipFile
numcodecs_wasm_sz3_2_url = (
"https://files.pythonhosted.org/packages/a0/c4/"
+ "8e0cdbda2ffd5a591bf31c15e274984e065dfbfff27a725acf481c2a2bde/"
+ "numcodecs_wasm_sz3-0.7.0-py3-none-any.whl"
)
with tempfile.TemporaryDirectory() as tmp:
numcodecs_wasm_sz3_2_wheel = Path(tmp) / Path(numcodecs_wasm_sz3_2_url).name
with urlopen(numcodecs_wasm_sz3_2_url) as response:
with numcodecs_wasm_sz3_2_wheel.open("wb") as wheel:
shutil.copyfileobj(response, wheel)
with ZipFile(numcodecs_wasm_sz3_2_wheel, "r") as wheel:
numcodecs_wasm_sz3_2_source = wheel.extractall(tmp)
(Path(tmp) / "numcodecs_wasm_sz3").rename(
Path(tmp) / "numcodecs_wasm_sz3_2"
)
sys.path.append(tmp)
from numcodecs_wasm_sz3_2 import Sz3 as _Sz3_2
sys.path.remove(tmp)
Copied!
class Sz3_2(_Sz3_2):
codec_id = "sz3.2.rs"
def get_config(self):
return {**super().get_config(), "id": type(self).codec_id}
class Sz3_3(_Sz3_3):
codec_id = "sz3.3.rs"
def get_config(self):
return {**super().get_config(), "id": type(self).codec_id}
numcodecs.registry.register_codec(Sz3_2)
numcodecs.registry.register_codec(Sz3_3)
class Sz3_2(_Sz3_2):
codec_id = "sz3.2.rs"
def get_config(self):
return {**super().get_config(), "id": type(self).codec_id}
class Sz3_3(_Sz3_3):
codec_id = "sz3.3.rs"
def get_config(self):
return {**super().get_config(), "id": type(self).codec_id}
numcodecs.registry.register_codec(Sz3_2)
numcodecs.registry.register_codec(Sz3_3)
Copied!
wvpa = xr.open_dataset(
Path() / "data" / "hoaps-c.r30.h06.wvpa.2020-08" / "data.nc",
engine="netcdf4",
decode_timedelta=True,
).wvpa.sel(time="2020-08-02")
wvpa.sizes
wvpa = xr.open_dataset(
Path() / "data" / "hoaps-c.r30.h06.wvpa.2020-08" / "data.nc",
engine="netcdf4",
decode_timedelta=True,
).wvpa.sel(time="2020-08-02")
wvpa.sizes
Frozen({'time': 4, 'lat': 320, 'lon': 720})
Copied!
from compression_safeguards.utils.cast import as_bits
def compute_corrections_percentage(my_wvpa: np.ndarray, orig_wvpa: np.ndarray) -> float:
return np.mean(
~((my_wvpa == orig_wvpa) | (np.isnan(my_wvpa) & np.isnan(orig_wvpa)))
)
def compute_corrections_percentage_nan(
my_wvpa: np.ndarray, orig_wvpa: np.ndarray
) -> float:
return np.mean(as_bits(my_wvpa.values) != as_bits(orig_wvpa.values))
from compression_safeguards.utils.cast import as_bits
def compute_corrections_percentage(my_wvpa: np.ndarray, orig_wvpa: np.ndarray) -> float:
return np.mean(
~((my_wvpa == orig_wvpa) | (np.isnan(my_wvpa) & np.isnan(orig_wvpa)))
)
def compute_corrections_percentage_nan(
my_wvpa: np.ndarray, orig_wvpa: np.ndarray
) -> float:
return np.mean(as_bits(my_wvpa.values) != as_bits(orig_wvpa.values))
Copied!
import observe
observations = []
import observe
observations = []
Lossless compression¶
We first compress the data losslessly with ZStandard at level 22, which gives maximum compression, to provide a baseline.
Copied!
zstd = Zstd(level=22)
with observe.observe(zstd, observations):
wvpa_zstd_enc = zstd.encode(wvpa.values)
wvpa_zstd = wvpa.copy(deep=True, data=zstd.decode(wvpa_zstd_enc))
wvpa_zstd_cr = wvpa.nbytes / np.asarray(wvpa_zstd_enc).nbytes
zstd = Zstd(level=22)
with observe.observe(zstd, observations):
wvpa_zstd_enc = zstd.encode(wvpa.values)
wvpa_zstd = wvpa.copy(deep=True, data=zstd.decode(wvpa_zstd_enc))
wvpa_zstd_cr = wvpa.nbytes / np.asarray(wvpa_zstd_enc).nbytes
Compressing water vapour path with lossy compressors¶
We configure each compressor with an absolute error bound of 0.1 kg/m^2.
Copied!
eb_abs = 0.1
eb_abs = 0.1
Copied!
zfp = Zfp(mode="fixed-accuracy", tolerance=eb_abs, non_finite="allow-unsafe")
with observe.observe(zfp, observations):
wvpa_zfp_enc = zfp.encode(wvpa.values)
wvpa_zfp = wvpa.copy(deep=True, data=zfp.decode(wvpa_zfp_enc))
wvpa_zfp_cr = wvpa.nbytes / wvpa_zfp_enc.nbytes
zfp = Zfp(mode="fixed-accuracy", tolerance=eb_abs, non_finite="allow-unsafe")
with observe.observe(zfp, observations):
wvpa_zfp_enc = zfp.encode(wvpa.values)
wvpa_zfp = wvpa.copy(deep=True, data=zfp.decode(wvpa_zfp_enc))
wvpa_zfp_cr = wvpa.nbytes / wvpa_zfp_enc.nbytes
Copied!
sz3_2 = Sz3_2(eb_mode="abs", eb_abs=eb_abs)
with observe.observe(sz3_2, observations):
wvpa_sz3_2_enc = sz3_2.encode(wvpa.values)
wvpa_sz3_2 = wvpa.copy(deep=True, data=sz3_2.decode(wvpa_sz3_2_enc))
wvpa_sz3_2_cr = wvpa.nbytes / np.asarray(wvpa_sz3_2_enc).nbytes
sz3_2 = Sz3_2(eb_mode="abs", eb_abs=eb_abs)
with observe.observe(sz3_2, observations):
wvpa_sz3_2_enc = sz3_2.encode(wvpa.values)
wvpa_sz3_2 = wvpa.copy(deep=True, data=sz3_2.decode(wvpa_sz3_2_enc))
wvpa_sz3_2_cr = wvpa.nbytes / np.asarray(wvpa_sz3_2_enc).nbytes
Copied!
sz3_3 = Sz3_3(eb_mode="abs", eb_abs=eb_abs)
with observe.observe(sz3_3, observations):
wvpa_sz3_3_enc = sz3_3.encode(wvpa.values)
wvpa_sz3_3 = wvpa.copy(deep=True, data=sz3_3.decode(wvpa_sz3_3_enc))
wvpa_sz3_3_cr = wvpa.nbytes / np.asarray(wvpa_sz3_3_enc).nbytes
sz3_3 = Sz3_3(eb_mode="abs", eb_abs=eb_abs)
with observe.observe(sz3_3, observations):
wvpa_sz3_3_enc = sz3_3.encode(wvpa.values)
wvpa_sz3_3 = wvpa.copy(deep=True, data=sz3_3.decode(wvpa_sz3_3_enc))
wvpa_sz3_3_cr = wvpa.nbytes / np.asarray(wvpa_sz3_3_enc).nbytes
Copied!
# inspired by H5Z-SPERR's treatment of NaN values:
# https://github.com/NCAR/H5Z-SPERR/blob/72ebcb00e382886c229c5ef5a7e237fe451d5fb8/src/h5z-sperr.c#L464-L473
# https://github.com/NCAR/H5Z-SPERR/blob/72ebcb00e382886c229c5ef5a7e237fe451d5fb8/src/h5zsperr_helper.cpp#L179-L212
sperr = CodecStack(
ReplaceFilterCodec(replacements={np.nan: "nan_mean"}), Sperr(mode="pwe", pwe=eb_abs)
)
with observe.observe(sperr, observations):
wvpa_sperr_enc = sperr.encode(wvpa.values)
wvpa_sperr = wvpa.copy(deep=True, data=sperr.decode(wvpa_sperr_enc))
wvpa_sperr_cr = wvpa.nbytes / np.asarray(wvpa_sperr_enc).nbytes
# inspired by H5Z-SPERR's treatment of NaN values:
# https://github.com/NCAR/H5Z-SPERR/blob/72ebcb00e382886c229c5ef5a7e237fe451d5fb8/src/h5z-sperr.c#L464-L473
# https://github.com/NCAR/H5Z-SPERR/blob/72ebcb00e382886c229c5ef5a7e237fe451d5fb8/src/h5zsperr_helper.cpp#L179-L212
sperr = CodecStack(
ReplaceFilterCodec(replacements={np.nan: "nan_mean"}), Sperr(mode="pwe", pwe=eb_abs)
)
with observe.observe(sperr, observations):
wvpa_sperr_enc = sperr.encode(wvpa.values)
wvpa_sperr = wvpa.copy(deep=True, data=sperr.decode(wvpa_sperr_enc))
wvpa_sperr_cr = wvpa.nbytes / np.asarray(wvpa_sperr_enc).nbytes
Copied!
zero = ZeroCodec()
with observe.observe(zero, observations):
wvpa_zero_enc = zero.encode(wvpa.values)
wvpa_zero = wvpa.copy(deep=True, data=zero.decode(wvpa_zero_enc))
zero = ZeroCodec()
with observe.observe(zero, observations):
wvpa_zero_enc = zero.encode(wvpa.values)
wvpa_zero = wvpa.copy(deep=True, data=zero.decode(wvpa_zero_enc))
Compressing water vapour path with the safeguarded lossy compressors¶
We configure the safeguards with an absolute error bound of 0.1 kg/m^2. We then compare (a) considering NaN values only as equal if their bit patterns match and (b) considering all NaN values as equal.
Copied!
wvpa_sg = dict()
wvpa_sg_cr = dict()
for codec_id, codec in {
zero.codec_id: zero,
zfp.codec_id: zfp,
"sz3.2.rs": sz3_2,
"sz3.3.rs": sz3_3,
"sperr.rs": sperr,
}.items():
sg = SafeguardedCodec(
codec=codec, safeguards=[dict(kind="eb", type="abs", eb=eb_abs, equal_nan=True)]
)
with observe.observe(sg, observations):
wvpa_sg_enc = sg.encode(wvpa.values)
wvpa_sg[codec_id] = wvpa.copy(deep=True, data=sg.decode(wvpa_sg_enc))
wvpa_sg_cr[codec_id] = wvpa.nbytes / np.asarray(wvpa_sg_enc).nbytes
wvpa_sg = dict()
wvpa_sg_cr = dict()
for codec_id, codec in {
zero.codec_id: zero,
zfp.codec_id: zfp,
"sz3.2.rs": sz3_2,
"sz3.3.rs": sz3_3,
"sperr.rs": sperr,
}.items():
sg = SafeguardedCodec(
codec=codec, safeguards=[dict(kind="eb", type="abs", eb=eb_abs, equal_nan=True)]
)
with observe.observe(sg, observations):
wvpa_sg_enc = sg.encode(wvpa.values)
wvpa_sg[codec_id] = wvpa.copy(deep=True, data=sg.decode(wvpa_sg_enc))
wvpa_sg_cr[codec_id] = wvpa.nbytes / np.asarray(wvpa_sg_enc).nbytes
Copied!
wvpa_sg_nan = dict()
wvpa_sg_nan_cr = dict()
for codec_id, codec in {
zero.codec_id: zero,
zfp.codec_id: zfp,
"sz3.2.rs": sz3_2,
"sz3.3.rs": sz3_3,
"sperr.rs": sperr,
}.items():
sg_nan = SafeguardedCodec(
codec=codec,
safeguards=[dict(kind="eb", type="abs", eb=eb_abs, equal_nan=False)],
)
with observe.observe(sg_nan, observations):
wvpa_sg_nan_enc = sg_nan.encode(wvpa.values)
wvpa_sg_nan[codec_id] = wvpa.copy(
deep=True, data=sg_nan.decode(wvpa_sg_nan_enc)
)
wvpa_sg_nan_cr[codec_id] = wvpa.nbytes / np.asarray(wvpa_sg_nan_enc).nbytes
wvpa_sg_nan = dict()
wvpa_sg_nan_cr = dict()
for codec_id, codec in {
zero.codec_id: zero,
zfp.codec_id: zfp,
"sz3.2.rs": sz3_2,
"sz3.3.rs": sz3_3,
"sperr.rs": sperr,
}.items():
sg_nan = SafeguardedCodec(
codec=codec,
safeguards=[dict(kind="eb", type="abs", eb=eb_abs, equal_nan=False)],
)
with observe.observe(sg_nan, observations):
wvpa_sg_nan_enc = sg_nan.encode(wvpa.values)
wvpa_sg_nan[codec_id] = wvpa.copy(
deep=True, data=sg_nan.decode(wvpa_sg_nan_enc)
)
wvpa_sg_nan_cr[codec_id] = wvpa.nbytes / np.asarray(wvpa_sg_nan_enc).nbytes
Copied!
wvpa_sg_lossless = dict()
wvpa_sg_lossless_cr = dict()
for codec_id, codec in {
zero.codec_id: zero,
zfp.codec_id: zfp,
"sz3.2.rs": sz3_2,
"sz3.3.rs": sz3_3,
"sperr.rs": sperr,
}.items():
sg = SafeguardedCodec(
codec=codec,
safeguards=[dict(kind="eb", type="abs", eb=eb_abs, equal_nan=True)],
# produce lossless corrections and refine them with iteration
compute=dict(unstable_iterative=True, unstable_lossless_corrections=True),
)
with observe.observe(sg, observations):
wvpa_sg_lossless_enc = sg.encode(wvpa.values)
wvpa_sg_lossless[codec_id] = wvpa.copy(
deep=True, data=sg.decode(wvpa_sg_lossless_enc)
)
wvpa_sg_lossless_cr[codec_id] = (
wvpa.nbytes / np.asarray(wvpa_sg_lossless_enc).nbytes
)
wvpa_sg_lossless = dict()
wvpa_sg_lossless_cr = dict()
for codec_id, codec in {
zero.codec_id: zero,
zfp.codec_id: zfp,
"sz3.2.rs": sz3_2,
"sz3.3.rs": sz3_3,
"sperr.rs": sperr,
}.items():
sg = SafeguardedCodec(
codec=codec,
safeguards=[dict(kind="eb", type="abs", eb=eb_abs, equal_nan=True)],
# produce lossless corrections and refine them with iteration
compute=dict(unstable_iterative=True, unstable_lossless_corrections=True),
)
with observe.observe(sg, observations):
wvpa_sg_lossless_enc = sg.encode(wvpa.values)
wvpa_sg_lossless[codec_id] = wvpa.copy(
deep=True, data=sg.decode(wvpa_sg_lossless_enc)
)
wvpa_sg_lossless_cr[codec_id] = (
wvpa.nbytes / np.asarray(wvpa_sg_lossless_enc).nbytes
)
Copied!
wvpa_sg_nan_lossless = dict()
wvpa_sg_nan_lossless_cr = dict()
for codec_id, codec in {
zero.codec_id: zero,
zfp.codec_id: zfp,
"sz3.2.rs": sz3_2,
"sz3.3.rs": sz3_3,
"sperr.rs": sperr,
}.items():
sg_nan = SafeguardedCodec(
codec=codec,
safeguards=[dict(kind="eb", type="abs", eb=eb_abs, equal_nan=False)],
# produce lossless corrections and refine them with iteration
compute=dict(unstable_iterative=True, unstable_lossless_corrections=True),
)
with observe.observe(sg_nan, observations):
wvpa_sg_nan_lossless_enc = sg_nan.encode(wvpa.values)
wvpa_sg_nan_lossless[codec_id] = wvpa.copy(
deep=True, data=sg_nan.decode(wvpa_sg_nan_lossless_enc)
)
wvpa_sg_nan_lossless_cr[codec_id] = (
wvpa.nbytes / np.asarray(wvpa_sg_nan_lossless_enc).nbytes
)
wvpa_sg_nan_lossless = dict()
wvpa_sg_nan_lossless_cr = dict()
for codec_id, codec in {
zero.codec_id: zero,
zfp.codec_id: zfp,
"sz3.2.rs": sz3_2,
"sz3.3.rs": sz3_3,
"sperr.rs": sperr,
}.items():
sg_nan = SafeguardedCodec(
codec=codec,
safeguards=[dict(kind="eb", type="abs", eb=eb_abs, equal_nan=False)],
# produce lossless corrections and refine them with iteration
compute=dict(unstable_iterative=True, unstable_lossless_corrections=True),
)
with observe.observe(sg_nan, observations):
wvpa_sg_nan_lossless_enc = sg_nan.encode(wvpa.values)
wvpa_sg_nan_lossless[codec_id] = wvpa.copy(
deep=True, data=sg_nan.decode(wvpa_sg_nan_lossless_enc)
)
wvpa_sg_nan_lossless_cr[codec_id] = (
wvpa.nbytes / np.asarray(wvpa_sg_nan_lossless_enc).nbytes
)
Compressing water vapour path with OptZConfig¶
We configure OptZConfig with a custom safety violations metric, implemented in Python, that computes the percentage of violations $V$. We then maximise the score
$$ \textrm{score} = \begin{cases} -\textrm{V} \quad &\text{if } \textrm{V} > 0 \\ \textrm{CR} \quad &\text{otherwise} \end{cases} $$
using the FRAZ search algorithm with 25 iterations, where CR is the achieved compression ratio. Since FRAZ seems to struggle with finding sufficient absolute error bounds spread across several orders of magnitude, we search for bounds in logarithmic space by wrapping each codec in an Exponential<CODEC> meta-codec.
Copied!
class SafetyViolationsMetric(numcodecs.abc.Codec):
codec_id = "safety-violations-metric"
def __init__(self):
self._data = None
def encode(self, buf):
# store the original data for later
self._data = np.array(buf, copy=True)
# return no metric
return np.empty(0, dtype=np.float64)
def decode(self, buf, out=None):
# compute the violations
violations = np.mean(
~(
(np.abs(buf - self._data) <= eb_abs)
| (np.isnan(buf) & np.isnan(self._data))
)
)
self._data = None
# return the violations score metric
return numcodecs.compat.ndarray_copy(np.float64(violations), out)
numcodecs.registry.register_codec(SafetyViolationsMetric)
class SafetyViolationsMetric(numcodecs.abc.Codec):
codec_id = "safety-violations-metric"
def __init__(self):
self._data = None
def encode(self, buf):
# store the original data for later
self._data = np.array(buf, copy=True)
# return no metric
return np.empty(0, dtype=np.float64)
def decode(self, buf, out=None):
# compute the violations
violations = np.mean(
~(
(np.abs(buf - self._data) <= eb_abs)
| (np.isnan(buf) & np.isnan(self._data))
)
)
self._data = None
# return the violations score metric
return numcodecs.compat.ndarray_copy(np.float64(violations), out)
numcodecs.registry.register_codec(SafetyViolationsMetric)
Copied!
class ExponentialZfp(Zfp):
codec_id = "e-zfp.rs"
def __new__(cls, tolerance: float, **kwargs):
codec = super().__new__(cls, tolerance=np.exp(tolerance), **kwargs)
codec._tolerance = tolerance
return codec
def get_config(self):
return {
**super().get_config(),
"id": type(self).codec_id,
"tolerance": self._tolerance,
}
class ExponentialSz3_2(Sz3_2):
codec_id = "e-sz3.2.rs"
def __new__(cls, eb_abs: float, **kwargs):
codec = super().__new__(cls, eb_abs=np.exp(eb_abs), **kwargs)
codec._eb_abs = eb_abs
return codec
def get_config(self):
return {
**super().get_config(),
"id": type(self).codec_id,
"eb_abs": self._eb_abs,
}
class ExponentialSz3_3(Sz3_3):
codec_id = "e-sz3.3.rs"
def __new__(cls, eb_abs: float, **kwargs):
codec = super().__new__(cls, eb_abs=np.exp(eb_abs), **kwargs)
codec._eb_abs = eb_abs
return codec
def get_config(self):
return {
**super().get_config(),
"id": type(self).codec_id,
"eb_abs": self._eb_abs,
}
# libpressio does not support the configuration format used by the CodecStack
# and ReplaceFilterCodec, so we inline their functionality here
class ExponentialMaskedSperr(Sperr):
codec_id = "e-sperr.rs"
def __new__(cls, pwe: float, **kwargs):
codec = super().__new__(cls, pwe=np.exp(pwe), **kwargs)
codec._pwe = pwe
codec._filter = ReplaceFilterCodec(replacements={np.nan: "nan_mean"})
return codec
def encode(self, buf):
return super().encode(self._filter.encode(buf))
def decode(self, buf, out=None):
return self._filter.decode(super().decode(buf, out=out), out=out)
def get_config(self):
return {**super().get_config(), "id": type(self).codec_id, "pwe": self._pwe}
numcodecs.registry.register_codec(ExponentialZfp)
numcodecs.registry.register_codec(ExponentialSz3_2)
numcodecs.registry.register_codec(ExponentialSz3_3)
numcodecs.registry.register_codec(ExponentialMaskedSperr)
class ExponentialZfp(Zfp):
codec_id = "e-zfp.rs"
def __new__(cls, tolerance: float, **kwargs):
codec = super().__new__(cls, tolerance=np.exp(tolerance), **kwargs)
codec._tolerance = tolerance
return codec
def get_config(self):
return {
**super().get_config(),
"id": type(self).codec_id,
"tolerance": self._tolerance,
}
class ExponentialSz3_2(Sz3_2):
codec_id = "e-sz3.2.rs"
def __new__(cls, eb_abs: float, **kwargs):
codec = super().__new__(cls, eb_abs=np.exp(eb_abs), **kwargs)
codec._eb_abs = eb_abs
return codec
def get_config(self):
return {
**super().get_config(),
"id": type(self).codec_id,
"eb_abs": self._eb_abs,
}
class ExponentialSz3_3(Sz3_3):
codec_id = "e-sz3.3.rs"
def __new__(cls, eb_abs: float, **kwargs):
codec = super().__new__(cls, eb_abs=np.exp(eb_abs), **kwargs)
codec._eb_abs = eb_abs
return codec
def get_config(self):
return {
**super().get_config(),
"id": type(self).codec_id,
"eb_abs": self._eb_abs,
}
# libpressio does not support the configuration format used by the CodecStack
# and ReplaceFilterCodec, so we inline their functionality here
class ExponentialMaskedSperr(Sperr):
codec_id = "e-sperr.rs"
def __new__(cls, pwe: float, **kwargs):
codec = super().__new__(cls, pwe=np.exp(pwe), **kwargs)
codec._pwe = pwe
codec._filter = ReplaceFilterCodec(replacements={np.nan: "nan_mean"})
return codec
def encode(self, buf):
return super().encode(self._filter.encode(buf))
def decode(self, buf, out=None):
return self._filter.decode(super().decode(buf, out=out), out=out)
def get_config(self):
return {**super().get_config(), "id": type(self).codec_id, "pwe": self._pwe}
numcodecs.registry.register_codec(ExponentialZfp)
numcodecs.registry.register_codec(ExponentialSz3_2)
numcodecs.registry.register_codec(ExponentialSz3_3)
numcodecs.registry.register_codec(ExponentialMaskedSperr)
Copied!
wvpa_optzconfig = dict()
wvpa_optzconfig_cr = dict()
for codec_id, (codec, parameter, lower_bound) in {
zfp.codec_id: (zfp, "tolerance", 1e-12), # tiny bound
"sz3.2.rs": (sz3_2, "eb_abs", 1e-12), # tiny bound
"sz3.3.rs": (sz3_3, "eb_abs", 1e-3), # decent guess
"sperr.rs": (sperr[-1], "pwe", 1e-12), # decent guess
}.items():
optzconfig = Pressio(
compressor_id="opt",
compressor_config={
"opt:output": ["composite:score"],
"opt:inputs": [f"numcodecs.rs:{parameter}"],
"opt:lower_bound": np.log(lower_bound),
"opt:upper_bound": np.log(eb_abs),
"opt:max_iterations": 25,
"opt:objective_mode_name": "max",
},
early_config={
"opt:compressor": "pressio",
"pressio:compressor": "numcodecs.rs",
**{
f"numcodecs.rs:{k}": f"e-{v}" if k == "id" else v
for k, v in codec.get_config().items()
},
"opt:search": "fraz",
"pressio:metric": "composite",
"composite:plugins": ["size", "numcodecs.rs-metric"],
"composite:scripts": [
"""
violations = metrics["numcodecs.rs-metric:decompression"]
if violations > 0 then
return "score", -violations
else
return "score", metrics["size:compression_ratio"]
end
"""
],
"numcodecs.rs-metric:id": "safety-violations-metric",
},
)
with observe.observe(optzconfig, observations):
wvpa_optzconfig_enc = optzconfig.encode(wvpa.values)
wvpa_optzconfig[codec_id] = wvpa.copy(
deep=True, data=optzconfig.decode(wvpa_optzconfig_enc)
)
wvpa_optzconfig_cr[codec_id] = wvpa.nbytes / wvpa_optzconfig_enc.nbytes
wvpa_optzconfig = dict()
wvpa_optzconfig_cr = dict()
for codec_id, (codec, parameter, lower_bound) in {
zfp.codec_id: (zfp, "tolerance", 1e-12), # tiny bound
"sz3.2.rs": (sz3_2, "eb_abs", 1e-12), # tiny bound
"sz3.3.rs": (sz3_3, "eb_abs", 1e-3), # decent guess
"sperr.rs": (sperr[-1], "pwe", 1e-12), # decent guess
}.items():
optzconfig = Pressio(
compressor_id="opt",
compressor_config={
"opt:output": ["composite:score"],
"opt:inputs": [f"numcodecs.rs:{parameter}"],
"opt:lower_bound": np.log(lower_bound),
"opt:upper_bound": np.log(eb_abs),
"opt:max_iterations": 25,
"opt:objective_mode_name": "max",
},
early_config={
"opt:compressor": "pressio",
"pressio:compressor": "numcodecs.rs",
**{
f"numcodecs.rs:{k}": f"e-{v}" if k == "id" else v
for k, v in codec.get_config().items()
},
"opt:search": "fraz",
"pressio:metric": "composite",
"composite:plugins": ["size", "numcodecs.rs-metric"],
"composite:scripts": [
"""
violations = metrics["numcodecs.rs-metric:decompression"]
if violations > 0 then
return "score", -violations
else
return "score", metrics["size:compression_ratio"]
end
"""
],
"numcodecs.rs-metric:id": "safety-violations-metric",
},
)
with observe.observe(optzconfig, observations):
wvpa_optzconfig_enc = optzconfig.encode(wvpa.values)
wvpa_optzconfig[codec_id] = wvpa.copy(
deep=True, data=optzconfig.decode(wvpa_optzconfig_enc)
)
wvpa_optzconfig_cr[codec_id] = wvpa.nbytes / wvpa_optzconfig_enc.nbytes
rank={0,1,} iter={0} input={-14.9668,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={1} input={-21.7226,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={2} input={-8.31428,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={3} input={-22.2717,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={4} input={-4.03186,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={5} input={-25.9225,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={6} input={-4.82515,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={7} input={-4.32668,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={8} input={-15.2103,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={9} input={-23.6072,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={10} input={-7.07512,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={11} input={-15.055,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={12} input={-23.8995,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={13} input={-6.36479,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={14} input={-25.5114,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={15} input={-21.8949,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={16} input={-18.4519,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={17} input={-2.64642,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={18} input={-14.543,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={19} input={-4.47936,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={20} input={-27.4262,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={21} input={-23.5642,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={22} input={-6.99472,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={23} input={-25.5847,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={24} input={-3.05598,} output={-0.682795,} objective={-0.682795}
final_iter={25} inputs={-14.9668,} output={-0.682795,}
rank={0,1,} iter={0} input={-14.9668,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={1} input={-21.7226,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={2} input={-8.31428,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={3} input={-22.2717,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={4} input={-4.03186,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={5} input={-25.9225,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={6} input={-4.82515,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={7} input={-4.32668,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={8} input={-15.2103,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={9} input={-23.6072,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={10} input={-7.07512,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={11} input={-15.055,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={12} input={-23.8995,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={13} input={-6.36479,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={14} input={-25.5114,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={15} input={-21.8949,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={16} input={-18.4519,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={17} input={-2.64642,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={18} input={-14.543,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={19} input={-4.47936,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={20} input={-27.4262,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={21} input={-23.5642,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={22} input={-6.99472,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={23} input={-25.5847,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={24} input={-3.05598,} output={-0.682795,} objective={-0.682795}
final_iter={25} inputs={-14.9668,} output={-0.682795,}
rank={0,1,} iter={0} input={-4.60517,} output={7.15771,} objective={7.15771}
rank={0,1,} iter={1} input={-5.8335,} output={6.36508,} objective={6.36508}
rank={0,1,} iter={2} input={-3.39562,} output={8.08699,} objective={8.08699}
rank={0,1,} iter={3} input={-2.30259,} output={9.05992,} objective={9.05992}
rank={0,1,} iter={4} input={-5.93333,} output={6.30179,} objective={6.30179}
rank={0,1,} iter={5} input={-2.617,} output={8.74736,} objective={8.74736}
rank={0,1,} iter={6} input={-6.59711,} output={5.8773,} objective={5.8773}
rank={0,1,} iter={7} input={-2.76123,} output={8.62748,} objective={8.62748}
rank={0,1,} iter={8} input={-2.6706,} output={8.73302,} objective={8.73302}
rank={0,1,} iter={9} input={-2.31789,} output={9.05033,} objective={9.05033}
rank={0,1,} iter={10} input={-2.33051,} output={9.04183,} objective={9.04183}
rank={0,1,} iter={11} input={-2.307,} output={9.06157,} objective={9.06157}
rank={0,1,} iter={12} input={-2.30544,} output={9.06008,} objective={9.06008}
rank={0,1,} iter={13} input={-6.90776,} output={5.68644,} objective={5.68644}
rank={0,1,} iter={14} input={-2.30699,} output={9.06019,} objective={9.06019}
rank={0,1,} iter={15} input={-4.62298,} output={7.16393,} objective={7.16393}
rank={0,1,} iter={16} input={-2.91334,} output={8.48801,} objective={8.48801}
rank={0,1,} iter={17} input={-3.46499,} output={8.02305,} objective={8.02305}
rank={0,1,} iter={18} input={-2.59288,} output={8.77466,} objective={8.77466}
rank={0,1,} iter={19} input={-2.886,} output={8.50805,} objective={8.50805}
rank={0,1,} iter={20} input={-5.0018,} output={6.91005,} objective={6.91005}
rank={0,1,} iter={21} input={-2.5965,} output={8.76981,} objective={8.76981}
rank={0,1,} iter={22} input={-3.52318,} output={7.94491,} objective={7.94491}
rank={0,1,} iter={23} input={-2.45175,} output={8.90523,} objective={8.90523}
rank={0,1,} iter={24} input={-2.30792,} output={9.05867,} objective={9.05867}
final_iter={25} inputs={-2.307,} output={9.06157,}
rank={0,1,} iter={0} input={-14.9668,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={1} input={-21.7226,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={2} input={-8.31428,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={3} input={-22.2717,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={4} input={-4.03186,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={5} input={-25.9225,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={6} input={-4.82515,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={7} input={-4.32668,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={8} input={-15.2103,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={9} input={-23.6072,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={10} input={-7.07512,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={11} input={-15.055,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={12} input={-23.8995,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={13} input={-6.36479,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={14} input={-25.5114,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={15} input={-21.8949,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={16} input={-18.4519,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={17} input={-2.64642,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={18} input={-14.543,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={19} input={-4.47936,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={20} input={-27.4262,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={21} input={-23.5642,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={22} input={-6.99472,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={23} input={-25.5847,} output={-0.682795,} objective={-0.682795}
rank={0,1,} iter={24} input={-3.05598,} output={-0.682795,} objective={-0.682795}
final_iter={25} inputs={-14.9668,} output={-0.682795,}
Copied!
old_cmap_and_norm = earthkit.plots.styles.colors.cmap_and_norm
old_cmap_and_norm = earthkit.plots.styles.colors.cmap_and_norm
Copied!
def my_cmap_and_norm(colors, levels, normalize=True, extend=None, extend_levels=True):
return old_cmap_and_norm(colors, levels, normalize, extend, True)
earthkit.plots.styles.colors.cmap_and_norm = my_cmap_and_norm
def my_cmap_and_norm(colors, levels, normalize=True, extend=None, extend_levels=True):
return old_cmap_and_norm(colors, levels, normalize, extend, True)
earthkit.plots.styles.colors.cmap_and_norm = my_cmap_and_norm
Copied!
def plot_wvpa(
da: xr.DataArray,
cr,
chart,
title,
eb_abs,
error=False,
corr=None,
inset=True,
):
# compute the default style that earthkit.maps would apply
source = earthkit.plots.sources.XarraySource(da)
style = copy.deepcopy(
earthkit.plots.styles.auto.guess_style(
source,
units=source.units,
)
)
style._levels = earthkit.plots.styles.levels.Levels(np.linspace(2, 73, 22))
style._legend_kwargs["ticks"] = np.linspace(2, 73, 5)
extend_left = np.nanmin(da) < 2
extend_right = np.nanmax(da) > 73
extend = {
(False, False): "neither",
(True, False): "min",
(False, True): "max",
(True, True): "both",
}[(extend_left, extend_right)]
style._legend_kwargs["extend"] = extend
chart.ax.set_global()
chart.ax.fill_between(
[0, 1],
[1, 1],
hatch="X",
edgecolor="white",
facecolor="lightgrey",
transform=chart.ax.transAxes,
zorder=-12,
)
chart.pcolormesh(
da.sel(time="2020-08-02T06:00"), style=style, zorder=-11, rasterized=True
)
with plt.rc_context(
{
"hatch.color": "white",
"hatch.linewidth": 0.5,
}
):
chart.contourf(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape[-2:],
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape[-2:],
),
z=np.isnan(wvpa.sel(time="2020-08-02T06:00")).values,
colors=["none"],
levels=[-0.5, 0.9, 1.5],
hatches=[None, "X"],
legend_style=None,
zorder=-10,
)
if error:
if corr is not None:
da_hatch = da.copy(
data=(da.values.view(np.uint32) == corr.values.view(np.uint32))
)
da_corr = (~da_hatch).astype(float)
old_process_projection_requirements = (
chart.ax.get_figure()._process_projection_requirements
)
def _process_projection_requirements(
*, axes_class=None, polar=False, projection=None, **kwargs
):
if axes_class is not None and projection is not None:
return axes_class, dict(projection=projection, **kwargs)
return old_process_projection_requirements(
axes_class=axes_class, polar=polar, projection=projection, **kwargs
)
chart.ax.get_figure()._process_projection_requirements = (
_process_projection_requirements
)
axin = chart.ax.inset_axes(
[0.025, 0.05, 1 / 3, 1 / 3],
xticklabels=[],
yticklabels=[],
axes_class=type(chart.ax),
projection=chart.ax.projection,
)
axin.set_global()
axin.pcolormesh(
da_corr.lon.values,
da_corr.lat.values,
np.squeeze(da_corr.sel(time="2020-08-02T06:00").values),
cmap=mpl.colors.ListedColormap(["white", "green", "lawngreen"]),
vmin=0,
vmax=2,
rasterized=True,
)
axin.coastlines(color="#555555")
axin.spines["geo"].set_edgecolor("black")
axin.set_title(
"Corrections",
path_effects=[PathEffects.withStroke(linewidth=3, foreground="white")],
)
elif inset:
da_err = ~((np.abs(da - wvpa) <= eb_abs) | (np.isnan(da) & np.isnan(wvpa)))
old_process_projection_requirements = (
chart.ax.get_figure()._process_projection_requirements
)
def _process_projection_requirements(
*, axes_class=None, polar=False, projection=None, **kwargs
):
if axes_class is not None and projection is not None:
return axes_class, dict(projection=projection, **kwargs)
return old_process_projection_requirements(
axes_class=axes_class, polar=polar, projection=projection, **kwargs
)
chart.ax.get_figure()._process_projection_requirements = (
_process_projection_requirements
)
axin = chart.ax.inset_axes(
[0.025, 0.05, 1 / 3, 1 / 3],
xticklabels=[],
yticklabels=[],
axes_class=type(chart.ax),
projection=chart.ax.projection,
)
axin.set_global()
axin.pcolormesh(
da_err.lon.values,
da_err.lat.values,
np.squeeze(da_err.sel(time="2020-08-02T06:00").values),
cmap=mpl.colors.ListedColormap(["white", "red"]),
vmin=0,
vmax=1,
rasterized=True,
)
axin.coastlines(color="#555555")
axin.spines["geo"].set_edgecolor("black")
axin.set_title(
"Violations",
path_effects=[PathEffects.withStroke(linewidth=3, foreground="white")],
)
chart.title(title)
if error:
err_v = np.mean(
~((np.abs(da - wvpa) <= eb_abs) | (np.isnan(da) & np.isnan(wvpa)))
)
err_v = (
0
if err_v == 0
else np.format_float_positional(100 * err_v, precision=1, min_digits=1)
+ "%"
)
if err_v == "0.0%":
err_v = "<0.05%"
t = chart.ax.text(
0.95,
0.1,
f"V={err_v}",
ha="right",
va="bottom",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
t = chart.ax.text(
0.95,
0.9,
rf"$\times$ {np.round(cr, 2)}"
if error
else humanize.naturalsize(wvpa.nbytes, binary=True),
ha="right",
va="top",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
chart.contour(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape[-2:],
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape[-2:],
),
z=np.isnan(wvpa.sel(time="2020-08-02T06:00")).values,
colors=["none"],
linecolors=["white"],
levels=[-0.5, 0.9, 1.5],
legend_style=None,
labels=False,
zorder=-10,
)
for m in earthkit.plots.schemas.schema.quickmap_subplot_workflow:
if m != "title":
getattr(chart, m)()
for m in earthkit.plots.schemas.schema.quickmap_figure_workflow:
if m != "legend":
getattr(chart, m)()
chart.legend()
counts, bins = np.histogram(da.values.flatten(), range=(2, 73), bins=21)
midpoints = bins[:-1] + np.diff(bins) / 2
cb = chart.ax.collections[1].colorbar
extend_width = (bins[-1] - bins[-2]) / (bins[-1] - bins[0])
cax = cb.ax.inset_axes(
[
0.0 - extend_width * extend_left,
1.25,
1.0 + extend_width * (0 + extend_left + extend_right + 2),
1.0,
]
)
cax.bar(
midpoints,
height=counts,
width=(bins[-1] - bins[0]) / len(counts),
color=cb.cmap(cb.norm(midpoints)),
)
if extend_left:
cax.bar(
bins[0] - (bins[1] - bins[0]) / 2,
height=np.sum(da < 2),
width=(bins[-1] - bins[0]) / len(counts),
color=cb.cmap(cb.norm(midpoints[0])),
)
if extend_right:
cax.bar(
bins[-1] + (bins[-1] - bins[-2]) / 2,
height=np.sum(da > 73),
width=(bins[-1] - bins[0]) / len(counts),
color=cb.cmap(cb.norm(midpoints[-1])),
)
cax.bar(
bins[-1] + (bins[-1] - bins[-2]) * (extend_right * 2 + 2 + 1) / 2,
height=np.sum(np.isnan(da)),
width=(bins[-1] - bins[0]) / len(counts),
color="lightgrey",
edgecolor="white",
lw=0,
hatch="XXXX",
)
with np.errstate(invalid="ignore"):
q1, q2, q3 = np.quantile(da.values.flatten(), [0.25, 0.5, 0.75])
cax.axvline(da.mean().item(), ls=":", ymin=0.1, ymax=0.9, c="w", lw=2)
cax.axvline(q1, ymin=0.25, ymax=0.75, c="w", lw=2)
cax.axvline(q2, ymin=0.1, ymax=0.9, c="w", lw=2)
cax.axvline(q3, ymin=0.25, ymax=0.75, c="w", lw=2)
cax.axvline(da.mean().item(), ymin=0.1, ymax=0.9, ls=":", c="k", lw=1)
cax.axvline(q1, ymin=0.25, ymax=0.75, c="k", lw=1)
cax.axvline(q2, ymin=0.1, ymax=0.9, c="k", lw=1)
cax.axvline(q3, ymin=0.25, ymax=0.75, c="k", lw=1)
cax.set_xlim(
2 - (bins[-1] - bins[-2]) * extend_left,
73 + (bins[-1] - bins[-2]) * (extend_right + 2),
)
cax.set_xticks([])
cax.set_yticks([])
cax.spines[:].set_visible(False)
def plot_wvpa(
da: xr.DataArray,
cr,
chart,
title,
eb_abs,
error=False,
corr=None,
inset=True,
):
# compute the default style that earthkit.maps would apply
source = earthkit.plots.sources.XarraySource(da)
style = copy.deepcopy(
earthkit.plots.styles.auto.guess_style(
source,
units=source.units,
)
)
style._levels = earthkit.plots.styles.levels.Levels(np.linspace(2, 73, 22))
style._legend_kwargs["ticks"] = np.linspace(2, 73, 5)
extend_left = np.nanmin(da) < 2
extend_right = np.nanmax(da) > 73
extend = {
(False, False): "neither",
(True, False): "min",
(False, True): "max",
(True, True): "both",
}[(extend_left, extend_right)]
style._legend_kwargs["extend"] = extend
chart.ax.set_global()
chart.ax.fill_between(
[0, 1],
[1, 1],
hatch="X",
edgecolor="white",
facecolor="lightgrey",
transform=chart.ax.transAxes,
zorder=-12,
)
chart.pcolormesh(
da.sel(time="2020-08-02T06:00"), style=style, zorder=-11, rasterized=True
)
with plt.rc_context(
{
"hatch.color": "white",
"hatch.linewidth": 0.5,
}
):
chart.contourf(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape[-2:],
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape[-2:],
),
z=np.isnan(wvpa.sel(time="2020-08-02T06:00")).values,
colors=["none"],
levels=[-0.5, 0.9, 1.5],
hatches=[None, "X"],
legend_style=None,
zorder=-10,
)
if error:
if corr is not None:
da_hatch = da.copy(
data=(da.values.view(np.uint32) == corr.values.view(np.uint32))
)
da_corr = (~da_hatch).astype(float)
old_process_projection_requirements = (
chart.ax.get_figure()._process_projection_requirements
)
def _process_projection_requirements(
*, axes_class=None, polar=False, projection=None, **kwargs
):
if axes_class is not None and projection is not None:
return axes_class, dict(projection=projection, **kwargs)
return old_process_projection_requirements(
axes_class=axes_class, polar=polar, projection=projection, **kwargs
)
chart.ax.get_figure()._process_projection_requirements = (
_process_projection_requirements
)
axin = chart.ax.inset_axes(
[0.025, 0.05, 1 / 3, 1 / 3],
xticklabels=[],
yticklabels=[],
axes_class=type(chart.ax),
projection=chart.ax.projection,
)
axin.set_global()
axin.pcolormesh(
da_corr.lon.values,
da_corr.lat.values,
np.squeeze(da_corr.sel(time="2020-08-02T06:00").values),
cmap=mpl.colors.ListedColormap(["white", "green", "lawngreen"]),
vmin=0,
vmax=2,
rasterized=True,
)
axin.coastlines(color="#555555")
axin.spines["geo"].set_edgecolor("black")
axin.set_title(
"Corrections",
path_effects=[PathEffects.withStroke(linewidth=3, foreground="white")],
)
elif inset:
da_err = ~((np.abs(da - wvpa) <= eb_abs) | (np.isnan(da) & np.isnan(wvpa)))
old_process_projection_requirements = (
chart.ax.get_figure()._process_projection_requirements
)
def _process_projection_requirements(
*, axes_class=None, polar=False, projection=None, **kwargs
):
if axes_class is not None and projection is not None:
return axes_class, dict(projection=projection, **kwargs)
return old_process_projection_requirements(
axes_class=axes_class, polar=polar, projection=projection, **kwargs
)
chart.ax.get_figure()._process_projection_requirements = (
_process_projection_requirements
)
axin = chart.ax.inset_axes(
[0.025, 0.05, 1 / 3, 1 / 3],
xticklabels=[],
yticklabels=[],
axes_class=type(chart.ax),
projection=chart.ax.projection,
)
axin.set_global()
axin.pcolormesh(
da_err.lon.values,
da_err.lat.values,
np.squeeze(da_err.sel(time="2020-08-02T06:00").values),
cmap=mpl.colors.ListedColormap(["white", "red"]),
vmin=0,
vmax=1,
rasterized=True,
)
axin.coastlines(color="#555555")
axin.spines["geo"].set_edgecolor("black")
axin.set_title(
"Violations",
path_effects=[PathEffects.withStroke(linewidth=3, foreground="white")],
)
chart.title(title)
if error:
err_v = np.mean(
~((np.abs(da - wvpa) <= eb_abs) | (np.isnan(da) & np.isnan(wvpa)))
)
err_v = (
0
if err_v == 0
else np.format_float_positional(100 * err_v, precision=1, min_digits=1)
+ "%"
)
if err_v == "0.0%":
err_v = "<0.05%"
t = chart.ax.text(
0.95,
0.1,
f"V={err_v}",
ha="right",
va="bottom",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
t = chart.ax.text(
0.95,
0.9,
rf"$\times$ {np.round(cr, 2)}"
if error
else humanize.naturalsize(wvpa.nbytes, binary=True),
ha="right",
va="top",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
chart.contour(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape[-2:],
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape[-2:],
),
z=np.isnan(wvpa.sel(time="2020-08-02T06:00")).values,
colors=["none"],
linecolors=["white"],
levels=[-0.5, 0.9, 1.5],
legend_style=None,
labels=False,
zorder=-10,
)
for m in earthkit.plots.schemas.schema.quickmap_subplot_workflow:
if m != "title":
getattr(chart, m)()
for m in earthkit.plots.schemas.schema.quickmap_figure_workflow:
if m != "legend":
getattr(chart, m)()
chart.legend()
counts, bins = np.histogram(da.values.flatten(), range=(2, 73), bins=21)
midpoints = bins[:-1] + np.diff(bins) / 2
cb = chart.ax.collections[1].colorbar
extend_width = (bins[-1] - bins[-2]) / (bins[-1] - bins[0])
cax = cb.ax.inset_axes(
[
0.0 - extend_width * extend_left,
1.25,
1.0 + extend_width * (0 + extend_left + extend_right + 2),
1.0,
]
)
cax.bar(
midpoints,
height=counts,
width=(bins[-1] - bins[0]) / len(counts),
color=cb.cmap(cb.norm(midpoints)),
)
if extend_left:
cax.bar(
bins[0] - (bins[1] - bins[0]) / 2,
height=np.sum(da < 2),
width=(bins[-1] - bins[0]) / len(counts),
color=cb.cmap(cb.norm(midpoints[0])),
)
if extend_right:
cax.bar(
bins[-1] + (bins[-1] - bins[-2]) / 2,
height=np.sum(da > 73),
width=(bins[-1] - bins[0]) / len(counts),
color=cb.cmap(cb.norm(midpoints[-1])),
)
cax.bar(
bins[-1] + (bins[-1] - bins[-2]) * (extend_right * 2 + 2 + 1) / 2,
height=np.sum(np.isnan(da)),
width=(bins[-1] - bins[0]) / len(counts),
color="lightgrey",
edgecolor="white",
lw=0,
hatch="XXXX",
)
with np.errstate(invalid="ignore"):
q1, q2, q3 = np.quantile(da.values.flatten(), [0.25, 0.5, 0.75])
cax.axvline(da.mean().item(), ls=":", ymin=0.1, ymax=0.9, c="w", lw=2)
cax.axvline(q1, ymin=0.25, ymax=0.75, c="w", lw=2)
cax.axvline(q2, ymin=0.1, ymax=0.9, c="w", lw=2)
cax.axvline(q3, ymin=0.25, ymax=0.75, c="w", lw=2)
cax.axvline(da.mean().item(), ymin=0.1, ymax=0.9, ls=":", c="k", lw=1)
cax.axvline(q1, ymin=0.25, ymax=0.75, c="k", lw=1)
cax.axvline(q2, ymin=0.1, ymax=0.9, c="k", lw=1)
cax.axvline(q3, ymin=0.25, ymax=0.75, c="k", lw=1)
cax.set_xlim(
2 - (bins[-1] - bins[-2]) * extend_left,
73 + (bins[-1] - bins[-2]) * (extend_right + 2),
)
cax.set_xticks([])
cax.set_yticks([])
cax.spines[:].set_visible(False)
Copied!
def table_wvpa(
da: xr.DataArray,
cr,
title,
eb_abs,
corr=None,
cr_nan=None,
):
err_inf = np.amax(np.abs(da - wvpa))
err_2 = np.sqrt(np.mean(np.square(da - wvpa)))
err_v = np.mean(~((np.abs(da - wvpa) <= eb_abs) | (np.isnan(da) & np.isnan(wvpa))))
err_v = (
0
if err_v == 0
else np.format_float_positional(100 * err_v, precision=1, min_digits=1) + "%"
)
if err_v == "0.0%":
err_v = "<0.05%"
corr = None if corr is None else compute_corrections_percentage(da, corr)
corr = (
""
if corr is None
else (
0
if corr == 0
else np.format_float_positional(100 * corr, precision=1, min_digits=1) + "%"
)
)
if corr == "0.0%":
corr = "<0.05%"
return pd.DataFrame(
{
"Compressor": [title[0]],
"Safeguarded": [title[1]],
"Corrections": [title[2]],
r"$L_{\infty}$": [
f"{err_inf:.02}",
],
r"$L_{2}$": [
f"{err_2:.02}",
],
"V": [err_v],
"C": [corr],
"CR (same NaN)": [
"" if cr_nan is None else rf"$\times$ {np.round(cr_nan, 2)}"
],
"CR (any NaN)": [
rf"$\times$ {np.round(cr, 2)}",
],
}
)
def table_wvpa(
da: xr.DataArray,
cr,
title,
eb_abs,
corr=None,
cr_nan=None,
):
err_inf = np.amax(np.abs(da - wvpa))
err_2 = np.sqrt(np.mean(np.square(da - wvpa)))
err_v = np.mean(~((np.abs(da - wvpa) <= eb_abs) | (np.isnan(da) & np.isnan(wvpa))))
err_v = (
0
if err_v == 0
else np.format_float_positional(100 * err_v, precision=1, min_digits=1) + "%"
)
if err_v == "0.0%":
err_v = "<0.05%"
corr = None if corr is None else compute_corrections_percentage(da, corr)
corr = (
""
if corr is None
else (
0
if corr == 0
else np.format_float_positional(100 * corr, precision=1, min_digits=1) + "%"
)
)
if corr == "0.0%":
corr = "<0.05%"
return pd.DataFrame(
{
"Compressor": [title[0]],
"Safeguarded": [title[1]],
"Corrections": [title[2]],
r"$L_{\infty}$": [
f"{err_inf:.02}",
],
r"$L_{2}$": [
f"{err_2:.02}",
],
"V": [err_v],
"C": [corr],
"CR (same NaN)": [
"" if cr_nan is None else rf"$\times$ {np.round(cr_nan, 2)}"
],
"CR (any NaN)": [
rf"$\times$ {np.round(cr, 2)}",
],
}
)
Copied!
fig = earthkit.plots.Figure(
size=(10, 19),
rows=5,
columns=2,
)
plot_wvpa(wvpa, 1.0, fig.add_map(0, 0), "Original", eb_abs=eb_abs)
plot_wvpa(
wvpa_zfp,
wvpa_zfp_cr,
fig.add_map(1, 0),
r"ZFP($\epsilon_{{abs}}$)",
eb_abs=eb_abs,
error=True,
)
plot_wvpa(
wvpa_sz3_2,
wvpa_sz3_2_cr,
fig.add_map(2, 0),
r"SZ3[v3.2]($\epsilon_{{abs}}$)",
eb_abs=eb_abs,
error=True,
)
plot_wvpa(
wvpa_sz3_3,
wvpa_sz3_3_cr,
fig.add_map(3, 0),
r"SZ3[v3.3]($\epsilon_{{abs}}$)",
eb_abs=eb_abs,
error=True,
)
plot_wvpa(
wvpa_sperr,
wvpa_sperr_cr,
fig.add_map(4, 0),
r"SPERR($\epsilon_{{abs}}$)",
eb_abs=eb_abs,
error=True,
)
plot_wvpa(
wvpa_sg["zero"],
wvpa_sg_cr["zero"],
fig.add_map(0, 1),
r"Safeguarded(0, $\epsilon_{{abs}} \cup \overset{{\text{{NaN}}}}{{\leftrightarrow}}$ )",
error=True,
eb_abs=eb_abs,
corr=wvpa_zero,
)
plot_wvpa(
wvpa_sg["zfp.rs"],
wvpa_sg_cr["zfp.rs"],
fig.add_map(1, 1),
r"Safeguarded(ZFP, $\epsilon_{{abs}} \cup \overset{{\text{{NaN}}}}{{\leftrightarrow}}$ )",
error=True,
eb_abs=eb_abs,
corr=wvpa_zfp,
)
plot_wvpa(
wvpa_sg["sz3.2.rs"],
wvpa_sg_cr["sz3.2.rs"],
fig.add_map(2, 1),
r"Safeguarded(SZ3[v3.2], $\epsilon_{{abs}} \cup \overset{{\text{{NaN}}}}{{\leftrightarrow}}$ )",
error=True,
eb_abs=eb_abs,
corr=wvpa_sz3_2,
)
plot_wvpa(
wvpa_sg["sz3.3.rs"],
wvpa_sg_cr["sz3.3.rs"],
fig.add_map(3, 1),
r"Safeguarded(SZ3, $\epsilon_{{abs}} \cup \overset{{\text{{NaN}}}}{{\leftrightarrow}}$ )",
error=True,
eb_abs=eb_abs,
corr=wvpa_sz3_3,
)
plot_wvpa(
wvpa_sg["sperr.rs"],
wvpa_sg_cr["sperr.rs"],
fig.add_map(4, 1),
r"Safeguarded(SPERR, $\epsilon_{{abs}} \cup \overset{{\text{{NaN}}}}{{\leftrightarrow}}$ )",
error=True,
eb_abs=eb_abs,
corr=wvpa_sperr,
)
fig.save(Path("plots") / "nan.pdf")
fig = earthkit.plots.Figure(
size=(10, 19),
rows=5,
columns=2,
)
plot_wvpa(wvpa, 1.0, fig.add_map(0, 0), "Original", eb_abs=eb_abs)
plot_wvpa(
wvpa_zfp,
wvpa_zfp_cr,
fig.add_map(1, 0),
r"ZFP($\epsilon_{{abs}}$)",
eb_abs=eb_abs,
error=True,
)
plot_wvpa(
wvpa_sz3_2,
wvpa_sz3_2_cr,
fig.add_map(2, 0),
r"SZ3[v3.2]($\epsilon_{{abs}}$)",
eb_abs=eb_abs,
error=True,
)
plot_wvpa(
wvpa_sz3_3,
wvpa_sz3_3_cr,
fig.add_map(3, 0),
r"SZ3[v3.3]($\epsilon_{{abs}}$)",
eb_abs=eb_abs,
error=True,
)
plot_wvpa(
wvpa_sperr,
wvpa_sperr_cr,
fig.add_map(4, 0),
r"SPERR($\epsilon_{{abs}}$)",
eb_abs=eb_abs,
error=True,
)
plot_wvpa(
wvpa_sg["zero"],
wvpa_sg_cr["zero"],
fig.add_map(0, 1),
r"Safeguarded(0, $\epsilon_{{abs}} \cup \overset{{\text{{NaN}}}}{{\leftrightarrow}}$ )",
error=True,
eb_abs=eb_abs,
corr=wvpa_zero,
)
plot_wvpa(
wvpa_sg["zfp.rs"],
wvpa_sg_cr["zfp.rs"],
fig.add_map(1, 1),
r"Safeguarded(ZFP, $\epsilon_{{abs}} \cup \overset{{\text{{NaN}}}}{{\leftrightarrow}}$ )",
error=True,
eb_abs=eb_abs,
corr=wvpa_zfp,
)
plot_wvpa(
wvpa_sg["sz3.2.rs"],
wvpa_sg_cr["sz3.2.rs"],
fig.add_map(2, 1),
r"Safeguarded(SZ3[v3.2], $\epsilon_{{abs}} \cup \overset{{\text{{NaN}}}}{{\leftrightarrow}}$ )",
error=True,
eb_abs=eb_abs,
corr=wvpa_sz3_2,
)
plot_wvpa(
wvpa_sg["sz3.3.rs"],
wvpa_sg_cr["sz3.3.rs"],
fig.add_map(3, 1),
r"Safeguarded(SZ3, $\epsilon_{{abs}} \cup \overset{{\text{{NaN}}}}{{\leftrightarrow}}$ )",
error=True,
eb_abs=eb_abs,
corr=wvpa_sz3_3,
)
plot_wvpa(
wvpa_sg["sperr.rs"],
wvpa_sg_cr["sperr.rs"],
fig.add_map(4, 1),
r"Safeguarded(SPERR, $\epsilon_{{abs}} \cup \overset{{\text{{NaN}}}}{{\leftrightarrow}}$ )",
error=True,
eb_abs=eb_abs,
corr=wvpa_sperr,
)
fig.save(Path("plots") / "nan.pdf")
Copied!
wvpa_table = pd.concat(
[
table_wvpa(
wvpa_sg_lossless["zero"],
wvpa_sg_lossless_cr["zero"],
[
"0",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"lossless",
],
eb_abs=eb_abs,
corr=wvpa_zero,
cr_nan=wvpa_sg_nan_lossless_cr["zero"],
),
table_wvpa(
wvpa_sg["zero"],
wvpa_sg_cr["zero"],
[
"0",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"one-shot",
],
eb_abs=eb_abs,
corr=wvpa_zero,
cr_nan=wvpa_sg_nan_cr["zero"],
),
table_wvpa(
wvpa_zfp,
wvpa_zfp_cr,
[r"ZFP($\epsilon_{abs}$)", "-", ""],
eb_abs=eb_abs,
),
table_wvpa(
wvpa_sg_lossless["zfp.rs"],
wvpa_sg_lossless_cr["zfp.rs"],
[
r"ZFP($\epsilon_{abs}$)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"lossless",
],
eb_abs=eb_abs,
corr=wvpa_zfp,
cr_nan=wvpa_sg_nan_lossless_cr["zfp.rs"],
),
table_wvpa(
wvpa_sg["zfp.rs"],
wvpa_sg_cr["zfp.rs"],
[
r"ZFP($\epsilon_{abs}$)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"one-shot",
],
eb_abs=eb_abs,
corr=wvpa_zfp,
cr_nan=wvpa_sg_nan_cr["zfp.rs"],
),
table_wvpa(
wvpa_optzconfig["zfp.rs"],
wvpa_optzconfig_cr["zfp.rs"],
[
r"OptZConfig(ZFP)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"",
],
eb_abs=eb_abs,
),
table_wvpa(
wvpa_sz3_2,
wvpa_sz3_2_cr,
[r"SZ3[v3.2]($\epsilon_{abs}$)", "-", ""],
eb_abs=eb_abs,
),
table_wvpa(
wvpa_sg_lossless["sz3.2.rs"],
wvpa_sg_lossless_cr["sz3.2.rs"],
[
r"SZ3[v3.2]($\epsilon_{abs}$)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"lossless",
],
eb_abs=eb_abs,
corr=wvpa_sz3_2,
cr_nan=wvpa_sg_nan_lossless_cr["sz3.2.rs"],
),
table_wvpa(
wvpa_sg["sz3.2.rs"],
wvpa_sg_cr["sz3.2.rs"],
[
r"SZ3[v3.2]($\epsilon_{abs}$)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"one-shot",
],
eb_abs=eb_abs,
corr=wvpa_sz3_2,
cr_nan=wvpa_sg_nan_cr["sz3.2.rs"],
),
table_wvpa(
wvpa_optzconfig["sz3.2.rs"],
wvpa_optzconfig_cr["sz3.2.rs"],
[
r"OptZConfig(SZ3[v3.2])",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"",
],
eb_abs=eb_abs,
),
table_wvpa(
wvpa_sz3_3,
wvpa_sz3_3_cr,
[r"SZ3[v3.3]($\epsilon_{abs}$)", "-", ""],
eb_abs=eb_abs,
),
table_wvpa(
wvpa_sg_lossless["sz3.3.rs"],
wvpa_sg_lossless_cr["sz3.3.rs"],
[
r"SZ3[v3.3]($\epsilon_{abs}$)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"lossless",
],
eb_abs=eb_abs,
corr=wvpa_sz3_3,
cr_nan=wvpa_sg_nan_lossless_cr["sz3.3.rs"],
),
table_wvpa(
wvpa_sg["sz3.3.rs"],
wvpa_sg_cr["sz3.3.rs"],
[
r"SZ3[v3.3]($\epsilon_{abs}$)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"one-shot",
],
eb_abs=eb_abs,
corr=wvpa_sz3_3,
cr_nan=wvpa_sg_nan_cr["sz3.3.rs"],
),
table_wvpa(
wvpa_optzconfig["sz3.3.rs"],
wvpa_optzconfig_cr["sz3.3.rs"],
[
r"OptZConfig(SZ3[v3.3])",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"",
],
eb_abs=eb_abs,
),
table_wvpa(
wvpa_sperr,
wvpa_sperr_cr,
[r"SPERR($\epsilon_{abs}$)", "-", ""],
eb_abs=eb_abs,
),
table_wvpa(
wvpa_sg_lossless["sperr.rs"],
wvpa_sg_lossless_cr["sperr.rs"],
[
r"SPERR($\epsilon_{abs}$)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"lossless",
],
eb_abs=eb_abs,
corr=wvpa_sperr,
cr_nan=wvpa_sg_nan_lossless_cr["sperr.rs"],
),
table_wvpa(
wvpa_sg["sperr.rs"],
wvpa_sg_cr["sperr.rs"],
[
r"SPERR($\epsilon_{abs}$)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"one-shot",
],
eb_abs=eb_abs,
corr=wvpa_sperr,
cr_nan=wvpa_sg_nan_cr["sperr.rs"],
),
table_wvpa(
wvpa_optzconfig["sperr.rs"],
wvpa_optzconfig_cr["sperr.rs"],
[
r"OptZConfig(SPERR)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"",
],
eb_abs=eb_abs,
),
table_wvpa(
wvpa_zstd,
wvpa_zstd_cr,
["ZSTD(22)", "-", ""],
eb_abs=eb_abs,
),
]
).set_index(["Compressor", "Safeguarded", "Corrections"])
Path("tables").joinpath("nan.tex").write_text(
wvpa_table.to_latex(escape=False)
.replace("%", r"\%")
.replace("\\cline{1-9} \\cline{2-9}\n\\bottomrule", "\\bottomrule")
)
wvpa_table
wvpa_table = pd.concat(
[
table_wvpa(
wvpa_sg_lossless["zero"],
wvpa_sg_lossless_cr["zero"],
[
"0",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"lossless",
],
eb_abs=eb_abs,
corr=wvpa_zero,
cr_nan=wvpa_sg_nan_lossless_cr["zero"],
),
table_wvpa(
wvpa_sg["zero"],
wvpa_sg_cr["zero"],
[
"0",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"one-shot",
],
eb_abs=eb_abs,
corr=wvpa_zero,
cr_nan=wvpa_sg_nan_cr["zero"],
),
table_wvpa(
wvpa_zfp,
wvpa_zfp_cr,
[r"ZFP($\epsilon_{abs}$)", "-", ""],
eb_abs=eb_abs,
),
table_wvpa(
wvpa_sg_lossless["zfp.rs"],
wvpa_sg_lossless_cr["zfp.rs"],
[
r"ZFP($\epsilon_{abs}$)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"lossless",
],
eb_abs=eb_abs,
corr=wvpa_zfp,
cr_nan=wvpa_sg_nan_lossless_cr["zfp.rs"],
),
table_wvpa(
wvpa_sg["zfp.rs"],
wvpa_sg_cr["zfp.rs"],
[
r"ZFP($\epsilon_{abs}$)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"one-shot",
],
eb_abs=eb_abs,
corr=wvpa_zfp,
cr_nan=wvpa_sg_nan_cr["zfp.rs"],
),
table_wvpa(
wvpa_optzconfig["zfp.rs"],
wvpa_optzconfig_cr["zfp.rs"],
[
r"OptZConfig(ZFP)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"",
],
eb_abs=eb_abs,
),
table_wvpa(
wvpa_sz3_2,
wvpa_sz3_2_cr,
[r"SZ3[v3.2]($\epsilon_{abs}$)", "-", ""],
eb_abs=eb_abs,
),
table_wvpa(
wvpa_sg_lossless["sz3.2.rs"],
wvpa_sg_lossless_cr["sz3.2.rs"],
[
r"SZ3[v3.2]($\epsilon_{abs}$)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"lossless",
],
eb_abs=eb_abs,
corr=wvpa_sz3_2,
cr_nan=wvpa_sg_nan_lossless_cr["sz3.2.rs"],
),
table_wvpa(
wvpa_sg["sz3.2.rs"],
wvpa_sg_cr["sz3.2.rs"],
[
r"SZ3[v3.2]($\epsilon_{abs}$)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"one-shot",
],
eb_abs=eb_abs,
corr=wvpa_sz3_2,
cr_nan=wvpa_sg_nan_cr["sz3.2.rs"],
),
table_wvpa(
wvpa_optzconfig["sz3.2.rs"],
wvpa_optzconfig_cr["sz3.2.rs"],
[
r"OptZConfig(SZ3[v3.2])",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"",
],
eb_abs=eb_abs,
),
table_wvpa(
wvpa_sz3_3,
wvpa_sz3_3_cr,
[r"SZ3[v3.3]($\epsilon_{abs}$)", "-", ""],
eb_abs=eb_abs,
),
table_wvpa(
wvpa_sg_lossless["sz3.3.rs"],
wvpa_sg_lossless_cr["sz3.3.rs"],
[
r"SZ3[v3.3]($\epsilon_{abs}$)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"lossless",
],
eb_abs=eb_abs,
corr=wvpa_sz3_3,
cr_nan=wvpa_sg_nan_lossless_cr["sz3.3.rs"],
),
table_wvpa(
wvpa_sg["sz3.3.rs"],
wvpa_sg_cr["sz3.3.rs"],
[
r"SZ3[v3.3]($\epsilon_{abs}$)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"one-shot",
],
eb_abs=eb_abs,
corr=wvpa_sz3_3,
cr_nan=wvpa_sg_nan_cr["sz3.3.rs"],
),
table_wvpa(
wvpa_optzconfig["sz3.3.rs"],
wvpa_optzconfig_cr["sz3.3.rs"],
[
r"OptZConfig(SZ3[v3.3])",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"",
],
eb_abs=eb_abs,
),
table_wvpa(
wvpa_sperr,
wvpa_sperr_cr,
[r"SPERR($\epsilon_{abs}$)", "-", ""],
eb_abs=eb_abs,
),
table_wvpa(
wvpa_sg_lossless["sperr.rs"],
wvpa_sg_lossless_cr["sperr.rs"],
[
r"SPERR($\epsilon_{abs}$)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"lossless",
],
eb_abs=eb_abs,
corr=wvpa_sperr,
cr_nan=wvpa_sg_nan_lossless_cr["sperr.rs"],
),
table_wvpa(
wvpa_sg["sperr.rs"],
wvpa_sg_cr["sperr.rs"],
[
r"SPERR($\epsilon_{abs}$)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"one-shot",
],
eb_abs=eb_abs,
corr=wvpa_sperr,
cr_nan=wvpa_sg_nan_cr["sperr.rs"],
),
table_wvpa(
wvpa_optzconfig["sperr.rs"],
wvpa_optzconfig_cr["sperr.rs"],
[
r"OptZConfig(SPERR)",
r"$\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$",
"",
],
eb_abs=eb_abs,
),
table_wvpa(
wvpa_zstd,
wvpa_zstd_cr,
["ZSTD(22)", "-", ""],
eb_abs=eb_abs,
),
]
).set_index(["Compressor", "Safeguarded", "Corrections"])
Path("tables").joinpath("nan.tex").write_text(
wvpa_table.to_latex(escape=False)
.replace("%", r"\%")
.replace("\\cline{1-9} \\cline{2-9}\n\\bottomrule", "\\bottomrule")
)
wvpa_table
| $L_{\infty}$ | $L_{2}$ | V | C | CR (same NaN) | CR (any NaN) | |||
|---|---|---|---|---|---|---|---|---|
| Compressor | Safeguarded | Corrections | ||||||
| 0 | $\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$ | lossless | 0.0 | 0.0 | 0 | 100.0% | $\times$ 7.57 | $\times$ 7.57 |
| one-shot | 0.1 | 0.05 | 0 | 100.0% | $\times$ 13.18 | $\times$ 13.18 | ||
| ZFP($\epsilon_{abs}$) | - | 0.03 | 0.0047 | 68.3% | $\times$ 5.03 | |||
| $\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$ | lossless | 0.03 | 0.0047 | 0 | 68.3% | $\times$ 3.35 | $\times$ 3.35 | |
| one-shot | 0.03 | 0.0047 | 0 | 68.3% | $\times$ 3.35 | $\times$ 4.22 | ||
| OptZConfig(ZFP) | $\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$ | 3.8e-06 | 7.9e-08 | 68.3% | $\times$ 2.05 | |||
| SZ3[v3.2]($\epsilon_{abs}$) | - | 0.1 | 0.057 | 68.3% | $\times$ 16.92 | |||
| $\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$ | lossless | 0.1 | 0.057 | 0 | 68.3% | $\times$ 2.31 | $\times$ 2.31 | |
| one-shot | 0.1 | 0.057 | 0 | 68.3% | $\times$ 2.31 | $\times$ 14.06 | ||
| OptZConfig(SZ3[v3.2]) | $\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$ | 2.4e-07 | 3.5e-09 | 68.3% | $\times$ 5.13 | |||
| SZ3[v3.3]($\epsilon_{abs}$) | - | 0.1 | 0.049 | 0 | $\times$ 9.06 | |||
| $\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$ | lossless | 0.1 | 0.049 | 0 | 0 | $\times$ 9.06 | $\times$ 9.06 | |
| one-shot | 0.1 | 0.049 | 0 | 0 | $\times$ 9.06 | $\times$ 9.06 | ||
| OptZConfig(SZ3[v3.3]) | $\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$ | 0.1 | 0.049 | 0 | $\times$ 9.06 | |||
| SPERR($\epsilon_{abs}$) | - | 0.1 | 0.041 | 68.3% | $\times$ 11.58 | |||
| $\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$ | lossless | 0.1 | 0.041 | 0 | 68.3% | $\times$ 2.39 | $\times$ 2.39 | |
| one-shot | 0.1 | 0.041 | 0 | 68.3% | $\times$ 2.39 | $\times$ 11.01 | ||
| OptZConfig(SPERR) | $\epsilon_{abs} \cup \overset{\text{NaN}}{\leftrightarrow}$ | 4.8e-07 | 4.2e-08 | 68.3% | $\times$ 2.59 | |||
| ZSTD(22) | - | 0.0 | 0.0 | 0 | $\times$ 6.0 |
Summary visualisation¶
Copied!
fig = earthkit.plots.Figure(
size=(10, 5),
rows=1,
columns=2,
)
chart = fig.add_map(0, 0)
# compute the default style that earthkit.maps would apply
source = earthkit.plots.sources.XarraySource(wvpa)
style = copy.deepcopy(
earthkit.plots.styles.auto.guess_style(
source,
units=source.units,
)
)
style._levels = earthkit.plots.styles.levels.Levels(np.linspace(2, 73, 22))
style._legend_kwargs["ticks"] = np.linspace(2, 73, 5)
style._legend_kwargs["extend"] = "both"
chart.ax.fill_between(
[0, 1],
[1, 1],
hatch="X",
edgecolor="white",
facecolor="lightgrey",
transform=chart.ax.transAxes,
zorder=-12,
)
# Original
da = wvpa.sel(time="2020-08-02T06:00", lon=slice(-180, -60))
chart.pcolormesh(da, style=style, zorder=-11)
t = chart.ax.text(
1 / 6,
0.9,
"Original",
ha="center",
va="top",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
# SZ3[v3.2]
da = wvpa.sel(time="2020-08-02T06:00", lon=slice(-60, +60))
dac = wvpa_sz3_2.sel(time="2020-08-02T06:00", lon=slice(-60, +60))
chart.pcolormesh(dac, style=style, zorder=-11)
with plt.rc_context(
{
"hatch.color": "white",
"hatch.linewidth": 0.5,
}
):
chart.contourf(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
levels=[-0.5, 0.9, 1.5],
hatches=[None, "X"],
legend_style=None,
zorder=-10,
)
t = chart.ax.text(
0.5,
0.9,
"SZ3[v3.2]",
ha="center",
va="top",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
t = chart.ax.text(
0.5,
0.5,
rf"$\times$ {np.round(wvpa_sz3_2_cr, 2)}",
ha="center",
va="center",
transform=chart.ax.transAxes,
color="mistyrose",
fontsize=20,
)
t.set_path_effects([PathEffects.withStroke(linewidth=3, foreground="black")])
err_v = np.mean(
~((np.abs(wvpa_sz3_2 - wvpa) <= eb_abs) | (np.isnan(wvpa_sz3_2) & np.isnan(wvpa)))
)
err_v = (
0
if err_v == 0
else np.format_float_positional(100 * err_v, precision=1, min_digits=1) + "%"
)
if err_v == "0.0%":
err_v = "<0.05%"
t = chart.ax.text(
0.5,
0.1,
f"V={err_v}",
ha="center",
va="bottom",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
# SZ3[v3.3]
da = wvpa.sel(time="2020-08-02T06:00", lon=slice(+60, 180))
dac = wvpa_sz3_3.sel(time="2020-08-02T06:00", lon=slice(+60, 180))
chart.pcolormesh(dac, style=style, zorder=-11)
with plt.rc_context(
{
"hatch.color": "white",
"hatch.linewidth": 0.5,
}
):
chart.contourf(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
levels=[-0.5, 0.9, 1.5],
hatches=[None, "X"],
legend_style=None,
zorder=-10,
)
t = chart.ax.text(
5 / 6,
0.9,
"SZ3[v3.3]",
ha="center",
va="top",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
t = chart.ax.text(
5 / 6,
0.5,
rf"$\times$ {np.round(wvpa_sz3_3_cr, 2)}",
ha="center",
va="center",
transform=chart.ax.transAxes,
color="lightgreen",
fontsize=20,
)
t.set_path_effects([PathEffects.withStroke(linewidth=3, foreground="black")])
err_v = np.mean(
~((np.abs(wvpa_sz3_3 - wvpa) <= eb_abs) | (np.isnan(wvpa_sz3_3) & np.isnan(wvpa)))
)
err_v = (
0
if err_v == 0
else np.format_float_positional(100 * err_v, precision=1, min_digits=1) + "%"
)
if err_v == "0.0%":
err_v = "<0.05%"
t = chart.ax.text(
5 / 6,
0.1,
f"V={err_v}",
ha="center",
va="bottom",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
da = wvpa.sel(time="2020-08-02T06:00")
chart.contour(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
linecolors=["white"],
levels=[-0.5, 0.9, 1.5],
legend_style=None,
labels=False,
zorder=-10,
)
chart.ax.set_rasterization_zorder(-10)
chart.ax.axvline(-60, c="white", ls=(2, (4, 4)), lw=2)
chart.ax.axvline(-60, c="black", ls=(6, (4, 4)), lw=2)
chart.ax.axvline(+60, c="white", ls=(2, (4, 4)), lw=2)
chart.ax.axvline(+60, c="black", ls=(6, (4, 4)), lw=2)
chart.title("Without safeguards")
for m in earthkit.plots.schemas.schema.quickmap_subplot_workflow:
if m != "title":
getattr(chart, m)()
for m in earthkit.plots.schemas.schema.quickmap_figure_workflow:
if m != "legend":
getattr(chart, m)()
chart = fig.add_map(0, 1)
chart.ax.fill_between(
[0, 1],
[1, 1],
hatch="X",
edgecolor="white",
facecolor="lightgrey",
transform=chart.ax.transAxes,
zorder=-12,
)
# Safeguarded(0)
da = wvpa.sel(time="2020-08-02T06:00", lon=slice(-180, -60))
dac = wvpa_sg["zero"].sel(time="2020-08-02T06:00", lon=slice(-180, -60))
chart.pcolormesh(dac, style=style, zorder=-11)
with plt.rc_context(
{
"hatch.color": "white",
"hatch.linewidth": 0.5,
}
):
chart.contourf(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
levels=[-0.5, 0.9, 1.5],
hatches=[None, "X"],
legend_style=None,
zorder=-10,
)
t = chart.ax.text(
1 / 6,
0.9,
"Sg(0)",
ha="center",
va="top",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
t = chart.ax.text(
1 / 6,
0.5,
rf"$\times$ {np.round(wvpa_sg_cr['zero'], 2)}",
ha="center",
va="center",
transform=chart.ax.transAxes,
color="lightgreen",
fontsize=20,
)
t.set_path_effects([PathEffects.withStroke(linewidth=3, foreground="black")])
err_v = np.mean(
~(
(np.abs(wvpa_sg["zero"] - wvpa) <= eb_abs)
| (np.isnan(wvpa_sg["zero"]) & np.isnan(wvpa))
)
)
err_v = (
0
if err_v == 0
else np.format_float_positional(100 * err_v, precision=1, min_digits=1) + "%"
)
if err_v == "0.0%":
err_v = "<0.05%"
t = chart.ax.text(
1 / 6,
0.1,
f"V={err_v}",
ha="center",
va="bottom",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
# Safeguarded(SZ3[v3.2])
da = wvpa.sel(time="2020-08-02T06:00", lon=slice(-60, +60))
dac = wvpa_sg["sz3.2.rs"].sel(time="2020-08-02T06:00", lon=slice(-60, +60))
chart.pcolormesh(dac, style=style, zorder=-11)
with plt.rc_context(
{
"hatch.color": "white",
"hatch.linewidth": 0.5,
}
):
chart.contourf(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
levels=[-0.5, 0.9, 1.5],
hatches=[None, "X"],
legend_style=None,
zorder=-10,
)
t = chart.ax.text(
0.5,
0.9,
"Sg(SZ3[v3.2])",
ha="center",
va="top",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
t = chart.ax.text(
0.5,
0.5,
rf"$\times$ {np.round(wvpa_sg_cr['sz3.2.rs'], 2)}",
ha="center",
va="center",
transform=chart.ax.transAxes,
color="lightgreen",
fontsize=20,
)
t.set_path_effects([PathEffects.withStroke(linewidth=3, foreground="black")])
err_v = np.mean(
~(
(np.abs(wvpa_sg["sz3.2.rs"] - wvpa) <= eb_abs)
| (np.isnan(wvpa_sg["sz3.2.rs"]) & np.isnan(wvpa))
)
)
err_v = (
0
if err_v == 0
else np.format_float_positional(100 * err_v, precision=1, min_digits=1) + "%"
)
if err_v == "0.0%":
err_v = "<0.05%"
t = chart.ax.text(
0.5,
0.1,
f"V={err_v}",
ha="center",
va="bottom",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
# Safeguarded(SZ3.3)
da = wvpa.sel(time="2020-08-02T06:00", lon=slice(+60, 180))
dac = wvpa_sg["sz3.3.rs"].sel(time="2020-08-02T06:00", lon=slice(+60, 180))
chart.pcolormesh(dac, style=style, zorder=-11)
with plt.rc_context(
{
"hatch.color": "white",
"hatch.linewidth": 0.5,
}
):
chart.contourf(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
levels=[-0.5, 0.9, 1.5],
hatches=[None, "X"],
legend_style=None,
zorder=-10,
)
t = chart.ax.text(
5 / 6,
0.9,
"Sg(SZ3[v3.3])",
ha="center",
va="top",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
t = chart.ax.text(
5 / 6,
0.5,
rf"$\times$ {np.round(wvpa_sg_cr['sz3.3.rs'], 2)}",
ha="center",
va="center",
transform=chart.ax.transAxes,
color="lightgreen",
fontsize=20,
)
t.set_path_effects([PathEffects.withStroke(linewidth=3, foreground="black")])
err_v = np.mean(
~(
(np.abs(wvpa_sg["sz3.3.rs"] - wvpa) <= eb_abs)
| (np.isnan(wvpa_sg["sz3.3.rs"]) & np.isnan(wvpa))
)
)
err_v = (
0
if err_v == 0
else np.format_float_positional(100 * err_v, precision=1, min_digits=1) + "%"
)
if err_v == "0.0%":
err_v = "<0.05%"
t = chart.ax.text(
5 / 6,
0.1,
f"V={err_v}",
ha="center",
va="bottom",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
da = wvpa.sel(time="2020-08-02T06:00")
chart.contour(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
linecolors=["white"],
levels=[-0.5, 0.9, 1.5],
legend_style=None,
labels=False,
zorder=-10,
)
chart.ax.set_rasterization_zorder(-10)
chart.ax.axvline(-60, c="white", ls=(2, (4, 4)), lw=2)
chart.ax.axvline(-60, c="black", ls=(6, (4, 4)), lw=2)
chart.ax.axvline(+60, c="white", ls=(2, (4, 4)), lw=2)
chart.ax.axvline(+60, c="black", ls=(6, (4, 4)), lw=2)
for m in earthkit.plots.schemas.schema.quickmap_subplot_workflow:
if m != "title":
getattr(chart, m)()
for m in earthkit.plots.schemas.schema.quickmap_figure_workflow:
if m != "legend":
getattr(chart, m)()
chart.title("Safeguarded: preserve missing values")
fig.legend()
fig.save(Path("plots") / "nan-summary.pdf")
fig = earthkit.plots.Figure(
size=(10, 5),
rows=1,
columns=2,
)
chart = fig.add_map(0, 0)
# compute the default style that earthkit.maps would apply
source = earthkit.plots.sources.XarraySource(wvpa)
style = copy.deepcopy(
earthkit.plots.styles.auto.guess_style(
source,
units=source.units,
)
)
style._levels = earthkit.plots.styles.levels.Levels(np.linspace(2, 73, 22))
style._legend_kwargs["ticks"] = np.linspace(2, 73, 5)
style._legend_kwargs["extend"] = "both"
chart.ax.fill_between(
[0, 1],
[1, 1],
hatch="X",
edgecolor="white",
facecolor="lightgrey",
transform=chart.ax.transAxes,
zorder=-12,
)
# Original
da = wvpa.sel(time="2020-08-02T06:00", lon=slice(-180, -60))
chart.pcolormesh(da, style=style, zorder=-11)
t = chart.ax.text(
1 / 6,
0.9,
"Original",
ha="center",
va="top",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
# SZ3[v3.2]
da = wvpa.sel(time="2020-08-02T06:00", lon=slice(-60, +60))
dac = wvpa_sz3_2.sel(time="2020-08-02T06:00", lon=slice(-60, +60))
chart.pcolormesh(dac, style=style, zorder=-11)
with plt.rc_context(
{
"hatch.color": "white",
"hatch.linewidth": 0.5,
}
):
chart.contourf(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
levels=[-0.5, 0.9, 1.5],
hatches=[None, "X"],
legend_style=None,
zorder=-10,
)
t = chart.ax.text(
0.5,
0.9,
"SZ3[v3.2]",
ha="center",
va="top",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
t = chart.ax.text(
0.5,
0.5,
rf"$\times$ {np.round(wvpa_sz3_2_cr, 2)}",
ha="center",
va="center",
transform=chart.ax.transAxes,
color="mistyrose",
fontsize=20,
)
t.set_path_effects([PathEffects.withStroke(linewidth=3, foreground="black")])
err_v = np.mean(
~((np.abs(wvpa_sz3_2 - wvpa) <= eb_abs) | (np.isnan(wvpa_sz3_2) & np.isnan(wvpa)))
)
err_v = (
0
if err_v == 0
else np.format_float_positional(100 * err_v, precision=1, min_digits=1) + "%"
)
if err_v == "0.0%":
err_v = "<0.05%"
t = chart.ax.text(
0.5,
0.1,
f"V={err_v}",
ha="center",
va="bottom",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
# SZ3[v3.3]
da = wvpa.sel(time="2020-08-02T06:00", lon=slice(+60, 180))
dac = wvpa_sz3_3.sel(time="2020-08-02T06:00", lon=slice(+60, 180))
chart.pcolormesh(dac, style=style, zorder=-11)
with plt.rc_context(
{
"hatch.color": "white",
"hatch.linewidth": 0.5,
}
):
chart.contourf(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
levels=[-0.5, 0.9, 1.5],
hatches=[None, "X"],
legend_style=None,
zorder=-10,
)
t = chart.ax.text(
5 / 6,
0.9,
"SZ3[v3.3]",
ha="center",
va="top",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
t = chart.ax.text(
5 / 6,
0.5,
rf"$\times$ {np.round(wvpa_sz3_3_cr, 2)}",
ha="center",
va="center",
transform=chart.ax.transAxes,
color="lightgreen",
fontsize=20,
)
t.set_path_effects([PathEffects.withStroke(linewidth=3, foreground="black")])
err_v = np.mean(
~((np.abs(wvpa_sz3_3 - wvpa) <= eb_abs) | (np.isnan(wvpa_sz3_3) & np.isnan(wvpa)))
)
err_v = (
0
if err_v == 0
else np.format_float_positional(100 * err_v, precision=1, min_digits=1) + "%"
)
if err_v == "0.0%":
err_v = "<0.05%"
t = chart.ax.text(
5 / 6,
0.1,
f"V={err_v}",
ha="center",
va="bottom",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
da = wvpa.sel(time="2020-08-02T06:00")
chart.contour(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
linecolors=["white"],
levels=[-0.5, 0.9, 1.5],
legend_style=None,
labels=False,
zorder=-10,
)
chart.ax.set_rasterization_zorder(-10)
chart.ax.axvline(-60, c="white", ls=(2, (4, 4)), lw=2)
chart.ax.axvline(-60, c="black", ls=(6, (4, 4)), lw=2)
chart.ax.axvline(+60, c="white", ls=(2, (4, 4)), lw=2)
chart.ax.axvline(+60, c="black", ls=(6, (4, 4)), lw=2)
chart.title("Without safeguards")
for m in earthkit.plots.schemas.schema.quickmap_subplot_workflow:
if m != "title":
getattr(chart, m)()
for m in earthkit.plots.schemas.schema.quickmap_figure_workflow:
if m != "legend":
getattr(chart, m)()
chart = fig.add_map(0, 1)
chart.ax.fill_between(
[0, 1],
[1, 1],
hatch="X",
edgecolor="white",
facecolor="lightgrey",
transform=chart.ax.transAxes,
zorder=-12,
)
# Safeguarded(0)
da = wvpa.sel(time="2020-08-02T06:00", lon=slice(-180, -60))
dac = wvpa_sg["zero"].sel(time="2020-08-02T06:00", lon=slice(-180, -60))
chart.pcolormesh(dac, style=style, zorder=-11)
with plt.rc_context(
{
"hatch.color": "white",
"hatch.linewidth": 0.5,
}
):
chart.contourf(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
levels=[-0.5, 0.9, 1.5],
hatches=[None, "X"],
legend_style=None,
zorder=-10,
)
t = chart.ax.text(
1 / 6,
0.9,
"Sg(0)",
ha="center",
va="top",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
t = chart.ax.text(
1 / 6,
0.5,
rf"$\times$ {np.round(wvpa_sg_cr['zero'], 2)}",
ha="center",
va="center",
transform=chart.ax.transAxes,
color="lightgreen",
fontsize=20,
)
t.set_path_effects([PathEffects.withStroke(linewidth=3, foreground="black")])
err_v = np.mean(
~(
(np.abs(wvpa_sg["zero"] - wvpa) <= eb_abs)
| (np.isnan(wvpa_sg["zero"]) & np.isnan(wvpa))
)
)
err_v = (
0
if err_v == 0
else np.format_float_positional(100 * err_v, precision=1, min_digits=1) + "%"
)
if err_v == "0.0%":
err_v = "<0.05%"
t = chart.ax.text(
1 / 6,
0.1,
f"V={err_v}",
ha="center",
va="bottom",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
# Safeguarded(SZ3[v3.2])
da = wvpa.sel(time="2020-08-02T06:00", lon=slice(-60, +60))
dac = wvpa_sg["sz3.2.rs"].sel(time="2020-08-02T06:00", lon=slice(-60, +60))
chart.pcolormesh(dac, style=style, zorder=-11)
with plt.rc_context(
{
"hatch.color": "white",
"hatch.linewidth": 0.5,
}
):
chart.contourf(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
levels=[-0.5, 0.9, 1.5],
hatches=[None, "X"],
legend_style=None,
zorder=-10,
)
t = chart.ax.text(
0.5,
0.9,
"Sg(SZ3[v3.2])",
ha="center",
va="top",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
t = chart.ax.text(
0.5,
0.5,
rf"$\times$ {np.round(wvpa_sg_cr['sz3.2.rs'], 2)}",
ha="center",
va="center",
transform=chart.ax.transAxes,
color="lightgreen",
fontsize=20,
)
t.set_path_effects([PathEffects.withStroke(linewidth=3, foreground="black")])
err_v = np.mean(
~(
(np.abs(wvpa_sg["sz3.2.rs"] - wvpa) <= eb_abs)
| (np.isnan(wvpa_sg["sz3.2.rs"]) & np.isnan(wvpa))
)
)
err_v = (
0
if err_v == 0
else np.format_float_positional(100 * err_v, precision=1, min_digits=1) + "%"
)
if err_v == "0.0%":
err_v = "<0.05%"
t = chart.ax.text(
0.5,
0.1,
f"V={err_v}",
ha="center",
va="bottom",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
# Safeguarded(SZ3.3)
da = wvpa.sel(time="2020-08-02T06:00", lon=slice(+60, 180))
dac = wvpa_sg["sz3.3.rs"].sel(time="2020-08-02T06:00", lon=slice(+60, 180))
chart.pcolormesh(dac, style=style, zorder=-11)
with plt.rc_context(
{
"hatch.color": "white",
"hatch.linewidth": 0.5,
}
):
chart.contourf(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
levels=[-0.5, 0.9, 1.5],
hatches=[None, "X"],
legend_style=None,
zorder=-10,
)
t = chart.ax.text(
5 / 6,
0.9,
"Sg(SZ3[v3.3])",
ha="center",
va="top",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
t = chart.ax.text(
5 / 6,
0.5,
rf"$\times$ {np.round(wvpa_sg_cr['sz3.3.rs'], 2)}",
ha="center",
va="center",
transform=chart.ax.transAxes,
color="lightgreen",
fontsize=20,
)
t.set_path_effects([PathEffects.withStroke(linewidth=3, foreground="black")])
err_v = np.mean(
~(
(np.abs(wvpa_sg["sz3.3.rs"] - wvpa) <= eb_abs)
| (np.isnan(wvpa_sg["sz3.3.rs"]) & np.isnan(wvpa))
)
)
err_v = (
0
if err_v == 0
else np.format_float_positional(100 * err_v, precision=1, min_digits=1) + "%"
)
if err_v == "0.0%":
err_v = "<0.05%"
t = chart.ax.text(
5 / 6,
0.1,
f"V={err_v}",
ha="center",
va="bottom",
transform=chart.ax.transAxes,
)
t.set_bbox(dict(facecolor="white", alpha=0.75, edgecolor="black"))
da = wvpa.sel(time="2020-08-02T06:00")
chart.contour(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
linecolors=["white"],
levels=[-0.5, 0.9, 1.5],
legend_style=None,
labels=False,
zorder=-10,
)
chart.ax.set_rasterization_zorder(-10)
chart.ax.axvline(-60, c="white", ls=(2, (4, 4)), lw=2)
chart.ax.axvline(-60, c="black", ls=(6, (4, 4)), lw=2)
chart.ax.axvline(+60, c="white", ls=(2, (4, 4)), lw=2)
chart.ax.axvline(+60, c="black", ls=(6, (4, 4)), lw=2)
for m in earthkit.plots.schemas.schema.quickmap_subplot_workflow:
if m != "title":
getattr(chart, m)()
for m in earthkit.plots.schemas.schema.quickmap_figure_workflow:
if m != "legend":
getattr(chart, m)()
chart.title("Safeguarded: preserve missing values")
fig.legend()
fig.save(Path("plots") / "nan-summary.pdf")
Detail visualisation¶
Copied!
for key, da in {
"original": wvpa,
"approximation": wvpa_sz3_2,
"corrected": wvpa_sg["sz3.2.rs"],
}.items():
fig = earthkit.plots.Figure(
size=(2, 2),
rows=1,
columns=1,
)
domain = earthkit.plots.geo.domains.Domain.from_bbox(bbox=[60, 180, -60, 60])
chart = fig.add_map(0, 0, domain=domain)
# compute the default style that earthkit.maps would apply
source = earthkit.plots.sources.XarraySource(wvpa)
style = copy.deepcopy(
earthkit.plots.styles.auto.guess_style(
source,
units=source.units,
)
)
style._levels = earthkit.plots.styles.levels.Levels(np.linspace(2, 73, 22))
style._legend_kwargs["ticks"] = np.linspace(2, 73, 5)
style._legend_kwargs["extend"] = "both"
with plt.rc_context(
{
"hatch.linewidth": 1,
}
):
chart.ax.fill_between(
[0, 1],
[1, 1],
hatch="X",
edgecolor="white",
facecolor="lightgrey",
transform=chart.ax.transAxes,
zorder=-12,
rasterized=False,
)
da = da.sel(time="2020-08-02T06:00")
chart.pcolormesh(da, style=style, zorder=-11, rasterized=True)
da = wvpa.sel(time="2020-08-02T06:00")
with plt.rc_context(
{
"hatch.color": "white",
"hatch.linewidth": 1,
}
):
chart.contourf(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
levels=[-0.5, 0.9, 1.5],
hatches=[None, "X"],
legend_style=None,
zorder=-10,
)
chart.contour(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
linecolors=["white"],
levels=[-0.5, 0.9, 1.5],
legend_style=None,
labels=False,
zorder=-10,
linewidths=2,
)
chart.coastlines()
fig.save(Path("plots") / f"nan-{key}.svg", transparent=True)
for key, da in {
"original": wvpa,
"approximation": wvpa_sz3_2,
"corrected": wvpa_sg["sz3.2.rs"],
}.items():
fig = earthkit.plots.Figure(
size=(2, 2),
rows=1,
columns=1,
)
domain = earthkit.plots.geo.domains.Domain.from_bbox(bbox=[60, 180, -60, 60])
chart = fig.add_map(0, 0, domain=domain)
# compute the default style that earthkit.maps would apply
source = earthkit.plots.sources.XarraySource(wvpa)
style = copy.deepcopy(
earthkit.plots.styles.auto.guess_style(
source,
units=source.units,
)
)
style._levels = earthkit.plots.styles.levels.Levels(np.linspace(2, 73, 22))
style._legend_kwargs["ticks"] = np.linspace(2, 73, 5)
style._legend_kwargs["extend"] = "both"
with plt.rc_context(
{
"hatch.linewidth": 1,
}
):
chart.ax.fill_between(
[0, 1],
[1, 1],
hatch="X",
edgecolor="white",
facecolor="lightgrey",
transform=chart.ax.transAxes,
zorder=-12,
rasterized=False,
)
da = da.sel(time="2020-08-02T06:00")
chart.pcolormesh(da, style=style, zorder=-11, rasterized=True)
da = wvpa.sel(time="2020-08-02T06:00")
with plt.rc_context(
{
"hatch.color": "white",
"hatch.linewidth": 1,
}
):
chart.contourf(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
levels=[-0.5, 0.9, 1.5],
hatches=[None, "X"],
legend_style=None,
zorder=-10,
)
chart.contour(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
linecolors=["white"],
levels=[-0.5, 0.9, 1.5],
legend_style=None,
labels=False,
zorder=-10,
linewidths=2,
)
chart.coastlines()
fig.save(Path("plots") / f"nan-{key}.svg", transparent=True)
Copied!
fig = earthkit.plots.Figure(
size=(2, 2),
rows=1,
columns=1,
)
domain = earthkit.plots.geo.domains.Domain.from_bbox(bbox=[60, 180, -60, 60])
chart = fig.add_map(0, 0, domain=domain)
# the actual correction is in binary, but a floating point difference looks nicer
correction = np.nan_to_num(wvpa_sg["sz3.2.rs"].values) - np.nan_to_num(
wvpa_sz3_2.values
)
# use NaN for no corrections, which we later fill in with neutral white
correction[correction == 0] = np.nan
correction = wvpa_sz3_2.copy(data=correction)
da = correction.sel(time="2020-08-02T06:00")
# compute the default style that earthkit.maps would apply
source = earthkit.plots.sources.XarraySource(wvpa)
style = copy.deepcopy(
earthkit.plots.styles.auto.guess_style(
source,
units=source.units,
)
)
# "invert" the colourmap so that the correction looks like it reverses the changes
style._levels = earthkit.plots.styles.levels.Levels(
np.linspace(np.nanmin(da.values), 0, 22)
)
chart.ax.fill_between(
[0, 1],
[1, 1],
facecolor="white",
transform=chart.ax.transAxes,
zorder=-12,
rasterized=False,
)
chart.pcolormesh(da, style=style, zorder=-11, rasterized=True)
da = wvpa.sel(time="2020-08-02T06:00")
chart.contour(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
linecolors=["lightgrey"],
levels=[-0.5, 0.9, 1.5],
legend_style=None,
labels=False,
zorder=-9,
linewidths=2,
)
chart.coastlines()
fig.save(Path("plots") / "nan-correction.svg", transparent=True)
fig = earthkit.plots.Figure(
size=(2, 2),
rows=1,
columns=1,
)
domain = earthkit.plots.geo.domains.Domain.from_bbox(bbox=[60, 180, -60, 60])
chart = fig.add_map(0, 0, domain=domain)
# the actual correction is in binary, but a floating point difference looks nicer
correction = np.nan_to_num(wvpa_sg["sz3.2.rs"].values) - np.nan_to_num(
wvpa_sz3_2.values
)
# use NaN for no corrections, which we later fill in with neutral white
correction[correction == 0] = np.nan
correction = wvpa_sz3_2.copy(data=correction)
da = correction.sel(time="2020-08-02T06:00")
# compute the default style that earthkit.maps would apply
source = earthkit.plots.sources.XarraySource(wvpa)
style = copy.deepcopy(
earthkit.plots.styles.auto.guess_style(
source,
units=source.units,
)
)
# "invert" the colourmap so that the correction looks like it reverses the changes
style._levels = earthkit.plots.styles.levels.Levels(
np.linspace(np.nanmin(da.values), 0, 22)
)
chart.ax.fill_between(
[0, 1],
[1, 1],
facecolor="white",
transform=chart.ax.transAxes,
zorder=-12,
rasterized=False,
)
chart.pcolormesh(da, style=style, zorder=-11, rasterized=True)
da = wvpa.sel(time="2020-08-02T06:00")
chart.contour(
x=np.broadcast_to(
da.lon.values.reshape(1, -1),
da.shape,
),
y=np.broadcast_to(
da.lat.values.reshape(-1, 1),
da.shape,
),
z=np.isnan(da).values,
colors=["none"],
linecolors=["lightgrey"],
levels=[-0.5, 0.9, 1.5],
legend_style=None,
labels=False,
zorder=-9,
linewidths=2,
)
chart.coastlines()
fig.save(Path("plots") / "nan-correction.svg", transparent=True)
Copied!
import json
with Path("observations").joinpath("nan.json").open("w") as f:
json.dump(observations, f)
import json
with Path("observations").joinpath("nan.json").open("w") as f:
json.dump(observations, f)
Copied!