<jupyter_start><jupyter_text>Train and Val<jupyter_code>data_dir = '/home/jupyter/data/'
args = {'tigge_dir':data_dir + f'tigge/32km/',
    'tigge_vars':['total_precipitation_ens10','total_column_water', '2m_temperature', 'convective_available_potential_energy', 'convective_inhibition'],
    'mrms_dir':data_dir + f'mrms/4km/RadarOnly_QPE_06H/',
    'rq_fn':data_dir + f'mrms/4km/RadarQuality.nc',
#     'const_fn':data_dir + 'tigge/32km/constants.nc',
#     'const_vars':['orog', 'lsm'],
    'data_period':('2018-01', '2019-12'),
    'val_days':1,
    'split':'train',
#     'pure_sr_ratio':8, 
    'tp_log':0.01, 
    'scale':True,
    'ensemble_mode':'stack_by_variable',
    'pad_tigge':15,
    'pad_tigge_channel': True, 
    'idx_stride': 8
    }

save_dir = '/home/jupyter/data/data_patches/'
# dataset_name = 'ensemble_tp_x10_added_vars_TCW-T-CAPE-CIN_log_trans_padded_15_channel'
ds_train = TiggeMRMSDataset(**args)
# pickle.dump(args, open(save_dir+'train/configs/dataset_args.pkl', 'wb'))
#save_images(ds_train, save_dir, 'train')
pickle.dump(ds_train, open(data_dir + f"saved_datasets/traindataset_{dataset_name}.pkl", "wb"))
pickle.dump(args, open(data_dir + f"saved_datasets/traindataset_{dataset_name}_args.pkl", "wb"))
val_args = args
val_args['maxs'] = ds_train.maxs
val_args['mins'] = ds_train.mins
val_args['split'] = 'valid'

#ds_valid = TiggeMRMSDataset(**val_args)
pickle.dump(val_args, open(save_dir+'valid/configs/dataset_args.pkl', 'wb'))
len(ds_valid)
save_images(ds_valid, save_dir, 'valid')
#pickle.dump(ds_valid, open(data_dir + f"saved_datasets/validdataset_{dataset_name}.pkl", "wb"))
#pickle.dump(val_args, open(data_dir + f"saved_datasets/validdataset_{dataset_name}_args.pkl", "wb"))
val_args = pickle.load(open('/home/jupyter/data/data_patches/valid/configs/dataset_args.pkl', 'rb'))

test_args = args
test_args['href_dir'] = data_dir + 'hrefv2/4km/total_precipitation/2020*.nc'
test_args['maxs'] = val_args['maxs']
test_args['mins'] = val_args['mins']
test_args.pop('val_days')
test_args.pop('split')
test_args['first_days'] = 5
test_args['data_period'] = ('2020-01', '2020-12')
# test_dataset_name = dataset_name + f"_first_days_{test_args['first_days']}"

ds_test = TiggeMRMSHREFDataset(**test_args)
save_images(ds_test, save_dir, 'test')
pickle.dump(test_args, open(save_dir+'test/configs/dataset_args.pkl', 'wb'))
len(ds_test)
pickle.dump(ds_test, open(data_dir + f"saved_datasets/testdataset_{test_dataset_name}.pkl", "wb"))
pickle.dump(test_args, open(data_dir + f"saved_datasets/testdataset_{test_dataset_name}_args.pkl", "wb"))
print("check")<jupyter_output><empty_output>