<jupyter_start><jupyter_text>ML Dataset statsI want to perform some verification on the ml dataset I generated.<jupyter_code>%load_ext autoreload
%autoreload 2
import dask
import dask.distributed
import logging
import pathlib
import xarray as xr
import torch
import torch.distributions
import seaborn as sns

from crims2s.dask import create_dask_cluster
from crims2s.util import add_biweekly_dim
#ML_DATASET_DIR = '***BASEDIR***/mlready/2021-08-08-test/'
ML_DATASET_DIR = '***BASEDIR***/mlready/2021-08-28-test'
ML_DATASET_DIR = '***BASEDIR***/mlready/2021-09-07-test-set'
ML_DATASET_DIR = '***BASEDIR***/mlready/2021-09-07-cube-train/'
ML_DATASET_DIR = '***BASEDIR***/mlready/2021-09-08-w-features/'
ML_DATASET_DIR = '***HOME***fast_scratch/s2s/mlready/2021-09-14-lean-dry-mask'
ML_DATASET_DIR = '***HOME***fast_scratch/s2s/mlready/2021-09-23-eccc-test'
ML_DATASET_DIR = '***HOME***fast_scratch/s2s/mlready/2021-09-26-ncep'<jupyter_output><empty_output><jupyter_text>Load dataset<jupyter_code>def preprocess_one_example(dataset):
    return dataset.expand_dims('forecast_time')
dataset_files = sorted([x for x in pathlib.Path(ML_DATASET_DIR).iterdir() if '.nc' in x.name])
dataset_files = [f for f in dataset_files if f.stem.endswith('0116')]
dataset_files[:10]
features = xr.open_mfdataset(dataset_files[:3], group='/features', concat_dim='forecast_time', preprocess=preprocess_one_example)
features
features.mean(dim=['realization', 'lead_time', 'forecast_time', 'latitude', 'longitude']).compute()
model = xr.open_mfdataset(dataset_files[:3], group='/model', concat_dim='forecast_time', preprocess=preprocess_one_example)
eccc = xr.open_mfdataset(dataset_files[:3], group='/eccc_parameters', concat_dim='forecast_time', preprocess=preprocess_one_example)
eccc.t2m_mu.isel(biweekly_forecast=1).plot()
ncep = xr.open_mfdataset(dataset_files[:3], group='/ncep_parameters', concat_dim='forecast_time', preprocess=preprocess_one_example)
ncep.t2m_mu.isel(biweekly_forecast=1, forecast_time=0).plot()
ncep.t2m_mu
obs = xr.open_mfdataset(dataset_files[:3], group='/obs', concat_dim='forecast_time', preprocess=preprocess_one_example)
parameters = xr.open_mfdataset(dataset_files[:3], group='/model_parameters', concat_dim='forecast_time', preprocess=preprocess_one_example)
parameters.t2m_mu.isel(biweekly_forecast=1, forecast_time=0).plot()
(ncep - parameters).tp_mu.isel(biweekly_forecast=0, forecast_time=0).plot()
dry_mask = preprocess_one_example(xr.open_dataset(dataset_files[0], group='/dry_mask'))
dry_mask.tp.astype(float).isel(lead_time=1).plot()
eccc.isel(biweekly_forecast=1).tp_cube_root_mu.plot()
dry_mask = xr.open_mfdataset(dataset_files[:3], group='/dry_mask', concat_dim='forecast_time', preprocess=preprocess_one_example)
parameters
model
parameters
obs.isnull().sum().compute()
features
obs
len(obs.lead_time)
46 / 7
some_computed = sample.isnull().sum(dim=['latitude', 'longitude', 'realization'])
some_computed.sum(dim=['variable', 'forecast_time']).compute().x.plot()
some_computed.isel(lead_time=[0, 1]).sum(dim='variable').compute()
sample.sel(variable='sst').isel(lead_time=slice(1, None)).sum(dim=['forecast_time', 'lead_time', 'realization']).compute().x.plot()<jupyter_output><empty_output><jupyter_text>Check y<jupyter_code>target = xr.open_mfdataset(dataset_files[:10], group='/terciles', concat_dim='forecast_time', preprocess=preprocess_one_example)
target = target.compute()
target
target.isnull().isel(forecast_time=0).sum(dim=['category', 'lead_time']).t2m.plot()<jupyter_output><empty_output><jupyter_text>Check obs<jupyter_code>obs = xr.open_mfdataset(dataset_files[:10], group='/obs', concat_dim='forecast_time', preprocess=preprocess_one_example)
obs.isnull().sum(dim=['forecast_time', 'lead_time']).t2m.plot()
model_params = xr.open_mfdataset(dataset_files, group='/model_parameters', preprocess=preprocess_one_example)
model_params<jupyter_output><empty_output><jupyter_text>Check parameters fit  For T2M<jupyter_code>model_biweekly = add_biweekly_dim(model, weeks_12=False)
model_biweekly
reworked_t2m = model_biweekly.isel(lead_time=slice(1, None)).t2m.squeeze().transpose('realization', 'lead_time', 'forecast_time', ...)
reworked_t2m
t2m_data = torch.tensor(reworked_t2m.data.compute())
t2m_mu = torch.tensor(parameters.t2m_mu.squeeze().data.compute())
t2m_sigma = torch.tensor(parameters.t2m_sigma.squeeze().data.compute())
t2m_mu.max()
t2m_mu.shape
t2m_data.shape
distribution = torch.distributions.Normal(t2m_mu, t2m_sigma)
-distribution.log_prob(t2m_data).mean()<jupyter_output><empty_output><jupyter_text>For TP<jupyter_code>tp_data = torch.tensor(model_biweekly.isel(lead_time=-1).tp.squeeze().transpose('realization', 'forecast_time', ...).data.compute())
tp_data.shape
parameters
tp_mu = torch.tensor(parameters.tp_cube_root_mu.squeeze().data.compute())
tp_sigma = torch.tensor(parameters.tp_cube_root_sigma.squeeze().data.compute())
tp_sigma.min()
distribution = torch.distributions.Normal(tp_mu, tp_sigma + 1e-9)
-distribution.log_prob(tp_data ** (1. / 3.)).mean()
tp_alpha = torch.tensor(parameters.tp_alpha.squeeze().data.compute())
tp_beta = torch.tensor(parameters.tp_beta.squeeze().data.compute())
distribution = torch.distributions.Gamma(tp_alpha, tp_beta)
-distribution.log_prob(tp_data + 1e-9).mean()
tp_data.min()
tp_mu.min()
sns.histplot(data=tp_data.flatten())<jupyter_output><empty_output><jupyter_text>NCEP<jupyter_code>sample = preprocess_one_example(xr.open_dataset(dataset_files[1], group='/ncep_parameters'))
sample<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup