# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: CMIP6 2019.10a
#     language: python
#     name: cmip6-201910a
# ---

def merge_and_update_indices(to_merge, df_consts, cols_to_merge, already_merged=pd.DataFrame()):
    """Function to fill in nan values of the indices to
    assets with corresponding constant variables
    with newly calculated indices (and variable_ids)"""
    df_merged = pd.merge(to_merge, 
                          df_consts, 
                          on=cols_to_merge, 
                          how='left', 
                          suffixes=('', '_const'))
    print(df_merged[df_merged['index_consts'].notnull()].count())
    merged = already_merged.append(df_merged[df_merged['index_consts'].notnull()])
    to_merge = df_merged[df_merged['index_consts'].isnull()]
    to_merge = to_merge.drop(columns=['index_consts','variable_id'])
    print(already_merged.shape,merged.shape,to_merge.shape)
    return merged, to_merge


# +
# just need one row per model/ensemble
df_merge_vars = (col_analysis.df.copy()
                .drop(columns=['table_id','variable_id','path','time_range','dcpp_init_year'])
                .drop_duplicates()
                .reset_index(drop=False))

# get the same type of dataframe for the const vars
cols_to_merge = [i for i in df_merge_vars.columns if i != 'index']
df_merge_const = (subset_const.df[cols_to_merge + ['variable_id']]
                  .reset_index(drop=False))

# perform first cut merge trying to match on all metadata
this_merged = pd.merge(df_merge_vars, 
                      df_merge_const, 
                      on=cols_to_merge, 
                      how='left', 
                      suffixes=('_vars', '_consts'))
merged = this_merged[this_merged['index_consts'].notnull()]
to_merge = this_merged[this_merged['index_consts'].isnull()]
print(f"{merged.shape[0]},{to_merge.shape[0]}")

## now drop ensemble member and merge again
# start with "preferred ensemble"
this_merge_const = df_merge_const[df_merge_const.member_id=='r1i1p1f1'].drop(columns=['member_id'])
cols_to_merge = [i for i in cols_to_merge if i != 'member_id']
merged, to_merge = merge_and_update_indices(to_merge, this_merge_const, cols_to_merge, already_merged = merged)
print(f"{merged.shape[0]},{to_merge.shape[0]}")

# now try with any ensemble
this_merge_const = df_merge_const.drop(columns=['member_id']).drop_duplicates(subset=cols_to_merge)
merged, to_merge = merge_and_update_indices(to_merge, this_merge_const, cols_to_merge, already_merged = merged)
print(f"{merged.shape[0]},{to_merge.shape[0]}")
# -

# now drop version, selecting latest version
cols_to_merge = [i for i in cols_to_merge if i != 'version']
this_merge_const = (this_merge_const.sort_values('version')
                    .groupby(cols_to_merge+['variable_id'])
                    .last()
                    .reset_index(drop=False))
merged = merge_and_update_indices(merged, this_merge_const, cols_to_merge)
print(f"{merged['index_consts'].count()}/{merged.shape[0]}")

StarcoderdataJupyterScriptsDedupFiltered