# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernel_info:
#     name: python2
#   kernelspec:
#     display_name: Python 2
#     language: python
#     name: python2
# ---

import xarray as xr
import numpy as np
import os,sys,glob

os.getcwd()
fl_wrfinput=glob.glob('wrfinput*')

fl_wrfinput

ds_disk = xr.open_dataset('wrfinput_d01.save.nc')

dict_rules_suews={
    'LAI_SUEWS':[('lai',1,3)],
    'albDecTr_SUEWS':0,
    'albEveTr_SUEWS':0,
    'albGrass_SUEWS':0,
    'DecidCap_SUEWS':0,
    'porosity_SUEWS':0,
    'GDD_SUEWS':[('gdd',1,5)],
    'HDD_SUEWS':[('hdd',1,6),('day',3,2)],
    'state_SUEWS':[('nsurf',1,7)],
    'soilmoist_SUEWS':[('nsurf',1,7)],
    'surf_var_SUEWS':[('nsurf',1,7)],
    'qn1_av_SUEWS':0,
    'qn1_s_SUEWS':0,
    'dqndt_SUEWS':0,
    'dqnsdt_SUEWS':0,
    'MeltWaterStore':[('nsurf',1,7)],
    'SnowAlb':0,
    'WUDay':[('wu',1,9)],
    'z0m_in':0,
    'zdm_in':0
}


# +
def gen_var_expand(name,rules,var_base=ds_disk['T2'].copy(deep=True)):
    var_new=var_base.rename(name.upper());
    var_attrs=var_base.attrs
    for x_rule in rules:
#         print x_rule
        (name_dim,axis,rep) = x_rule
        var_new=np.repeat(
            var_new.expand_dims(
            name_dim,axis=axis),rep,axis=axis)
    
    var_new.attrs=var_attrs
    var_new.attrs['MemoryOrder']='XYZ'
    var_new.attrs[u'description']=name
    var_new.attrs['stagger']='Z'
    
    return var_new
    
def gen_var_keep(name,rules,var_base=ds_disk['T2'].copy(deep=True)):
    var_new=var_base.rename(name)
    var_new.attrs[u'description']=name
    
    return var_new

def gen_var(name,rules,var_base=ds_disk['T2'].copy()):
    if type(rules) is list:
        var_new=gen_var_expand(name,rules,var_base)
    else:
        var_new=gen_var_keep(name,rules,var_base)
    
    return var_new
        
            
# -

ds_new=xr.Dataset({name.upper():gen_var(name,rules) for name,rules in dict_rules_suews.items()})
# ds_disk.merge(ds_new)
# ds_disk[['T2','Q2']].merge()
# ds_new[['z0m_in','zdm_in']]
ds_new['GDD_SUEWS'].attrs

ds_disk['T2'].attrs

ds_disk.SMOIS.attrs

# ds_disk.merge(ds_new)
ds_disk['SMOIS'].copy()
ds_new['SnowAlb'].attrs
ds_new['LAI_SUEWS'].attrs
for var in ds_disk.variables:
    if 'FieldType' in ds_disk[var].attrs:
        print var,ds_disk[var].ndim,ds_disk[var].attrs['FieldType']

ds_merged=ds_disk.update(ds_new)
for var in ds_merged.data_vars.keys():
    if 'coordinates' in ds_merged[var].attrs:
        print var
        del ds_merged[var].attrs['coordinates']
# ds_merged['ALBBCK']=ds_merged['ALBBCK']+0.2
ds_merged.to_netcdf('wrfinput_d01.new1.nc',mode='w',format='NETCDF4_CLASSIC')

rules=zip(*[('hdd',2,6),('day',4,2)])
var_base=ds_disk['T2'].copy()
var_new=var_base.copy();
var_new=var_new.expand_dims(rules[0],axis=rules[1])
var_new=np.repeat(var_new,)

ds_new.var()

ds_disk['LAI_SUEWS'].shape
ds_disk['SMOIS'].coords['XLAT']

xx=ds_disk['T2'].copy()
yy=np.repeat(xx.expand_dims('nsurf',axis=2),5,axis=2)
yy

ds_disk[['T2','Q2']]

yy=np.repeat(yy,5,axis=2)
yy

ds_disk['LAI_SUEWS']=xx

xx=ds_disk['T2'].copy()
yy=np.repeat(xx.expand_dims('nsurf',axis=2),5,axis=2)
yy

xx=xr.DataArray(np.ones((1,5,7,5)))
xx

ds_disk.coords['XLAT']

for var in ds_disk.data_vars.keys():
    if 'coordinates' in ds_disk[var].attrs:
        print var
        del ds_disk[var].attrs['coordinates']

ds_disk.to_netcdf('xx.nc')

da_test.name='LAI_SUEWS'

da_test.attrs.update(description='test var for SUEWS')
da_test

ds_disk[['T2','Q2']].to_netcdf('xx.nc',mode='w',format='NETCDF3_CLASSIC')

del ds_disk['Q2'].attrs['coordinates']
del ds_disk['T2'].attrs['coordinates']

ds_xx=xr.open_dataset('xx.nc')

ds_xx['T2']



StarcoderdataJupyterScriptsDedupFiltered