<jupyter_start><jupyter_text>with mask<jupyter_code>ds = xr.merge(dd[k][feature_ml[k]] for k in dd)\
       .groupby('time.month').mean(dim=["time"]).rename({"month":"time"}).where(mask["mask"])
loc_name = list(r_mask)
df_train_ls = []
df_test_ls = []

for loc in loc_name:
    # df_train
    train_tmp = ds.sel(time=ds.time.isin([1,2,3,5,6,7,9,10,11,12]))\
                  .where(r_mask[loc]).to_dataframe().reset_index().dropna()
    train_tmp["region"] = loc
    df_train_ls.append(train_tmp.copy())
    # df_test
    test_tmp = ds.sel(time=ds.time.isin([4,8]))\
                  .where(r_mask[loc]).to_dataframe().reset_index().dropna()
    test_tmp["region"] = loc
    df_test_ls.append(test_tmp.copy())
    
train = pd.concat(df_train_ls)
test = pd.concat(df_test_ls)
del df_train_ls, df_test_ls
gc.collect()
print("shape of train:", train.shape)
print("shape of test:", test.shape)<jupyter_output>shape of train: (39230, 42)
shape of test: (7846, 42)<jupyter_text>save and load<jupyter_code># save
train.to_parquet('../data/c_r_monthly_train.gzip', compression='gzip') 
test.to_parquet('../data/c_r_monthly_test.gzip', compression='gzip') 

print("shape of train:", train.shape)
print("shape of test:", test.shape)

# load 
train_l = pd.read_parquet('../data/c_r_monthly_train.gzip')
test_l = pd.read_parquet('../data/c_r_monthly_test.gzip')

print("shape of train_l:", train_l.shape)
print("shape of test_l:", test_l.shape)

# test
np.testing.assert_array_equal(train[feature_ls].values, train_l[feature_ls].values)
np.testing.assert_array_equal(test[feature_ls].values, test_l[feature_ls].values)
print("passed the test!")<jupyter_output>shape of train: (39230, 42)
shape of test: (7846, 42)
shape of train_l: (39230, 42)
shape of test_l: (7846, 42)
passed the test!
{'fraction_non_alphanumeric': 0.1749156355455568, 'fraction_numerical': 0.030371203599550055, 'mean_word_length': 5.198606271777003, 'pattern_counts': {'":': 1, '<': 7, '<?xml version=': 0, '>': 7, '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': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '531127', 'license': 'MIT', 'n_tokens_mistral': 744, 'n_tokens_neox': 702, 'n_words': 113, 'path': '2_automl/ml_data_prep/clusters_regions_monthly_data_prep.ipynb', 'repo_name': 'zzheng93/code_DSI_India_AutoML', 'size': '1778'}