<jupyter_start><jupyter_text>Create a coarsened dataset<jupyter_code>def coarsen(ds, ds_coarse, grid=True, kpp=False, debugW=False):
    attributes = {}
    for var in (ds_coarse.variables):
        attributes[var] = ds_coarse[var].attrs
        
    coarse_grid = xgcm.Grid(ds_coarse, periodic=['X'], boundary='extend', metrics={
        ('X',): ['dxC', 'dxG','dxF', 'dxV'], # X distances
        ('Y',): ['dyC', 'dyG','dyF', 'dyU'], # Y distances
        ('Z',): ['drC', 'drF'], # Z distances
        ('X', 'Y'): ['rA', 'rAz', 'rAs', 'rAw'] # Areas
    })
    
    ds_coarsened = xr.Dataset()

    print("Coarsening `ETAN`     ", end='\r')
    ds_coarsened['ETAN'] = ds.ETAN.coarsen(dim={'XC': 10, 'YC': 10}).mean().compute()
    print("Coarsening `MXLDEPTH`     ", end='\r')
    ds_coarsened['MXLDEPTH'] = ds.MXLDEPTH.coarsen(dim={'XC': 10, 'YC': 10}).mean().compute()
    print("Coarsening `SALT`      ", end='\r')
    SALTcontent = (ds.SALT*ds.drF * ds.rA * ds.hFacC).coarsen(dim={'XC': 10, 'YC': 10}).sum().compute()
    newSALT = (SALTcontent/(ds_coarse.drF * ds_coarse.rA * ds_coarse.hFacC)).compute().fillna(0)
    ds_coarsened['SALT'] = newSALT.where(np.isfinite(newSALT), 0).compute()
    print("Coarsening `THETA`  ", end='\r')
    THETAcontent = (ds.THETA*ds.drF * ds.rA * ds.hFacC).coarsen(dim={'XC': 10, 'YC': 10}).sum().compute()
    newTHETA = (THETAcontent/(ds_coarse.drF * ds_coarse.rA * ds_coarse.hFacC)).compute().fillna(0)
    ds_coarsened['THETA'] = newTHETA.where(np.isfinite(newTHETA), 0).compute()
    if kpp == True:
        variables.append("KPPhbl")
        variables.append("KPPdiffS")
        print("Coarsening `KPPhbl`  ", end='\r')
        ds_coarsened['KPPhbl'] = ds.KPPhbl.coarsen(dim={'XC': 10, 'YC': 10}).mean().compute()
        print("Coarsening `KPPdiffS`  ", end='\r')
        ds_coarsened['KPPdiffS'] = ds.KPPdiffS.coarsen(dim={'XC': 10, 'YC': 10}).mean().compute()
    for var in (ds_coarse.variables):
        ds_coarse[var].attrs = attributes[var]
    print("Coarsening `UVEL`      ", end='\r')
    u_transport = ds.UVEL * ds.dyG * ds.hFacW * ds.drF
    u_transport_coarse = u_transport[:, :, :, ::10].coarsen(dim={'YC': 10}).sum()
    newUVEL = (u_transport_coarse / (ds_coarse.dyG * ds_coarse.hFacW * ds_coarse.drF)).fillna(0)
    ds_coarsened['UVEL'] = newUVEL.where(np.isfinite(newUVEL), 0).compute()
    print("Coarsening `VVEL`      ", end='\r')
    v_transport = ds.VVEL * ds.dxG * ds.hFacS * ds.drF
    v_transport_coarse = v_transport[:, :, ::10, :].coarsen(dim={'XC': 10}).sum()
    newVVEL = (v_transport_coarse / (ds_coarse.dxG * ds_coarse.hFacS * ds_coarse.drF)).fillna(0)
    ds_coarsened['VVEL'] = newVVEL.where(np.isfinite(newVVEL), 0).compute()
    print("Computing `WVEL` from non-divergence", end='\r')
    wvel_coarsened = ds.WVEL.coarsen(dim={'XC': 10, 'YC': 10}).mean()
    ds_coarsened['WVEL'] = (coarse_grid.cumsum(coarse_grid.diff(u_transport_coarse, 'X') + coarse_grid.diff(v_transport_coarse, 'Y', boundary='fill'), axis='Z', boundary='fill')/ds_coarse.rA + wvel_coarsened.isel(Zl=0)).compute()

    for coord in list(ds_coarse.coords):
        if coord not in list(ds_coarsened):
            ds_coarsened[coord] = ds_coarse[coord]
            ds_coarsened = ds_coarsened.set_coords(coord)

    for var in list(ds_coarsened.variables):
        ds_coarsened[var].attrs = attributes[var]        
    
    print("Done                                 ", end='\r')
    
    
    returnVars = [ds_coarsened]
    if grid:
        returnVars.append(coarse_grid)
    if debugW:
        returnVars.append(wvel_coarsened)
    
    if len(returnVars) == 1:
        return ds_coarsened
    else:
        return returnVars
ds_coarsened, coarse_grid, wvel_coarsened = coarsen(ds, ds_coarse, debugW=True)
ds_coarsened.to_netcdf(path="/projects/0/topios/hydrodynamic_data/MITgcm/ACC_channel/ACC_ridge_fine_2y_coarsened.nc")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.18833162743091095, 'fraction_numerical': 0.012282497441146366, 'mean_word_length': 3.593419506462985, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '557347', 'license': 'MIT', 'n_tokens_mistral': 1620, 'n_tokens_neox': 1570, 'n_words': 228, 'path': 'data_processing/coarsening.ipynb', 'repo_name': 'daanreijnders/isoneutral-dispersion', 'size': '3908'}