"""Kernel analog forecasting of the El Nino Southern Oscillation. The training and test data used for experiment ENSO_FROM_ERA5_NINO34SST are included under /examples/enso/data. """ import jax.numpy as jnp import nc_time_axis as nc_time_axis import nlsa.jax.distance as dst import nlsa.jax.kernels as knl import nlsa.jax.scalars as scls import nlsa_models.climate as clim import nlsa_models.era5 as era5 import numpy as np import sys from collections.abc import Sequence from dataclasses import dataclass from enum import StrEnum, auto from jax import Array from nlsa.function_algebra import compose2 from nlsa.io_actions import ( IO, h5it, netcdfit, pickleit, plotit, plotem, timeit, ) from nlsa.jax.kernels import ( BsKernelPars, ConePars, DmKernelPars, KernelEigen, KernelEigenShardings, KernelPars, TuneInfo, TunePars, ) from nlsa.jax.sharding import NamedSharder from nlsa.jax.vector_algebra import L2FnAlgebraShardings from nlsa.jax.utils import fst from nlsa_models.climate import ( Covariate, Climatology, DataPars, NPData, ProjectionPars, Response, SkillScores, Time, TimeSampling, ) from pathlib import Path from tabulate import tabulate from typing import Literal, TypedDict type Plots = Literal[ "all", "bandwidth_tuning", "bandwidth_func", "kernel_tuning", "laplacian_spec", "kernel_eigen", "kernel_modes", "expansion_coeffs", "running_pred", "pred_timeseries", "skill_scores", ] class Experiment(StrEnum): """Experiments provided in this script.""" NINO34_FROM_ERA5_NINO34SST = auto() """Prediction of the Nino 3.4 index from ERA5 Nino 3.4 SST.""" ELNINO_2027_FROM_ERA5_NINO34SST = auto() """Prediction of the 2026-27 super El Nino from ERA5 Nino 3.4 SST.""" EXPERIMENT: Experiment = Experiment.ELNINO_2027_FROM_ERA5_NINO34SST IDX_GPU: int | Sequence[int] | None = 0 XLA_MEM_FRACTION: str | None = "0.95" JAX_CACHE_DIR: str | None = "jax_cache" FP: Literal["f32", "f64"] = "f32" CONE_KERNEL: bool = True KERNEL_TUNING_GRAD_METHOD: Literal["explicit", "automatic"] = "automatic" KERNEL_NORMALIZATION: Literal["diffusion_maps", "bistochastic"] = ( "diffusion_maps" ) MATPLOTLIB_BACKEND: Literal["Agg"] | None = None OUTPUT_DATA_DIR = "examples/enso/data" NUM_TABULATE = 40 EXTRACT_DATA_MODE: Literal["calc", "calcsave", "read"] = "read" TUNE_KERNEL_MODE: Literal["calc", "calcsave", "read"] = "calc" KERNEL_EIGEN_MODE: Literal["calc", "calcsave", "read"] = "calc" KERNEL_PROJECTION_MODE: Literal["calc", "calcsave", "read"] = "calc" KAF_EXPANSION_COEFFS_MODE: Literal["calc", "calcsave", "read"] = "calc" KAF_PREDS_MODE: Literal["calc", "calcsave", "read"] = "calc" SKILL_SCORES_MODE: Literal["calc", "calcsave", "read"] = "calc" PLOT_MODE: Literal["save", "show", "saveshow"] | None = "show" WHICH_PLOTS: set[Plots] = {"running_pred", "pred_timeseries", "skill_scores"} DELAY_PLOT_MODE: Literal["backward", "central"] = "central" PLT_DATE_RANGE: tuple[str, str] | None = None KERNEL_EIGS_PLT: Sequence[int] | Literal["interactive"] | None = "interactive" KERNEL_MODES_PLT: Sequence[int] | Literal["interactive"] | None = "interactive" LEAD_TIMES_PLT: Sequence[int] | Literal["interactive"] | None = "interactive" INITIALIZATION_TIMES_PLT: Sequence[int] | Literal["interactive"] | None = ( "interactive" ) jax_env = clim.initialize_jax( idx_gpu=IDX_GPU, xla_mem_fraction=XLA_MEM_FRACTION, fp=FP, ) clim.initialize_matplotlib(backend=MATPLOTLIB_BACKEND) @dataclass(frozen=True, slots=True) class PredPars: """Dataclass containing prediction parameters.""" dt: int """Prediction timestep.""" # WARNING: The current KAF implementation ignores this parameter and # assumes dt=1. num_steps: int """Number of timesteps for prediction.""" # TODO: Replace tuple[int, int] by slice[int, int, int]. which_eigs: int | tuple[int, int] | list[int] """Kernel eigenfunctions eigenfunctions used in the prediction function.""" def __str__(self) -> str: """Create string representing prediction parameters.""" match self.which_eigs: case int(): eigs_str = "-".join(map(str, (0, self.which_eigs))) case (_, _): eigs_str = "-".join(map(str, self.which_eigs)) case list(): eigs_str = "_".join(map(str, self.which_eigs)) return "_".join((f"pred{self.num_steps}", "eigs" + eigs_str)) @dataclass(frozen=True, slots=True) class TrainPars[T: TimeSampling]: """Dataclass containing the training parameter values.""" data: DataPars[T] """Training data parameters.""" tune: TunePars """Kernel tuning parameters.""" kernel: KernelPars """Kernel eigendecomposition parameters.""" pred: PredPars """Prediction parameters.""" cone: ConePars | None = None """Cone kernel parameters.""" bw_tune: TunePars | None = None """Tuning parameters for kernel bandwidth function.""" def __str__(self) -> str: """Create string representation of training parameters.""" if self.cone is not None: cone_str = str(self.cone) else: cone_str = "" if self.bw_tune is None: bw_tune_str = "" else: bw_tune_str = str(self.bw_tune) return "_".join( filter( None, ( str(self.data), str(self.tune), bw_tune_str, cone_str, str(self.kernel), str(self.pred), ), ) ) # TODO: Complete this def tabulate(self, show: bool = True) -> str: """Create tabulated summary of the properties of a TrainPars object.""" tables = [self.data.tabulate(name="Training data", show=show)] return "".join(tables) @dataclass(frozen=True, slots=True) class TestPars[T: TimeSampling]: """Dataclass containing test parameter values.""" num_pred_steps: int """Number of prediction steps.""" data: DataPars[T] """Test data parameters.""" # TODO: Complete this def tabulate(self, show: bool = True) -> str: """Create tabulated summary of the properties of a TestPars object.""" tables = ( self.data.tabulate(name="Test data", show=show), tabulate([{"Number of prediction steps": self.num_pred_steps}]), ) return "".join(tables) @dataclass(frozen=True, slots=True) class Pars[T: TimeSampling]: """Dataclass containing the parameter values used in this example.""" train: TrainPars[T] """Training parameters.""" test: TestPars[T] """Test parameters.""" kernel_projection: ProjectionPars[T] | None = None """Projection parameters.""" def tabulate(self, show: bool = True) -> str: """Create tabulated summary of the properties of a Pars object.""" if self.kernel_projection is not None: kernel_proj_tab = self.kernel_projection.tabulate(show=show) else: kernel_proj_tab = None tables = ( self.train.tabulate(show=show), self.test.tabulate(show=show), kernel_proj_tab, ) return "".join(filter(None, tables)) class CommonPars(TypedDict): """Helper TypedDict to check common training/test parameter values.""" num_half_delays: int velocity_covariate: bool velocity_fd_order: Literal[2, 4, 6, 8] | None num_before: int @dataclass(frozen=True, slots=True) class TrainShardings: """Dataclass containing shardings related to the training data.""" l2: L2FnAlgebraShardings = L2FnAlgebraShardings() """Shardings for the L2 space associated with the training data.""" kernel_eigen: KernelEigenShardings = KernelEigenShardings() """Shardings for the kernel eigenvalue problem.""" @dataclass(frozen=True, slots=True) class TestShardings: """Dataclass containing shardings related to the test data.""" l2: L2FnAlgebraShardings = L2FnAlgebraShardings() """Shardings for the L2 space associated with the test data.""" @dataclass(frozen=True, slots=True) class Shardings: """Dataclass containing shardings used in this example.""" train: TrainShardings = TrainShardings() """Training shardings.""" test: TestShardings = TestShardings() """Test shardings.""" def from_experiment( experiment: Experiment, cone_kernel: bool, kernel_normalization: Literal["diffusion_maps", "bistochastic"], ) -> tuple[Pars[Literal["monthly"]], Shardings]: """Prepare parameters and shardings for the numerical experiment.""" match experiment: case Experiment.NINO34_FROM_ERA5_NINO34SST: cone_pars = ConePars(zeta=0.99) if cone_kernel else None era5_io = era5.IO( input_path="/storage/data/era5/monthly_1940-2025_nc", file_format="nc", ) time_sampling = "monthly" fd_order = 4 num_pred_steps = 24 train_date_range = ("1940-01-01", "2019-12-31") test_date_range = ("2018-01-01", "2025-12-31") common_pars: CommonPars = { "num_half_delays": 6, "velocity_covariate": True if cone_pars is not None else False, "velocity_fd_order": fd_order if cone_pars is not None else None, "num_before": 0, } climatology_date_range = train_date_range covariate_rolling_window = None covariate_rolling_mode = "center" train_covariate_time = Time( date_range=train_date_range, sampling=time_sampling, custom_climatology_date_range=climatology_date_range, rolling_window=covariate_rolling_window, rolling_mode=covariate_rolling_mode, ) test_covariate_time = Time( date_range=test_date_range, sampling=time_sampling, custom_climatology_date_range=climatology_date_range, rolling_window=covariate_rolling_window, rolling_mode=covariate_rolling_mode, ) covariate_climatology = Climatology(remove=False, standardize=True) covariate_vars = [era5.Var.SST] covariate_domain = era5.nino34_domain(step_lon=4, step_lat=4) projection_rolling_window = None projection_rolling_mode = "center" projection_time = Time( date_range=train_date_range, sampling=time_sampling, custom_climatology_date_range=climatology_date_range, rolling_window=projection_rolling_window, rolling_mode=projection_rolling_mode, ) projection_climatology = Climatology( remove=False, standardize=False ) projection_vars = [era5.Var.SST] projection_domain = era5.global_domain(step_lon=4, step_lat=4) response_rolling_window = None response_rolling_mode = "center" train_response_time = Time( date_range=train_date_range, sampling=time_sampling, rolling_window=response_rolling_window, rolling_mode=response_rolling_mode, custom_climatology_date_range=climatology_date_range, ) test_response_time = Time( date_range=test_date_range, sampling=time_sampling, rolling_window=response_rolling_window, rolling_mode=response_rolling_mode, custom_climatology_date_range=climatology_date_range, ) response_climatology = Climatology(remove=True, standardize=False) response_var = era5.Var.SST response_domain = era5.nino34_domain(sampling="area_averaged") train_covariate_specs = ( era5.DataSpecs( vars=covariate_vars, domain=covariate_domain, time=train_covariate_time, io=era5_io, climatology=covariate_climatology, ), ) projection_specs = era5.DataSpecs( vars=projection_vars, domain=projection_domain, time=projection_time, io=era5_io, climatology=projection_climatology, ) train_response_specs = era5.DataSpecs( vars=[response_var], domain=response_domain, time=train_response_time, io=era5_io, climatology=response_climatology, ) test_covariate_specs = ( era5.DataSpecs( vars=covariate_vars, domain=covariate_domain, time=test_covariate_time, io=era5_io, climatology=covariate_climatology, ), ) test_response_specs = era5.DataSpecs( vars=[response_var], domain=response_domain, time=test_response_time, io=era5_io, climatology=response_climatology, ) train_data_pars = DataPars( covariate=Covariate(specs=train_covariate_specs), response=Response(specs=train_response_specs), **common_pars, ) test_data_pars = DataPars( covariate=Covariate(specs=test_covariate_specs), response=Response(specs=test_response_specs), **common_pars, ) bw_tune_pars = TunePars( manifold_dim=None, num_bandwidths=128, log10_bandwidth_lims=(-3, 3), bandwidth_scl=1, ) if cone_pars is not None: tune_pars = TunePars( manifold_dim=None, num_bandwidths=128, log10_bandwidth_lims=(-3, 3), bandwidth_scl=2, ) else: tune_pars = TunePars( manifold_dim=None, num_bandwidths=128, log10_bandwidth_lims=(-3, 3), bandwidth_scl=1, ) match kernel_normalization: case "diffusion_maps": kernel_pars = DmKernelPars( normalization="fokkerplanck", eigensolver="eigh", num_eigs=512, ) case "bistochastic": kernel_pars = BsKernelPars( eigensolver="svd", num_eigs=512, ) pred_pars = PredPars( dt=1, which_eigs=512, num_steps=num_pred_steps ) kernel_projection_pars = ProjectionPars( data=projection_specs, time_origin=train_data_pars.delay_embedding_center, which_pcs=[0, 1, 2, 3], lead_steps=0, ) if len(jax_env.devices) > 1: sharder = NamedSharder( devices=jax_env.devices, shape=(len(jax_env.devices),), axis_names=("x"), ) x_sharding = sharder.sharding("x") replicating = sharder.sharding(None) l2_shardings = L2FnAlgebraShardings( data=x_sharding, vectors=x_sharding ) l2_tst_shardings = L2FnAlgebraShardings( data=replicating, vectors=replicating ) kernel_eigen_shardings = KernelEigenShardings( eigenvalues=replicating, eigenvectors=x_sharding, weights=x_sharding, ) train_shardings = TrainShardings( l2=l2_shardings, kernel_eigen=kernel_eigen_shardings, ) test_shardings = TestShardings(l2=l2_tst_shardings) else: train_shardings = TrainShardings() test_shardings = TestShardings() shardings = Shardings(train=train_shardings, test=test_shardings) case Experiment.ELNINO_2027_FROM_ERA5_NINO34SST: cone_pars = ConePars(zeta=0.99) if cone_kernel else None era5_io = era5.IO( input_path="/storage/data/era5/monthly_1940-202608_nc", file_format="nc", ) time_sampling = "monthly" fd_order = 4 num_pred_steps = 24 train_date_range = ("1940-01-01", "2026-08-31") test_date_range = ("2009-01-01", "2026-08-31") common_pars: CommonPars = { "num_half_delays": 6, "velocity_covariate": True if cone_pars is not None else False, "velocity_fd_order": fd_order if cone_pars is not None else None, "num_before": 0, } climatology_date_range = train_date_range covariate_rolling_window = None covariate_rolling_mode = "center" train_covariate_time = Time( date_range=train_date_range, sampling=time_sampling, custom_climatology_date_range=climatology_date_range, rolling_window=covariate_rolling_window, rolling_mode=covariate_rolling_mode, ) test_covariate_time = Time( date_range=test_date_range, sampling=time_sampling, custom_climatology_date_range=climatology_date_range, rolling_window=covariate_rolling_window, rolling_mode=covariate_rolling_mode, ) covariate_climatology = Climatology(remove=False, standardize=True) covariate_vars = [era5.Var.SST] covariate_domain = era5.nino34_domain(step_lon=4, step_lat=4) projection_rolling_window = None projection_rolling_mode = "center" projection_time = Time( date_range=train_date_range, sampling=time_sampling, custom_climatology_date_range=climatology_date_range, rolling_window=projection_rolling_window, rolling_mode=projection_rolling_mode, ) projection_climatology = Climatology( remove=False, standardize=False ) projection_vars = [era5.Var.SST] projection_domain = era5.global_domain(step_lon=4, step_lat=4) response_rolling_window = None response_rolling_mode = "center" train_response_time = Time( date_range=train_date_range, sampling=time_sampling, rolling_window=response_rolling_window, rolling_mode=response_rolling_mode, custom_climatology_date_range=climatology_date_range, ) test_response_time = Time( date_range=test_date_range, sampling=time_sampling, rolling_window=response_rolling_window, rolling_mode=response_rolling_mode, custom_climatology_date_range=climatology_date_range, ) response_climatology = Climatology(remove=True, standardize=False) response_var = era5.Var.SST response_domain = era5.nino34_domain(sampling="area_averaged") train_covariate_specs = ( era5.DataSpecs( vars=covariate_vars, domain=covariate_domain, time=train_covariate_time, io=era5_io, climatology=covariate_climatology, ), ) projection_specs = era5.DataSpecs( vars=projection_vars, domain=projection_domain, time=projection_time, io=era5_io, climatology=projection_climatology, ) train_response_specs = era5.DataSpecs( vars=[response_var], domain=response_domain, time=train_response_time, io=era5_io, climatology=response_climatology, ) test_covariate_specs = ( era5.DataSpecs( vars=covariate_vars, domain=covariate_domain, time=test_covariate_time, io=era5_io, climatology=covariate_climatology, ), ) test_response_specs = era5.DataSpecs( vars=[response_var], domain=response_domain, time=test_response_time, io=era5_io, climatology=response_climatology, ) train_data_pars = DataPars( covariate=Covariate(specs=train_covariate_specs), response=Response(specs=train_response_specs), num_after=num_pred_steps, **common_pars, ) test_data_pars = DataPars( covariate=Covariate(specs=test_covariate_specs), response=Response(specs=test_response_specs), **common_pars, ) bw_tune_pars = TunePars( manifold_dim=None, num_bandwidths=128, log10_bandwidth_lims=(-3, 3), bandwidth_scl=1, ) if cone_pars is not None: tune_pars = TunePars( manifold_dim=None, num_bandwidths=128, log10_bandwidth_lims=(-3, 3), bandwidth_scl=2, ) else: tune_pars = TunePars( manifold_dim=None, num_bandwidths=128, log10_bandwidth_lims=(-3, 3), bandwidth_scl=1, ) match kernel_normalization: case "diffusion_maps": kernel_pars = DmKernelPars( normalization="fokkerplanck", eigensolver="eigh", num_eigs=512, ) case "bistochastic": kernel_pars = BsKernelPars( eigensolver="svd", num_eigs=512, ) pred_pars = PredPars( dt=1, which_eigs=512, num_steps=num_pred_steps ) kernel_projection_pars = ProjectionPars( data=projection_specs, time_origin=train_data_pars.delay_embedding_center, which_pcs=[0, 1, 2, 3], lead_steps=0, ) if len(jax_env.devices) > 1: sharder = NamedSharder( devices=jax_env.devices, shape=(len(jax_env.devices),), axis_names=("x"), ) x_sharding = sharder.sharding("x") replicating = sharder.sharding(None) l2_shardings = L2FnAlgebraShardings( data=x_sharding, vectors=x_sharding ) l2_tst_shardings = L2FnAlgebraShardings( data=replicating, vectors=replicating ) kernel_eigen_shardings = KernelEigenShardings( eigenvalues=replicating, eigenvectors=x_sharding, weights=x_sharding, ) train_shardings = TrainShardings( l2=l2_shardings, kernel_eigen=kernel_eigen_shardings, ) test_shardings = TestShardings(l2=l2_tst_shardings) else: train_shardings = TrainShardings() test_shardings = TestShardings() shardings = Shardings(train=train_shardings, test=test_shardings) train_pars = TrainPars( data=train_data_pars, bw_tune=bw_tune_pars, cone=cone_pars, tune=tune_pars, kernel=kernel_pars, pred=pred_pars, ) test_pars = TestPars( data=test_data_pars, num_pred_steps=pred_pars.num_steps ) pars = Pars( train=train_pars, test=test_pars, kernel_projection=kernel_projection_pars, ) return pars, shardings pars, shardings = from_experiment( EXPERIMENT, CONE_KERNEL, KERNEL_NORMALIZATION ) io = IO(root=Path.cwd() / OUTPUT_DATA_DIR) extract_data_arrays = timeit( pickleit( clim.extract_data_arrays, io=io, mode=EXTRACT_DATA_MODE, fname="data", cls=NPData, ) ) compute_kernel_bandwidth = timeit( pickleit( knl.tune_bandwidth, io=io, mode=TUNE_KERNEL_MODE, fname="tune_info", cls=TuneInfo, ) ) compute_kernel_eigen = timeit( pickleit( knl.compute_eigen, io=io, mode=KERNEL_EIGEN_MODE, fname="kernel_eigen", cls=KernelEigen, ) ) compute_kernel_eigen_projections = timeit( netcdfit( clim.compute_eigenfunction_projections, io=io, mode=KERNEL_PROJECTION_MODE, fname="kernel_eigen_proj", ) ) compute_kaf_expansion_coeffs = timeit( pickleit( clim.compute_kaf_expansion_coeffs, io=io, mode=KAF_EXPANSION_COEFFS_MODE, fname="kaf_coeffs", cls=Array, ) ) compute_kaf_preds = timeit( pickleit( knl.compute_kaf_preds, io=io, mode=KAF_PREDS_MODE, fname="kaf_preds", cls=Array, ) ) compute_skill_scores = timeit( h5it( clim.compute_skill_scores, io=io, mode=SKILL_SCORES_MODE, fname="pred_scores", cls=SkillScores, callback=clim.to_skill_scores, ) ) plot_kernel_tuning = plotit( knl.plot_kernel_tuning, io=io, mode=PLOT_MODE, fname="bandwidth_tuning_func", ) plot_bandwidth_function = plotit( clim.plot_bandwidth_function, io=io, mode=PLOT_MODE, fname="bandwidth_func", ) plot_laplace_spectrum = plotit( knl.plot_laplacian_spectrum, io=io, mode=PLOT_MODE, fname="lapl_spec" ) make_kernel_evecs_plotter = plotem( clim.make_kernel_evecs_plotter, io=io, mode=PLOT_MODE, fname="kernel_eigen", ) make_kernel_mode_plotter = plotem( clim.make_2d_mode_plotter_with_leads, io=io, mode=PLOT_MODE, fname="kernel_modes", ) plot_kaf_expansion_coeffs = plotit( knl.plot_kaf_expansion_coeffs, io=io, mode=PLOT_MODE, fname="kaf_expansion_coeffs", ) make_running_pred_plotter = plotem( clim.make_running_pred_plotter, io=io, mode=PLOT_MODE, fname="pred_running", ) make_pred_timeseries_plotter = plotem( clim.make_pred_timeseries_plotter, io=io, mode=PLOT_MODE, fname="pred_timeseries", ) plot_forecast_skill_scores = plotit( clim.plot_forecast_skill_scores, io=io, mode=PLOT_MODE, fname="pred_scores", ) def main(): """Kernel analog forecasting of ESTCP data.""" global io # Display information about the computation to be performed jax_env.tabulate() pars.tabulate() # Generate training and test data io @= str(pars.train.data.covariate) io /= str(pars.train.data.response) io /= str(pars.train.data) train_data = extract_data_arrays( pars.train.data, dtype=np.dtype(np.float64), ).to_device(dtype=jax_env.real_dtype, shardings=shardings.train.l2.data) io @= str(pars.test.data.covariate) io /= str(pars.test.data.response) io /= str(pars.test.data) test_data = extract_data_arrays( pars.test.data, dtype=np.dtype(np.float64) ).to_device(dtype=jax_env.real_dtype, shardings=shardings.test.l2.data) # Make scalar field and L2 space builders scl_r = scls.scalar_field(jax_env.real_dtype) impl_l2 = clim.make_data_driven_l2_space( data_pars=pars.train.data, dtype=jax_env.real_dtype, shardings=shardings.train.l2, ) impl_l2_tst = clim.make_data_driven_l2_space( data_pars=pars.test.data, dtype=jax_env.real_dtype, shardings=shardings.test.l2, ) # Set kernel shape function io @= str(pars.train.data.covariate) io /= str(pars.train.data) shape_func = jnp.exp match KERNEL_TUNING_GRAD_METHOD: case "explicit": neg_grad_shape_func = jnp.exp case "automatic": neg_grad_shape_func = None # Create and tune bandwidth function if pars.train.bw_tune is not None: io /= str(pars.train.bw_tune) if pars.train.data.velocity_covariate: bw_sqdist = compose2(dst.sqeuclidean, (fst, fst)) else: bw_sqdist = dst.sqeuclidean bw_tune_info = compute_kernel_bandwidth( pars.train.bw_tune, impl_l2, shape_func, bw_sqdist, train_data, neg_grad_shape_func, ) bandwidth_func = knl.make_data_driven_bandwidth_function( impl_l2, shape_func, bw_sqdist, bw_tune_info ) bw_tune_info.tabulate(name="Bandwidth function tuning") # Plot bandwidth function tuning and bandwidth function if PLOT_MODE is not None and not { "all", "bandwidth_tuning", }.isdisjoint(WHICH_PLOTS): plot_kernel_tuning(bw_tune_info, title="Bandwidth function tuning") if PLOT_MODE is not None and not {"all", "bandwidth_func"}.isdisjoint( WHICH_PLOTS ): plot_bandwidth_function( pars.train.data, impl_l2, bandwidth_func, train_data, shardings.train.l2.data, pars.test.data, impl_l2_tst, test_data, shardings.test.l2.data, delay_plot_mode=DELAY_PLOT_MODE, plt_date_range=PLT_DATE_RANGE, ) else: bandwidth_func = None # Create and tune kernel if pars.train.cone is not None: io /= str(pars.train.cone) sqdist = dst.make_sqcone( pars.train.cone.zeta, pars.train.cone.threshold ) else: sqdist = dst.sqeuclidean if bandwidth_func is not None: sqdist = knl.make_data_driven_scaled_sqdist( scl_r, sqdist, bandwidth_func ) else: sqdist = sqdist io /= str(pars.train.tune) tune_info = compute_kernel_bandwidth( pars.train.tune, impl_l2, shape_func, sqdist, train_data, neg_grad_shape_func, ) tune_info.tabulate(name="Kernel tuning") # Plot kernel tuning function if PLOT_MODE is not None and not { "all", "kernel_tuning", }.isdisjoint(WHICH_PLOTS): plot_kernel_tuning(tune_info, title="Kernel tuning") # Solve kernel eigenvalue problem io /= str(pars.train.kernel) kernel = knl.make_data_driven_rbf_kernel( scl_r, shape_func, sqdist, tune_info.bandwidth ) kernel_eigen = compute_kernel_eigen( pars.train.kernel, impl_l2, kernel, train_data, tune_info.bandwidth, pars.train.data.num_samples, jax_env.real_dtype, shardings=shardings.train.kernel_eigen, ) if len(jax_env.devices) > 1: kernel_eigen.inspect_array_shardings() kernel_eigen.tabulate(num_tabulate=NUM_TABULATE) # Plot spectrum of Laplacian eigenvalues if PLOT_MODE is not None and not {"all", "laplacian_spec"}.isdisjoint( WHICH_PLOTS ): plot_laplace_spectrum(kernel_eigen) # Plot representative kernel eigenfunctions if ( PLOT_MODE is not None and not {"all", "kernel_eigen"}.isdisjoint(WHICH_PLOTS) and KERNEL_EIGS_PLT is not None ): _, plot_kernel_eig = make_kernel_evecs_plotter( (pars.train.data, pars.train.kernel), impl_l2, train_data, kernel_eigen, shardings.train.l2.data, pars.test.data, impl_l2_tst, test_data, shardings.test.l2.data, kernel, delay_plot_mode=DELAY_PLOT_MODE, plt_date_range=PLT_DATE_RANGE, ) if KERNEL_EIGS_PLT == "interactive": while True: i = input( "Select kernel eigenfunction " f"0-{pars.train.kernel.num_eigs - 1} to plot, " "or press Enter to continue. " ) if i == "": break else: try: plot_kernel_eig(int(i)) except ValueError: print("Invalid input.") else: for i in KERNEL_EIGS_PLT: plot_kernel_eig(i) if "show" in PLOT_MODE: input("Press any key to continue...") # Compute projections onto kernel eigenfunctions if pars.kernel_projection is not None: io /= str(pars.kernel_projection.data) io /= str(pars.kernel_projection) ds_kernel_proj = compute_kernel_eigen_projections( pars.kernel_projection, kernel_eigen.dual_evecs, weights=1 / pars.train.data.num_delay_samples, ) # Plot representative kernel modes if ( PLOT_MODE is not None and not {"all", "kernel_modes"}.isdisjoint(WHICH_PLOTS) and KERNEL_MODES_PLT is not None ): _, plot_kernel_mode = make_kernel_mode_plotter( ds_kernel_proj, vars=pars.kernel_projection.data.varnames, ocean_only_vars=pars.kernel_projection.data.varnames, ) if KERNEL_MODES_PLT == "interactive": while True: i = input( "Select kernel mode " f"0-{ds_kernel_proj.sizes['mode'] - 1} to plot, " "or press Enter to continue. " ) if i == "": break else: try: plot_kernel_mode(int(i)) except ValueError: print("Invalid input.") else: for i in KERNEL_MODES_PLT: plot_kernel_mode(i) if "show" in PLOT_MODE: input("Press any key to continue...") io -= 2 # Compute expansion coefficients of the response function io /= str(pars.train.data.response) io /= str(pars.train.pred) coeffs = compute_kaf_expansion_coeffs( (pars.train.data, pars.train.kernel), impl_l2, train_data, kernel, kernel_eigen, pars.train.pred.num_steps, pars.train.pred.which_eigs, ) # Make heatmap of KAF response coefficients if PLOT_MODE is not None and not {"all", "expansion_coeffs"}.isdisjoint( WHICH_PLOTS ): plot_kaf_expansion_coeffs( coeffs, title=f"Response: {pars.train.data.response}", ) # Perform time series prediction io /= str(pars.test.data) preds = compute_kaf_preds( pars.train.kernel, impl_l2, train_data, kernel, kernel_eigen, coeffs, impl_l2_tst, test_data, pars.train.pred.which_eigs, ) # Plot running forecast if ( PLOT_MODE is not None and not {"all", "running_pred"}.isdisjoint(WHICH_PLOTS) and LEAD_TIMES_PLT is not None ): _, plot_pred = make_running_pred_plotter( pars.test.data, test_data, preds ) if LEAD_TIMES_PLT == "interactive": while True: i = input( "Select lead time " f"0-{pars.test.num_pred_steps} to plot, " "or press Enter to continue. " ) if i == "": break else: try: plot_pred(int(i)) except ValueError: print("Invalid input.") else: for i in LEAD_TIMES_PLT: plot_pred(i) if "show" in PLOT_MODE: input("Press any key to continue...") # Plot time series forecast if ( PLOT_MODE is not None and not {"all", "pred_timeseries"}.isdisjoint(WHICH_PLOTS) and INITIALIZATION_TIMES_PLT is not None ): _, plot_pred_ts = make_pred_timeseries_plotter( pars.test.data, test_data, preds ) if INITIALIZATION_TIMES_PLT == "interactive": while True: i = input( "Select initialization time " f"0-{pars.test.data.num_samples} to plot, " "or press Enter to continue. " ) if i == "": break else: try: plot_pred_ts(int(i)) except ValueError: print("Invalid input.") else: for i in INITIALIZATION_TIMES_PLT: plot_pred_ts(i) if "show" in PLOT_MODE: input("Press any key to continue...") # Compute forecast skill scores skill_scores = compute_skill_scores( pars.test.data, test_data, preds, ) ts = jnp.arange(pars.test.num_pred_steps + 1) match pars.train.data.time_sampling: case "daily": timestep_str = "days" case "monthly": timestep_str = "months" print( tabulate( jnp.vstack((ts, skill_scores["nrmses"], skill_scores["accs"])).T, headers=[ f"Lead time ({timestep_str})", "Normalized RMSE", "Anomaly Correlation", ], floatfmt=".4f", ) ) # Plot forecast skill scores if PLOT_MODE is not None and not {"all", "skill_scores"}.isdisjoint( WHICH_PLOTS ): plot_forecast_skill_scores(pars.test.data, skill_scores) if __name__ == "__main__": if len(sys.argv) == 2 and (sys.argv[1] == "--help" or sys.argv[1] == "-h"): print(__doc__) else: main()