<jupyter_start><jupyter_text>A new preprocessing scriptWhat does the preprocessing script have to do? Prepare the raw SPCAM data for use in the NN data generator.This entails:1. Computing derivative variables2. Flatten to sample dimension3. Crop levels and potentially lat/lon4. Shuffle training datasetThe preprocessing script should be in the main directory or in a scripts directory. With all the functions part of the cbrain module.Then there is the question how to do the normalization. Like where do I compute the normalization A and b? And how do I do the config file and train valid?  Read the multi-file dataset: xarray vs plain netcdfI don't know how to do benchmarking properly because of the RAM issue.<jupyter_code>DATADIR = '/local/S.Rasp/sp32fbp_andkua/'
!ls $DATADIR/*h1* | head -10<jupyter_output>/local/S.Rasp/sp32fbp_andkua//AndKua_aqua_SPCAM3.0_sp_fbp32.cam2.h1.0000-01-01-00000.nc
/local/S.Rasp/sp32fbp_andkua//AndKua_aqua_SPCAM3.0_sp_fbp32.cam2.h1.0000-01-02-00000.nc
/local/S.Rasp/sp32fbp_andkua//AndKua_aqua_SPCAM3.0_sp_fbp32.cam2.h1.0000-01-03-00000.nc
/local/S.Rasp/sp32fbp_andkua//AndKua_aqua_SPCAM3.0_sp_fbp32.cam2.h1.0000-01-04-00000.nc
/local/S.Rasp/sp32fbp_andkua//AndKua_aqua_SPCAM3.0_sp_fbp32.cam2.h1.0000-01-05-00000.nc
/local/S.Rasp/sp32fbp_andkua//AndKua_aqua_SPCAM3.0_sp_fbp32.cam2.h1.0000-01-06-00000.nc
/local/S.Rasp/sp32fbp_andkua//AndKua_aqua_SPCAM3.0_sp_fbp32.cam2.h1.0000-01-07-00000.nc
/local/S.Rasp/sp32fbp_andkua//AndKua_aqua_SPCAM3.0_sp_fbp32.cam2.h1.0000-01-08-00000.nc
/local/S.Rasp/sp32fbp_andkua//AndKua_aqua_SPCAM3.0_sp_fbp32.cam2.h1.0000-01-09-00000.nc
/local/S.Rasp/sp32fbp_andkua//AndKua_aqua_SPCAM3.0_sp_fbp32.cam2.h1.0000-01-10-00000.nc
ls: write error: Broken pipe<jupyter_text>xarray<jupyter_code>import xarray as xr
%%time
ds = xr.open_mfdataset(
    DATADIR + 'AndKua_aqua_SPCAM3.0_sp_fbp32.cam2.h1.0000-01-*-00000.nc',
    decode_times=False,
    concat_dim='time'
)
%%time
ds = xr.open_mfdataset(
    DATADIR + 'AndKua_aqua_SPCAM3.0_sp_fbp32.cam2.h1.0000-*-*-00000.nc',
    decode_times=False,
    concat_dim='time',
    #parallel=True  # This makes the reading significantly slower...
)<jupyter_output>CPU times: user 18.7 s, sys: 1.17 s, total: 19.8 s
Wall time: 26.4 s<jupyter_text>I mean this is super fast, is this a fluke?<jupyter_code>ds.time
ds.TAP<jupyter_output><empty_output><jupyter_text>Plain netcdfDon't even need to do that if it's that fast... Need to time that later!  Prepare feature and target datasets, actually make them one!<jupyter_code>input_vars = ['TBP', 'QBP', 'VBP', 'PS', 'SOLIN', 'SHFLX', 'LHFLX']
output_vars = ['TPHYSTND', 'PHQ', 'FSNT', 'FSNS', 'FLNT', 'FLNS', 'PRECT']
all_vars = input_vars + output_vars<jupyter_output><empty_output><jupyter_text>Compute derivative variables and cut to correct time stepsOnly include what is currently used!  Compute time step<jupyter_code>dt = ds.time[:2].diff('time').values * 24 * 60 * 60
dt
type(dt)
import numpy as np
dt = dt.astype(np.float16)<jupyter_output><empty_output><jupyter_text>BP variables<jupyter_code># Dictionary defining the physical tendency associated with each variable
# Must be hardcoded!
phy_dict = {
    'TAP': 'TPHYSTND',
    'QAP': 'PHQ',
    'QCAP': 'PHCLDLIQ',
    'QIAP': 'PHCLDICE',
    'VAP': 'VPHYSTND',
    'UAP': 'UPHYSTND'
}
def compute_bp(ds, var):
    """GCM state at beginning of time step before physics.
    ?BP = ?AP - physical tendency * dt
    
    Args:
        ds: entire xarray dataset
        var: BP variable name

    Returns:
        bp: xarray dataarray containing just BP variable, with the first time step cut.
    """
    base_var = var[:-2] + 'AP'
    return (ds[base_var] - ds[phy_dict[base_var]] * dt)[1:]   # Not the first time step
def compute_flx(ds, var):
    """Cuts last time step from flux variables"""
    return ds[var][:-1]
%time TBP = compute_bp(ds, 'TBP')
%time SHFLX = compute_flx(ds, 'SHFLX')
TBP
SHFLX
%time TBP_stack = TBP.stack(sample=('time', 'lat', 'lon'))   # Ok, this fills up my memory... WTH?
TBP_stack
%time SHFLX_stack = SHFLX.stack(sample=('time', 'lat', 'lon'))   # BAM 20G. WTF!?
SHFLX_stack
TBP.lev.size<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12680115273775217, 'fraction_numerical': 0.06652257444764649, 'mean_word_length': 4.53120849933599, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 23, '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': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '522580', 'license': 'MIT', 'n_tokens_mistral': 1859, 'n_tokens_neox': 1686, 'n_words': 372, 'path': 'notebooks/stephans-devlog/3.3-Rewriting-the-data-preprocessing-script.ipynb', 'repo_name': 'ankitesh97/CBRAIN-CAM', 'size': '4164'}