{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":56537,"databundleVersionId":8015876,"sourceType":"competition"}],"dockerImageVersionId":30698,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-05-23T12:50:54.770516Z","iopub.execute_input":"2024-05-23T12:50:54.771018Z","iopub.status.idle":"2024-05-23T12:50:54.777726Z","shell.execute_reply.started":"2024-05-23T12:50:54.770973Z","shell.execute_reply":"2024-05-23T12:50:54.775968Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!git clone https://github.com/leap-stc/ClimSim","metadata":{"execution":{"iopub.status.busy":"2024-05-23T12:50:54.780431Z","iopub.execute_input":"2024-05-23T12:50:54.780898Z","iopub.status.idle":"2024-05-23T12:51:09.486769Z","shell.execute_reply.started":"2024-05-23T12:50:54.780856Z","shell.execute_reply":"2024-05-23T12:51:09.485436Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.chdir('/kaggle/working/ClimSim')","metadata":{"execution":{"iopub.status.busy":"2024-05-23T12:51:09.488594Z","iopub.execute_input":"2024-05-23T12:51:09.48903Z","iopub.status.idle":"2024-05-23T12:51:09.494865Z","shell.execute_reply.started":"2024-05-23T12:51:09.488996Z","shell.execute_reply":"2024-05-23T12:51:09.493733Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install netCDF4\nimport netCDF4","metadata":{"execution":{"iopub.status.busy":"2024-05-23T12:51:09.497769Z","iopub.execute_input":"2024-05-23T12:51:09.498215Z","iopub.status.idle":"2024-05-23T12:51:21.859369Z","shell.execute_reply.started":"2024-05-23T12:51:09.498176Z","shell.execute_reply":"2024-05-23T12:51:21.858086Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# import necessary packages\n\nfrom climsim_utils.data_utils import *\n\n# set variable names\n\nv2_inputs = ['state_t',\n             'state_q0001',\n             'state_q0002',\n             'state_q0003',\n             'state_u',\n             'state_v',\n             'state_ps',\n             'pbuf_SOLIN',\n             'pbuf_LHFLX',\n             'pbuf_SHFLX',\n             'pbuf_TAUX',\n             'pbuf_TAUY',\n             'pbuf_COSZRS',\n             'cam_in_ALDIF',\n             'cam_in_ALDIR',\n             'cam_in_ASDIF',\n             'cam_in_ASDIR',\n             'cam_in_LWUP',\n             'cam_in_ICEFRAC',\n             'cam_in_LANDFRAC',\n             'cam_in_OCNFRAC',\n             'cam_in_SNOWHICE',\n             'cam_in_SNOWHLAND',\n             'pbuf_ozone', # outside of the upper troposphere lower stratosphere (UTLS, corresponding to indices 5-21), variance in minimal for these last 3 \n             'pbuf_CH4',\n             'pbuf_N2O']\n\nv2_outputs = ['ptend_t',\n              'ptend_q0001',\n              'ptend_q0002',\n              'ptend_q0003',\n              'ptend_u',\n              'ptend_v',\n              'cam_out_NETSW',\n              'cam_out_FLWDS',\n              'cam_out_PRECSC',\n              'cam_out_PRECC',\n              'cam_out_SOLS',\n              'cam_out_SOLL',\n              'cam_out_SOLSD',\n              'cam_out_SOLLD']\n\nvertically_resolved = ['state_t', \n                       'state_q0001', \n                       'state_q0002', \n                       'state_q0003', \n                       'state_u', \n                       'state_v', \n                       'pbuf_ozone', \n                       'pbuf_CH4', \n                       'pbuf_N2O', \n                       'ptend_t', \n                       'ptend_q0001', \n                       'ptend_q0002', \n                       'ptend_q0003', \n                       'ptend_u', \n                       'ptend_v']\n\nablated_vars = ['ptend_q0001',\n                'ptend_q0002',\n                'ptend_q0003',\n                'ptend_u',\n                'ptend_v']\n\nv2_vars = v2_inputs + v2_outputs\n\ntrain_col_names = []\nablated_col_names = []\nfor var in v2_vars:\n    if var in vertically_resolved:\n        for i in range(60):\n            train_col_names.append(var + '_' + str(i))\n            if i < 12 and var in ablated_vars:\n                ablated_col_names.append(var + '_' + str(i))\n    else:\n        train_col_names.append(var)\n\ninput_col_names = []\nfor var in v2_inputs:\n    if var in vertically_resolved:\n        for i in range(60):\n            input_col_names.append(var + '_' + str(i))\n    else:\n        input_col_names.append(var)\n\noutput_col_names = []\nfor var in v2_outputs:\n    if var in vertically_resolved:\n        for i in range(60):\n            output_col_names.append(var + '_' + str(i))\n    else:\n        output_col_names.append(var)\n\nassert(len(train_col_names) == 17 + 60*9 + 60*6 + 8)\nassert(len(input_col_names) == 17 + 60*9)\nassert(len(output_col_names) == 60*6 + 8)\nassert(len(set(output_col_names).intersection(set(ablated_col_names))) == len(ablated_col_names))\n\n# initialize data_utils object\n\ngrid_path = './grid_info/ClimSim_low-res_grid-info.nc'\nnorm_path = './preprocessing/normalizations/'\n\ngrid_info = xr.open_dataset(grid_path)\ninput_mean = xr.open_dataset(norm_path + 'inputs/input_mean.nc')\ninput_max = xr.open_dataset(norm_path + 'inputs/input_max.nc')\ninput_min = xr.open_dataset(norm_path + 'inputs/input_min.nc')\noutput_scale = xr.open_dataset(norm_path + 'outputs/output_scale.nc')\n\ndata = data_utils(grid_info = grid_info, \n                  input_mean = input_mean, \n                  input_max = input_max, \n                  input_min = input_min, \n                  output_scale = output_scale)\n\ndata.set_to_v2_vars()\n\n# do not normalize\ndata.normalize = False\n\n# create training data\n\n# set data path for training data\ndata.data_path = '/ocean/projects/atm200007p/jlin96/neurips_proj/e3sm_train/'\n\n# set regular expressions for selecting training data\ndata.set_regexps(data_split = 'train', \n                 regexps = ['E3SM-MMF.mli.000[1234567]-*-*-*.nc', # years 1 through 7\n                            'E3SM-MMF.mli.0008-01-*-*.nc']) # first month of year 8\n\n# set temporal subsampling\ndata.set_stride_sample(data_split = 'train', stride_sample = 7)\n\n# create list of files to extract data from\ndata.set_filelist(data_split = 'train')\n\n# save numpy files of training data\ndata_loader = data.load_ncdata_with_generator(data_split = 'train')\nnpy_iterator = list(data_loader.as_numpy_iterator())\nnpy_input = np.concatenate([npy_iterator[x][0] for x in range(len(npy_iterator))])\nnpy_output = np.concatenate([npy_iterator[x][1] for x in range(len(npy_iterator))])\ntrain_npy = np.concatenate([npy_input, npy_output], axis = 1)\ntrain_index = [\"train_\" + str(x) for x in range(train_npy.shape[0])]\n\ntrain = pd.DataFrame(train_npy, index = train_index, columns = train_col_names)\ntrain.index.name = 'sample_id'\nprint('dropping cam_in_SNOWHICE because of strange values')\ntrain.drop('cam_in_SNOWHICE', axis=1, inplace=True)\n\n# ASSERT, SHAPE, CSV, PRINT\nassert sum(train.isnull().any()) == 0\nprint(train.shape)\ntrain.to_csv('recreated/user_data/train.csv')\nprint('finished creating train data')\n\n\n# create train output weighting\n\ntrain_output_weighting = train[output_col_names].copy()\nweighting_dict = {}\nfor col_name in output_col_names:\n    weighting_dict[col_name] = 1/max(train_output_weighting[col_name].std(), 1e-15)\n    train_output_weighting[col_name] = weighting_dict[col_name]\ntrain_output_weighting[ablated_col_names] = 0\n\n# ASSERT, SHAPE, CSV, PRINT\nassert sum(train_output_weighting.isnull().any()) == 0\nprint(train_output_weighting.shape)\ntrain_output_weighting.to_csv('recreated/user_data/train_output_weighting.csv')\nprint('finished creating train_output_weighting data')\n\n# creating sample submission\nsample_submission = pd.DataFrame(index = shuffled_test_index, columns = output_col_names)\nsample_submission.index.name = 'sample_id'\nfor col_name in output_col_names:\n    sample_submission[col_name] = weighting_dict[col_name]\nsample_submission[ablated_col_names] = 0\n\n# ASSERT, SHAPE, CSV, PRINT\nassert sum(sample_submission.isnull().any()) == 0\nprint(sample_submission.shape)\nsample_submission.head(625000).to_csv('recreated/user_data/sample_submission.csv')\nprint('finished creating sample_submission data')","metadata":{"execution":{"iopub.status.busy":"2024-05-23T12:56:26.169616Z","iopub.execute_input":"2024-05-23T12:56:26.170134Z","iopub.status.idle":"2024-05-23T12:56:26.413364Z","shell.execute_reply.started":"2024-05-23T12:56:26.17009Z","shell.execute_reply":"2024-05-23T12:56:26.411633Z"},"trusted":true},"execution_count":null,"outputs":[]}]}