source: trunk/TOOLS/GENERATE_VEGET_MAP/ESA-LUH2/post_PFT.py @ 6267

# $HeadURL$
# $Date$
# $Revision$

import datetime, netCDF4, numpy as np, os

path_luh = "/home/surface7/vuichard/LandUseMap_CMIP6/ESA2.0.7b_19PFT_quarterdeg_final//landusemap_%(year)s.nc"
path_orc = "/home/satellites7/vbastri/LC/orchidee/ESA2.0.7b/19PFT_quarterdeg_final/PFTmap_2010.nc"
path_out = "/home/satellites7/vbastri/LC/orchidee/ESA-LUH2-final/PFTmap_%(year)s.nc"

# Define global attributes to be written in the output netcdf-files
atts = dict(description = "ORCHIDEE PFT map based on the ESA CCI Land Cover map v2.0.7b for 2010 aggregated by the user-tool v3.13 with the cross-walking table v2.4, followed by historical reconstrution using G. Hurtt data LU2v2h with merging rules v2.0",
            contact = "V.Bastrikov (vladislav.bastrikov@lsce.ipsl.fr), P.Peylin (philippe.peylin@lsce.ipsl.fr), N.Vuichard (nicolas.vuichard@lsce.ipsl.fr)",
            web = "https://orchidas.lsce.ipsl.fr/dev/LCCCI.php",
            date = datetime.datetime.now().strftime("%d/%m/%Y %H:%M"),
            resolution = "Regular lat/lon 0.25-degree")

# Generate latitudes/longitudes arrays and calculate areas of grid cells
step = 0.25
lat = np.arange(90 - step/2., -90, -step)
lon = np.arange(-180 + step/2., 180, step)
area = 2 * 6.371229E6**2 * np.repeat(step, len(lon))[None,:] * np.pi/180 * np.sin(0.5 * np.repeat(step, len(lat))[:,None] * np.pi/180) * np.cos(lat[:,None] * np.pi/180)

# Load ORCHIDEE PFT data for year 2010 (with 19 PFTs)
nc = netCDF4.Dataset(path_orc)
orc = nc.variables["maxvegetfrac"][0]
nc.close()

# Function to extend land by 1 pixel outside land-sea mask
def extend(veg):
    newveg = np.ma.masked_all((len(veg), len(lat), len(lon)))
    loc_area = np.where(veg[0].mask == False, area, 0)
    wveg = veg * loc_area
    for ilat in range(len(lat)):
        for ilon in range(len(lon)):
            if veg[0,ilat,ilon] is not np.ma.masked:
                newveg[:,ilat,ilon] = veg[:,ilat,ilon]
            else:
                if wveg[:, max(0, ilat-1) : ilat+2, max(0, ilon-1) : ilon+2].count() == 0: continue
                newveg[:,ilat,ilon] = wveg[:, max(0, ilat-1) : ilat+2, max(0, ilon-1) : ilon+2].sum(axis=-1).sum(axis=-1) / loc_area[max(0, ilat-1) : ilat+2, max(0, ilon-1) : ilon+2].sum()
    return newveg

for year in range(850, 2016):
    print year
    nc = netCDF4.Dataset(path_luh % dict(year=year))
    veg = nc.variables["YEARLYMAXVEGETFRAC_NORMALIZE"][0,:,::-1,:]
    nc.close

    print "Veg range [%s - %s], total range [%s - %s], %s pts" % (veg.min(), veg.max(), veg.sum(axis=0).min(), veg.sum(axis=0).max(), veg.count() / 15.)

    # Make LSM uniform across pfts
    vegsum = veg.sum(axis=0)
    for pft in range(15):
        ys, xs = np.where((vegsum.mask == False) & (veg[pft].mask == True))
        if len(ys) > 0:
            print "Add %s pts to PFT%s" % (len(ys), pft+1)
            veg[pft, ys, xs] = 0
    print "Veg range [%s - %s], total range [%s - %s], %s pts" % (veg.min(), veg.max(), veg.sum(axis=0).min(), veg.sum(axis=0).max(), veg.count() / 15.)
    
    # Clip to [0,1] and normalize
    veg = np.ma.clip(veg, 0, 1)
    veg = veg / veg.sum(axis=0)
    print "Clip to [0,1] and normalize"
    print "Veg range [%s - %s], total range [%s - %s], %s pts" % (veg.min(), veg.max(), veg.sum(axis=0).min(), veg.sum(axis=0).max(), veg.count() / 15.)

    # Fill masked inland pts with bare soil
    ys, xs = np.ma.where((orc[-2] > 0) & (veg[0].mask == True))
    print "Add %s pts as bare soil" % len(ys)
    veg[0, ys, xs] = 1
    for pft in range(1,15): veg[pft, ys, xs] = 0
    print "Veg range [%s - %s], total range [%s - %s], %s pts" % (veg.min(), veg.max(), veg.sum(axis=0).min(), veg.sum(axis=0).max(), veg.count() / 15.)

    # Extend by 1 pixel outside the land
    for idx in range(1):
        print "Extend by 1 pixel"
        veg = extend(veg)
        print "Veg range [%s - %s], total range [%s - %s], %s pts" % (veg.min(), veg.max(), veg.sum(axis=0).min(), veg.sum(axis=0).max(), veg.count() / 15.)

    # Write output file
    if os.path.exists(path_out % dict(year=year)): raw_input("The file already exists! Press CTRL-C to exit or any key to continue")
    ncout = netCDF4.Dataset(path_out % dict(year=year), "w", format = "NETCDF3_CLASSIC")
    ncout.setncatts(atts)
    ncout.createDimension("lat", len(lat))
    ncout.createDimension("lon", len(lon))
    ncout.createDimension("time_counter", None)
    ncout.createDimension("veget", len(veg))
    ncout.createVariable("lat", "f8", ("lat",))
    ncout.createVariable("lon", "f8", ("lon",))
    ncout.createVariable("maxvegetfrac", "f4", ("time_counter", "veget", "lat", "lon"), fill_value = 1.e+20)
    ncout.createVariable("time_counter", "f4", ("time_counter",))
    ncout.createVariable("veget", "i4", ("veget",))
    ncout.variables["lat"].setncatts(dict(long_name="Latitude", valid_max=90., valid_min=-90., units="degrees_north", axis="Y"))
    ncout.variables["lon"].setncatts(dict(long_name="Longitude", valid_max=180., valid_min=-180., units="degrees_east", axis="X", modulo=360., topology="circular"))
    ncout.variables["maxvegetfrac"].setncatts(dict(name = "maxvegetfrac", long_name = "Vegetation types", units = "-"))
    ncout.variables["time_counter"].setncatts(dict(units = "years since 0-1-1", calendar = "gregorian", axis = "T"))
    ncout.variables["veget"].setncatts(dict(valid_min = 1, valid_max = len(veg), long_name = "Vegetation classes", units = "-", axis = "Z"))
    ncout.variables["lat"][:] = lat
    ncout.variables["lon"][:] = lon
    ncout.variables["maxvegetfrac"][0] = veg
    ncout.variables["time_counter"][0] = year
    ncout.variables["veget"][:] = range(1, len(veg)+1)
    ncout.close()