source: trunk/TOOLS/GENERATE_VEGET_MAP/ESA-LUH2/gen2orc.py

import argparse, datetime, numpy as np, netCDF4, sys
from grid import Grid
from cwt import cwt_kg31to7, cwt_gen2orc, cwt_gen2orc16
from basics import printStat, writePFTMAP
CWT1 = cwt_kg31to7
CWT2 = cwt_gen2orc # cwt_gen2orc for standard 15 PFTs, cwt_gen2orc16 for additional urbain PFT
debug = False

parser = argparse.ArgumentParser()
parser.add_argument("year", action="store", type=int, help="year to process)")
parser.add_argument("nlat", action="store", type=int, help="number of latitude steps (143 for LMDZ resolution)")
parser.add_argument("-part", action="store", type=int, default=0, help="part (1 or 2 for nlat=14400 to split in two parts", required = False)
args = parser.parse_args()
year = args.year
nlat = args.nlat
part = args.part
nlon = 144 if nlat == 143 else nlat * 2
if part == 0: lat1=0; lat2=nlat
elif part == 1: lat1=0; lat2=int(nlat/2)
elif part == 2: lat1=int(nlat/2); lat2=nlat
step = 180. / nlat
luh_year0 = 1992
esa = "v2.0.7b" if year <= 2015 else "v2.1.1"
print("GEN2ORC for %s, %sx%s, %s:%s" % (year, nlat, nlon, lat1, lat2))

atts = dict(description = "ORCHIDEE PFT map derived from ESACCI Land Cover map %(esa)s for %(year)s, aggregated using user-tool v4.3, cross-walking table v2.4, C3/C4 for grasses from C.Still (May 2018), C3/C4 for crops from LUH2 (2022)" % dict(esa = esa, year = year),
            PFT1  = "bare soil",
            PFT2  = "tropical broadleaf evergreen",
            PFT3  = "tropical broadleaf raingreen",
            PFT4  = "temperate needleleaf evergreen",
            PFT5  = "temperate broadleaf evergreen",
            PFT6  = "temperate broadleaf summergreen",
            PFT7  = "boreal needleleaf evergreen",
            PFT8  = "boreal broadleaf summergreen",
            PFT9  = "boreal needleleaf summergreen",
            PFT10 = "temperate C3 grass",
            PFT11 = "C4 grass",
            PFT12 = "C3 agriculture",
            PFT13 = "C4 agriculture",
            PFT14 = "tropical C3 grass",
            PFT15 = "boreal C3 grass",
            # PFT16 = "urban", # add this attribute for 16 PFTs (with additional urbain PFT)
            contact = "V.Bastrikov (vladislav.bastrikov@lsce.ipsl.fr), P.Peylin (peylin@lsce.ipsl.fr)",
            web = "https://orchidas.lsce.ipsl.fr/dev/lccci/",
            date = datetime.datetime.now().strftime("%d/%m/%Y %H:%M"))

path_gen = "/home/surface5/vbastri/BDD/ESACCI-LC/GENERIC_PFT/%(step)s/ESACCI-LC-L4-LCCS-Map-300m-P1Y-aggregated-%(step).6fDeg-%(year)s-%(esa)s.nc" % dict(step = step, year = year, esa = esa)
path_luh = "/home/surface5/vbastri/BDD/LUH2/TRENDY2022/LUH2v11_1992-2022_%(nlat)sx%(nlon)s.nc" % dict(nlat = nlat, nlon = nlon)
path_kg  = "/home/surface5/vbastri/BDD/KG2017/KG_%(nlat)sx%(nlon)s.nc" % dict(nlat = nlat, nlon = nlon)
path_c4  = "/home/surface5/vbastri/BDD/C4FRAC/C4frac_%(nlat)sx%(nlon)s.nc" % dict(nlat = nlat, nlon = nlon)
path_wb  = "/home/surface5/vbastri/BDD/ESACCI-LC/WATER/WB_%(nlat)sx%(nlon)s.nc" % dict(nlat = nlat, nlon = nlon)
path_out = "PFTmap_%(year)s_%(nlat)sx%(nlon)s%(end)s.nc" % dict(year = year, nlat = nlat, nlon = nlon, end = "_%(lat1)s-%(lat2)s" % dict(lat1 = lat1, lat2 = lat2) if part!=0 else "")
GEN_LIST = ["Tree_Broadleaf_Evergreen", "Tree_Broadleaf_Deciduous", "Tree_Needleleaf_Evergreen", "Tree_Needleleaf_Deciduous", "Shrub_Broadleaf_Evergreen", "Shrub_Broadleaf_Deciduous", "Shrub_Needleleaf_Evergreen", "Shrub_Needleleaf_Deciduous", "Natural_Grass", "Crops", "Bare_soil", "Water", "Snow_Ice", "Urban", "NoData"]
GEN_KEYS = ["TrBrEv", "TrBrDe", "TrNeEv", "TrNeDe", "ShBrEv", "ShBrDe", "ShNeEv", "ShNeDe", "NatGr", "Crops", "BS",  "Water", "SnowIce", "Urban", "NoData"]
KG_KEYS = CWT1[0][1:]
ORC_KEYS = CWT2[0][2:]

# Function to safely divide two streams: if denominator equals zero, division result is set to zero
# During division zero values in denominator are changed to 1 so that python does not warn on errors
def divide(num, den): return np.where(den == 0, 0, num / np.where(den == 0, 1, den))

# Function to make plots for data streams
def plot(**kwargs):
    for name, data in kwargs.items():
        for idx, item in enumerate(data if len(data.shape)>2 else [data]):
            grid.plotmap(np.ma.masked_where(item==0, item), dpi=150, vmin=0, vmax=1,
                         filename="%s%s.png" % (name.lower(), idx if len(data.shape)>2 else ""), title="%s%s" % (name.upper(), idx if len(data.shape)>2 else ""),
                         info="min=%.3g\nmax=%.3g\nmean=%.3g\nnzpts=%s" % (item.min(), item.max(), grid.meanmask(np.ma.masked_where(item==0, item)), (item>0).sum()))

# Read and normalize generic PFTmap
print("\nREAD GENERIC PFTMAP")
nc = netCDF4.Dataset(path_gen)
lat = nc.variables["lat"][lat1:lat2]
lon = nc.variables["lon"][:]
gen = np.zeros((len(GEN_LIST), len(lat), len(lon)))
for idx, item in enumerate(GEN_LIST): gen[idx] = nc.variables[item][lat1:lat2].astype("f8").filled(0)
nc.close()
printStat(gen=gen)
gen = divide(gen, gen.sum(axis=0))
printStat(genn=gen)
if debug: grid = Grid(lat=lat, lon=lon)
if debug: plot(gen=gen, gensum=gen.sum(axis=0))

# Read and normalize Koeppen-Geiger map with 31 classes
print("\nREAD KG")
nc = netCDF4.Dataset(path_kg)
if not np.allclose(lat, nc.variables["lat"][lat1:lat2]) or not np.allclose(lon, nc.variables["lon"][:]): print("Different grid"); sys.exit()
kg31 = nc.variables["kg"][:,lat1:lat2].astype("f8").filled(0)
nc.close()
printStat(kg31=kg31)
kg31 = divide(kg31, kg31.sum(axis=0))
printStat(kg31n=kg31)

# Aggregate KG into 7 groups
print("Check CWT1 kg31to7 = %s" % np.all([np.array(item[1:]).sum() == 1 for item in CWT1[1:]]))
kg = np.zeros((len(KG_KEYS), len(lat), len(lon)))
for idx, rule in enumerate(CWT1[1:]): kg = kg + kg31[idx] * np.array(rule[1:])[:,None,None]
printStat(kg=kg)
if debug: plot(kg=kg)
# Fill non-defined values in KG with the dominant class across each latitude (or last class if no data in a latitude)
domclass = np.where(kg.sum(axis=(0,2)) == 0, len(kg)-1, kg.sum(axis=2).argmax(axis=0))
for idx in range(len(lat)): kg[domclass[idx], idx, kg[:,idx,:].sum(axis=0) == 0] = 1
printStat(kgf=kg)
if debug: plot(kgf=kg)

print("\nREAD C4")
nc = netCDF4.Dataset(path_c4)
if not np.allclose(lat, nc.variables["lat"][lat1:lat2]) or not np.allclose(lon, nc.variables["lon"][:]): print("Different grid"); sys.exit()
c4 = nc.variables["C4frac"][lat1:lat2].astype("f8").filled(0)
nc.close()
printStat(c4=c4)
if debug: plot(c4=c4)

print("\nREAD WB")
nc = netCDF4.Dataset(path_wb)
if not np.allclose(lat, nc.variables["lat"][lat1:lat2]) or not np.allclose(lon, nc.variables["lon"][:]): print("Different grid"); sys.exit()
wb = nc.variables["wb"][:,lat1:lat2].astype("f8").filled(0)
nc.close()
printStat(wb=wb)
wb = divide(wb, wb.sum(axis=0))
printStat(wbn=wb)
ocean = np.where(wb[0] > 0.999, 1, 0) # ocean 0/1-mask
coast = np.where((wb[0] > 0) & (wb[0] <= 0.999), 1, 0) # coast 0/1-mask
inland = np.where(wb[0] == 0, wb[2], 0) # inland water fraction excluding coasts
print("ocean_npts=%s, coast_npts=%s, inland_npts=%s" % (ocean.sum(), coast.sum(), (inland>0).sum()))
if debug: plot(wb=wb, ocean=ocean, coast=coast, inland=inland)

# read C4 map from LUH2 for crops
print("\nREAD LUH")
nc = netCDF4.Dataset(path_luh)
if not np.allclose(lat, nc.variables["lat"][lat1:lat2]) or not np.allclose(lon, nc.variables["lon"][:]): print("Different grid"); sys.exit()
c3ann = nc.variables["c3ann"][year-luh_year0, lat1:lat2].astype("f8").filled(0)
c4ann = nc.variables["c4ann"][year-luh_year0, lat1:lat2].astype("f8").filled(0)
c3per = nc.variables["c3per"][year-luh_year0, lat1:lat2].astype("f8").filled(0)
c4per = nc.variables["c4per"][year-luh_year0, lat1:lat2].astype("f8").filled(0)
c3nfx = nc.variables["c3nfx"][year-luh_year0, lat1:lat2].astype("f8").filled(0)
total = c3ann + c4ann + c3per + c4per + c3nfx
luhc4 = divide(c4ann + c4per, total)
nc.close()
printStat(luhc4=luhc4)
if debug: plot(luhc4=luhc4)

# aggregate generic PFTs to ORCHIDEE PFTs
print("\nAGGREGATE GENERIC PFTs to ORCHIDEE PFTs")
print("Check CWT2 gen2orc = %s" % np.all([np.array(item[2:]).sum() == 1 for item in CWT2[1:]]))
orc = np.zeros((len(ORC_KEYS), len(lat), len(lon)))
for rule in CWT2[1:]:
    idx = GEN_KEYS.index(rule[0])
    cond = rule[1]
    fracs = np.array(rule[2:])
    if cond in KG_KEYS: coef = kg[KG_KEYS.index(cond)]
    elif cond in ["C3", "C4"]: coef = c4 if cond == "C4" else 1-c4
    elif cond in ["luhC3", "luhC4"]: coef = luhc4 if cond == "luhC4" else 1-luhc4
    elif "/" in cond: coef = kg[KG_KEYS.index(cond.split("/")[0])] * (c4 if cond.split("/")[1] == "C4" else 1-c4)
    elif cond == "inland": coef = inland
    elif cond == "": coef = 1
    else: print("Unknown condition: %s" % cond); sys.exit()
    orc = orc + gen[idx] * coef * fracs[:,None,None]
printStat(orc=orc)

# set small fractions to zero (< 0.001)
orc = np.where(orc < 0.001, 0, orc)
# set 1 to bare soil on the coast if no other PFTs
orc[0] = np.where((coast == 1) & (orc[1:].sum(axis=0) < 0.001), 1, orc[0])
# normalize
orc = divide(orc, orc.sum(axis=0))
# apply land-sea mask
orc = np.ma.array(orc)
orc.mask = (ocean == 1)
# filter out Antarctica (lat < S60)
orc[:, lat < -60, :] = np.ma.masked

printStat(orcn=orc)
if debug: plot(orc=orc, orcsum=orc.sum(axis=0))

writePFTMAP(orc, lat=lat, lon=lon, year=year, filename=path_out, atts=atts)