import os
from netCDF4 import Dataset
import numpy as np
import pandas as pd
import xarray
from .....utils.classes.fluxes import Flux
# JvP 20210609: added import statements, MODULE_NAME module level logger
import logging
import sys
import netCDF4 as nc4
MODULE_NAME = __name__[__name__.index(
'TM5'):] if 'TM5' in __name__ else __name__
logger = logging.getLogger(MODULE_NAME)
# Original write function by A. Berchet
[docs]
def write_AB(self, name, flx_file, flx, mode="a"):
"""Write flux to TM5 emission compatible files.
Args:
self (Fluxes): the Fluxes plugin
flx_file (str): the file where to write fluxes
flx (xarray.DataArray): fluxes data to write
mode (str): 'w' to overwrite, 'a' to append
"""
print("WWWWWWWWWWWWWWWWWWWWW")
print(flx_file, name)
print("EEEEEEEEEEEEEEEEEEEEE")
print(flx)
# Here you need to write a function that dumps an xarray "flx" to flx_file
# The structure of "flx" is
# <xarray.DataArray (time: 15, lev: 1, lat: 45, lon: 60)>
return
# New write function by J.C.A. van Peet
[docs]
def write(self, name, flx_file, flx, mode="w", **kwargs):
"""
PURPOSE
Write flux to TM5 emission compatible files.
ARGS
self (Fluxes) = the Fluxes plugin
flx_file (str) = the file where to write fluxes
flx (xarray.DataArray) = fluxes data to write
mode (str) = 'w' to overwrite, 'a' to append
VERSION HISTORY
2.0 09-06-2021 by J.C.A. van Peet
Adapted for TM5.
1.0 ??-??-???? by A. Berchet
See original code above.
"""
# Set the name of this function
PROG_NAME = MODULE_NAME+".write"
# Local logger
logger = logging.getLogger(PROG_NAME)
logger.setLevel(logging.DEBUG)
# Printing options for arrays using np.array2string
# https://docs.scipy.org/doc/numpy/reference/generated/numpy.array2string.html
print_ops = {'max_line_width': 140, 'prefix': ' '*13+' = ',
'formatter': {'float': lambda x: ' %7.2f' % x,
'int': lambda x: ' %7i' % x,
'bool': lambda x: 'T' if x else 'F'}}
# Some debug statements...
logger.debug("")
logger.debug("*"*30)
logger.debug(PROG_NAME+" => DEBUG:")
logger.debug(" name = %s", name)
logger.debug(" flx_file = %s", flx_file)
logger.debug(" flx = %s", flx)
logger.debug(" flx.values = %s",
np.array2string(flx.values, **print_ops))
logger.debug(" flx.time = %s", flx.coords['time'].values)
logger.debug(" mode = %s", mode)
# Set emission grid, and define latitude and longitude.
region = 'glb600x400'
d_lon = 6.0
d_lat = 4.0
n_lon = 60
n_lat = 45
lons = np.linspace(-180.0, 180.0, num=n_lon+1)
lats = np.linspace(-90.0, 90.0, num=n_lat+1)
lon = 0.5*(lons[0:-1]+lons[1:])
lat = 0.5*(lats[0:-1]+lats[1:])
logger.debug(" lon = %s", np.array2string(lon, **print_ops))
logger.debug(" lat = %s", np.array2string(lat, **print_ops))
# Check the shape of the flux data
if (flx.values.shape[2:] != (n_lat, n_lon)):
logger.critical("")
logger.critical("*"*30)
logger.critical(
PROG_NAME+" => ERROR: flux shape does not match with n_lat and n_lon!")
logger.critical(" flx.values.shape = %s", flx.values.shape)
logger.critical(" n_lat = %s", n_lat)
logger.critical(" n_lon = %s", n_lon)
logger.critical("*"*30)
logger.critical("")
raise RuntimeError
# end if
# Convert the time coordinate from the flx xarray to time_[start|mid|end].
# Use the pandas Timestamp function so you can add an offset of a month
time_start = np.array([pd.Timestamp(bagger)
for bagger in flx.coords['time'].values])
time_end = time_start + pd.DateOffset(months=1)
time_mid = time_start + 0.5*(time_end-time_start)
logger.debug("time_start = %s", time_start)
logger.debug("time_end = %s", time_end)
logger.debug("time_mid = %s", time_mid)
# The fluxes from TM5 should be written to file prior to the iteration,
# therefore you will never use "mode = append", but only "mode = write".
# But mode is a parameter to this function, so check if it is set to "a"
# and issue an error if it is.
if (mode != 'w'):
logger.critical("")
logger.critical("*"*30)
logger.critical(PROG_NAME+" => ERROR: mode should be 'w'!")
logger.critical(" mode = %s", mode)
logger.critical("*"*30)
logger.critical("")
raise RuntimeError
# end if
# Here you need to write a function that dumps an xarray "flx" to flx_file
# The structure of "flx" is
# <xarray.DataArray (time: 15, lev: 1, lat: 45, lon: 60)>
# Delete any pre-existing files, and open the output file. Usually,
# selecting "w" when opening a file should create a new file overwriting
# any existing file. But if the file is opened in hdfview (which happens
# regularly), this issues an error. Deleting the file explicitly supresses
# that error. But os.remove raises an error if the file does not exist
# (at least in some Python versions I've tried this), so wrap it in a try
# statement to suppress it.
try:
os.remove(flx_file)
except:
pass
# end try
# Open file
nc_id = nc4.Dataset(flx_file, mode)
# Global dimension
# itime = 6: year, month, day, hour, minute, second
nc_id.createDimension('itime', 6)
# Create group
region_group = nc_id.createGroup(region)
# Create group dimensions
region_group.createDimension('latitude', n_lat)
region_group.createDimension('longitude', n_lon)
# Create CH4 group
ch4_group = region_group.createGroup('CH4')
# Create source group
source_group = ch4_group.createGroup('total')
# Add attributes
source_group.optimize = np.int32(1)
source_group.time_resolution = 'monthly '
# Create time dimensions
source_group.createDimension('nt', flx.values.shape[0])
# Create variables
emis_var = source_group.createVariable(
'emission', np.float64, ('nt', 'latitude', 'longitude'))
time_start_var = source_group.createVariable(
'time_start', np.int16, ('nt', 'itime'))
time_mid_var = source_group.createVariable(
'time_mid', np.int16, ('nt', 'itime'))
time_end_var = source_group.createVariable(
'time_end', np.int16, ('nt', 'itime'))
# Fill variables
# Note that flx.values has shape (time, n_lev, n_lat, n_lon) and that the
# second dimension is always one: emissions happen at the surface...
emis_var[()] = flx.values[:, 0, :, :]
time_start_var[()] = np.array([[bagger.year, bagger.month, bagger.day,
bagger.hour, bagger.minute, bagger.second]
for bagger in time_start])
time_mid_var[()] = np.array([[bagger.year, bagger.month, bagger.day,
bagger.hour, bagger.minute, bagger.second]
for bagger in time_mid])
time_end_var[()] = np.array([[bagger.year, bagger.month, bagger.day,
bagger.hour, bagger.minute, bagger.second]
for bagger in time_end])
# Variable attributes
emis_var.units = '(kg CH4) s-1'
# Close file
nc_id.close()
# ... final debug statements ...
logger.debug("*"*30)
logger.debug("")
# logger.debug("")
# logger.debug("*"*30)
# logger.debug(PROG_NAME+" => Computer says no!")
# logger.debug("*"*30)
# logger.debug("")
# try:
# raise RuntimeError
# except RuntimeError as e:
# #logger.exception("OOPS!")
# logger.critical(e, exc_info=True)
# #raise # => Will display the traceback on screen a second time
# sys.exit() # => Just exit.
# end try
return
# end function write
