<jupyter_start><jupyter_text>Table of Contents1&nbsp;&nbsp;Preprocessing pipeline1.1&nbsp;&nbsp;Merge Aquaplanet files1.2&nbsp;&nbsp;Compute differences and new vars1.3&nbsp;&nbsp;Create stacked feature and target dataset1.4&nbsp;&nbsp;Reshape variables1.5&nbsp;&nbsp;Test chunking1.6&nbsp;&nbsp;Normalize data1.7&nbsp;&nbsp;Save data1.8&nbsp;&nbsp;Shuffle data long sample dimension2&nbsp;&nbsp;Shuffle manually3&nbsp;&nbsp;Old stuff below4&nbsp;&nbsp;Preprocessed files This is simply a notebook to help me test stuff for the preprocessing script.<jupyter_code>import netCDF4 as nc
import os
import matplotlib.pyplot as plt
import xarray as xr
import numpy as np
import timeit
%matplotlib inline<jupyter_output><empty_output><jupyter_text>Preprocessing pipeline<jupyter_code>feature_vars = ['TAP', 'QAP', 'SHFLX', 'LHFLX', 'SOLIN', 'dTdt_adiabatic', 'dQdt_adiabatic']
target_vars =  ['SPDT', 'SPDQ', 'QRL', 'QRS', 'PRECT', 'FLUT'] 
fn = '/project/meteo/w2w/A6/S.Rasp/SP-CAM/Aquaplanet/AndKua_aqua_SPCAM3.0.cam2.h1.0000-01-01-00000.nc'
import netCDF4 as nc
rg = nc.Dataset(fn)
rg.variables['SOLIN']<jupyter_output><empty_output><jupyter_text>Merge Aquaplanet files<jupyter_code>data_dir = '/project/meteo/w2w/A6/S.Rasp/SP-CAM/Aquaplanet/AndKua_aqua_SPCAM3.0.cam2.h1.0000-01-0[1-3]*'
%ls $data_dir
merged_ds = xr.open_mfdataset(data_dir, decode_times=False)
merged_ds
merged_ds = merged_ds.isel(lev=np.arange(9, 30))
plt.plot(np.diff(merged_ds.time))
cut_time_steps = np.where(np.diff(merged_ds.time) > 2.09e-2)[0]
cut_time_steps<jupyter_output><empty_output><jupyter_text>Compute differences and new vars<jupyter_code>dTdt_adiabatic = (merged_ds['TAP'].diff('time', n=1)) / (0.5 * 60 * 60) - merged_ds['TPHYSTND']
dQdt_adiabatic = (merged_ds['QAP'].diff('time', n=1)) / (0.5 * 60 * 60) - merged_ds['PHQ']
dTdt_adiabatic
plt.plot(dTdt_adiabatic[:, 0, 0, 0])
[s for s in list(dTdt_adiabatic.coords) if 'lev' in s][0]
dTdt_adiabatic[:5, 0, 0, 0].values
((merged_ds['TAP'][1:6, 0, 0, 0].values - merged_ds['TAP'][:5, 0, 0, 0].values) / 
 (0.5 * 60 * 60) - merged_ds['TPHYSTND'][1:6, 0, 0, 0].values)<jupyter_output><empty_output><jupyter_text>Create stacked feature and target dataset<jupyter_code>name_list = (['dTdt_adiabatic_lvl%0i' % lev for lev in range(9,30)] +
             ['dQdt_adiabatic_lvl%i' % lev for lev in range(9,30)])
merged_ds.coords['lev']
# rename times
da_list = [dTdt_adiabatic, dQdt_adiabatic]

for var in feature_vars[:-2]:
    da_list.append(merged_ds[var][:-1])
    if 'lev' in merged_ds[var].dims:
        name_list += [(var + '_lvl%i') % lev for lev in range(9,30)]
    else:
        name_list += [var]
for da in da_list:
    da.coords['time'] = np.arange(da.coords['time'].size)
ilev = 0
for da in da_list:
    if 'lev' in da.coords:
        da.coords['lev'] = np.arange(ilev, ilev+da.coords['lev'].size)
        ilev += da.coords['lev'].size
    else:
        da.expand_dims('lev')
        da.coords['lev'] = ilev
        ilev += 1
feature_da = xr.concat(da_list, dim='lev')
feature_da.name = 'features'
feature_da
clean_time_steps = np.array(feature_da.coords['time'])
clean_time_steps, clean_time_steps.shape
clean_time_steps = np.delete(clean_time_steps, cut_time_steps)
clean_time_steps, clean_time_steps.shape
feature_da.isel(time=clean_time_steps)
name_da = xr.DataArray(name_list, coords=[feature_da.coords['lev']])
name_da.rename({'lev': 'feature_lev'})
merged_ds.TAP.mean(axis=0)
feature_ds = xr.Dataset({'features': feature_da, 'feature_names': name_da})
feature_ds
# Define conversion dict
L_V = 2.5e6   # Latent heat of vaporization is actually 2.26e6
C_P = 1e3 # Specific heat capacity of air at constant pressure
conversion_dict = {
    'SPDT': C_P,
    'SPDQ': L_V,
    'QRL': C_P,
    'QRS': C_P,
    'PRECT': 1e3*24*3600 * 1e-3,
    'FLUT': 1. * 1e-5,
}
da_list = []
for var in target_vars:
    da_list.append(merged_ds[var][1:] * conversion_dict[var])
ilev = 0
for da in da_list:
    if 'lev' in da.coords:
        da.coords['lev'] = np.arange(ilev, ilev+da.coords['lev'].size)
        ilev += da.coords['lev'].size
    else:
        da.expand_dims('lev')
        da.coords['lev'] = ilev
        ilev += 1
target_da = xr.concat(da_list, dim='lev')
target_da<jupyter_output><empty_output><jupyter_text>Reshape variables<jupyter_code>feature_da = feature_da.stack(sample=('time', 'lat', 'lon'))
feature_da
target_da = target_da.stack(sample=('time', 'lat', 'lon'))
feature_da = feature_da.transpose('sample', 'lev')
target_da = target_da.transpose('sample', 'lev')
feature_da<jupyter_output><empty_output><jupyter_text>Test chunking<jupyter_code>tmp_mean = feature_da.mean(axis=0)
t1 = timeit.default_timer()
tmp_mean.values
t2 = timeit.default_timer()
print('Time', t2-t1)
rechunked = feature_da.chunk({'sample': 100000, 'lev': 87})
rechunked
tmp_mean = rechunked.mean(axis=0)
t1 = timeit.default_timer()
tmp_mean.values
t2 = timeit.default_timer()
print('Time', t2-t1)
target_da.chunk({'sample': 100000, 'lev': 87})<jupyter_output><empty_output><jupyter_text>Normalize data<jupyter_code>feature_da = (feature_da - feature_da.mean(axis=0)) / feature_da.std(axis=0)
feature_da[:5, 0].values<jupyter_output><empty_output><jupyter_text>Save data<jupyter_code>target_da.mean(axis=0).values
feature_da = feature_da.reset_index('sample')
feature_da.to_netcdf('/local/S.Rasp/tmp/features_test.nc')
target_da = target_da.reset_index('sample')
target_da.to_netcdf('/local/S.Rasp/tmp/targets_test.nc')
import netCDF4 as nc
rg = nc.Dataset('/local/S.Rasp/tmp/SPCAM_outputs_train_prestack.nc')
rg.variables['targets'][0, :20]
target_da[0, 9:20].values<jupyter_output><empty_output><jupyter_text>Shuffle data long sample dimension<jupyter_code>feature_da
feature_da.coords['sample'].size
a = np.arange(feature_da.coords['sample'].size)
a
np.random.shuffle(a)
a
feature_da.coords['sample'] = a
feature_da
feature_da = feature_da.sortby('sample')
feature_da[:5, :5].values
feature_da = feature_da.chunk({'sample': 100000, 'lev': 87})
feature_da
means = feature_shuf.mean(axis=0)
means.load()
means.rename('feature_mean')
means.attrs['asdf'] ='asdf'
means
norm_ds = xr.Dataset({'1': means, '2': means})
norm_ds
means.to_netcdf('/local/S.Rasp/tmp/means_test.nc')
rg2 = nc.Dataset('/local/S.Rasp/tmp/means_test.nc')
rg2
rg2.close()<jupyter_output><empty_output><jupyter_text>Shuffle manually<jupyter_code>rg = nc.Dataset('/local/S.Rasp/cbrain_data/train_sample_features.nc')
rg
rand_idxs = np.arange(rg.dimensions['sample'].size)
np.random.shuffle(rand_idxs)
rand_idxs[:5]
# Create an equivalent file 
rg_shuffle = nc.Dataset('/local/S.Rasp/cbrain_data/train_sample_shuffle_features.nc', 'w')
for dim_name, dim in rg.dimensions.items():
    print(dim_name)
    rg_shuffle.createDimension(dim_name, dim.size)
rg_shuffle
for var_name, var in rg.variables.items():
    print(var_name)
    rg_shuffle.createVariable(var_name, var.dtype, var.dimensions)
    if var_name != 'features':
        rg_shuffle.variables[var_name][:] = var[:]
rg_shuffle
# Shuffle
from tqdm import tqdm
for i in tqdm(range(rg.dimensions['sample'].size)):
    i_rand = rand_idxs[i]
    rg_shuffle.variables['features'][i] = rg.variables['features'][i_rand]<jupyter_output>0%|          | 0/3530752 [00:00<?, ?it/s][A
  0%|          | 102/3530752 [00:00<58:02, 1013.93it/s][A
  0%|          | 202/3530752 [00:00<58:22, 1008.07it/s][A
  0%|          | 305/3530752 [00:00<58:05, 1012.85it/s][A
  0%|          | 408/3530752 [00:00<57:51, 1016.94it/s][A
  0%|          | 510/3530752 [00:00<57:49, 1017.58it/s][A
  0%|          | 612/3530752 [00:00<57:47, 1018.12it/s][A
  0%|          | 713/3530752 [00:00<57:55, 1015.57it/s][A
  0%|          | 815/3530752 [00:00<58:01, 1013.91it/s][A
  0%|          | 917/3530752 [00:00<58:14, 1009.98it/s][A
  0%|          | 1018/3530752 [00:01<58:25, 1006.85it/s][A
  0%|          | 1121/3530752 [00:01<58:18, 1008.92it/s][A
  0%|          | 1223/3530752 [00:01<58:10, 1011.29it/s][A
  0%|          | 1326/3530752 [00:01<58:10, 1011.18it/s][A
  0%|          | 1427/3530752 [00:01<1:05:20, 900.32it/s][A
  0%|          | 1519/3530752 [00:01<1:07:55, 865.94it/s][A
  0%|          | 1621/3530752 [00:01<1:04:56, 905.68it/s][[...]<jupyter_text>Old stuff below<jupyter_code>rg = nc.Dataset(data_dir + 'AndKua_aqua_SPCAM3.0.cam2.h1.0000-01-01-00000.nc')
rg
rg.variables['TAP'].long_name
for vn, v in rg.variables.items():
    print(vn)
    try:
        print(v.long_name)
    except:
        print('No long name.')
rg.dimensions['lev'].size
QRS = rg.variables['QRS'][:]
QRS.shape
plt.imshow(rg.variables['CLDTOT'][1, :, :])
plt.imshow(QRS[1, 0, :, :])
plt.colorbar()
rg.variables['PHQ']
time = rg.variables['TAP']
time.shape
time[:]
lats = rg.variables['lat'][:]
lats.shape
e = lats.reshape((1, 64, 1))
e.shape
e = np.repeat(e, 5, axis=0)
e.shape
e = np.repeat(e, 10, axis=2)
e.shape
e[0, 0, :]
c = np.repeat(lats, 5 * 10)
c.shape
d = c.reshape((5, 64, 10))
d.shape
d[0,:,0]
cond = (lats > -90) & (lats < 90)
cond
import numpy as np
lat_idxs = np.where((lats > -90) & (lats < 90))
lat_idxs
tap = rg.variables['TAP']
'LAT' in rg.variables.keys()
tap.dimensions
tap[:][:, :, lat_idxs, :]
a = np.arange(20)
a
b = np.arange(0, 0+19)
b
a[b]
nc.Dataset(data_dir + 'AndKua_aqua_SPCAM3.0.cam2.h1.0001-01-01-00000.nc').variables['time'][:]
nc.Dataset('/Volumes/STICK/data/SPCAM/sample/SPCAM_outputs.nc')
64 * 128 * 48
4.5e6 / (64 * 128 * 48)<jupyter_output><empty_output><jupyter_text>Preprocessed files<jupyter_code>data_dir = '/project/meteo/w2w/A6/S.Rasp/SP-CAM/preprocessed_data/'
old_rg = nc.Dataset(data_dir + 'SPCAM_outputs.nc')
new_rg = nc.Dataset(data_dir + 'SPCAM_outputs_flat.nc')
old_rg
new_rg
old_rg.variables['TAP'][:, 0]
new_rg.variables['TAP'][:, 0]<jupyter_output><empty_output><jupyter_text>Nice!<jupyter_code>det_rg = nc.Dataset(data_dir + 'SPCAM_outputs_detailed.nc')
det_rg
det_rg.variables['TAP'][0, 0, :, 0, 0]
mean_rg = nc.Dataset(data_dir + 'SPCAM_mean_detailed.nc')
mean_rg.variables['TAP'][:]
t = mean_rg.variables['TAP']
t.dtype
import timeit
t1 = timeit.default_timer()
t2 = timeit.default_timer()
print('Total time:', t2 - t1, 's')<jupyter_output>Total time: 6.809702608734369e-05 s
{'fraction_non_alphanumeric': 0.17531058962729246, 'fraction_numerical': 0.07306251232498522, 'mean_word_length': 5.084583083383324, 'pattern_counts': {'":': 0, '<': 74, '<?xml version=': 0, '>': 58, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '14', 'dir': 'jupyter-structured-clean-dedup', 'id': '561891', 'license': 'MIT', 'n_tokens_mistral': 4802, 'n_tokens_neox': 4343, 'n_words': 692, 'path': 'notebooks/dev/old_notebooks/explore_aqua_files.ipynb', 'repo_name': 'ankitesh97/CBRAIN-CAM', 'size': '10142'}