USDA-ARS-NWRC/awsm

View on GitHub
awsm/models/pysnobal/pysnobal_io.py

Summary

Maintainability
A
0 mins
Test Coverage
import os
from copy import copy
import logging
from datetime import datetime

import netCDF4 as nc
import pandas as pd
from spatialnc.proj import add_proj

FREEZE = 273.16


class PysnobalIO():

    OUTPUT_VARIABLES = {
        'net_radiation': {
            'units': 'W m-2',
            'description': 'Average net all-wave radiation',
            'ipysnobal_var': 'R_n_bar'
        },
        'sensible_heat': {
            'units': 'W m-2',
            'description': 'Average sensible heat transfer',
            'ipysnobal_var': 'H_bar'
        },
        'latent_heat': {
            'units': 'W m-2',
            'description': 'Average latent heat exchange',
            'ipysnobal_var': 'L_v_E_bar'
        },
        'snow_soil': {
            'units': 'W m-2',
            'description': 'Average snow/soil heat exchange',
            'ipysnobal_var': 'G_bar'
        },
        'precip_advected': {
            'units': 'W m-2',
            'description': 'Average advected heat from precipitation',
            'ipysnobal_var': 'M_bar'
        },
        'sum_energy_balance': {
            'units': 'W m-2',
            'description': 'Average sum of energy balance terms for snowcover',
            'ipysnobal_var': 'delta_Q_bar'
        },
        'evaporation': {
            'units': 'kg m-2 (equivalent to mm of water)',
            'description': 'Total evaporation and sublimation per unit area from surface of snowpack',  # noqa
            'ipysnobal_var': 'E_s_sum'
        },
        'snowmelt': {
            'units': 'kg m-2 (equivalent to mm of water )',
            'description': 'Total snowmelt per unit area occurring within the snowpack',  # noqa
            'ipysnobal_var': 'melt_sum'
        },
        'surface_water_input': {
            'units': 'kg m-2 (equivalent to mm of water)',
            'description': 'Surface water input is liquid water output from bottom of snowpack or rain on bare ground per unit area',  # noqa
            'ipysnobal_var': 'ro_pred_sum'
        },
        'cold_content': {
            'units': 'J m-2',
            'description': 'Snowcover cold content',
            'ipysnobal_var': 'cc_s'
        },
        'thickness': {
            'units': 'm',
            'description': 'Thickness of the snowcover',
            'ipysnobal_var': 'z_s'
        },
        'snow_density': {
            'units': 'kg m-3',
            'description': 'Average snow density of the snowcover',
            'ipysnobal_var': 'rho'
        },
        'specific_mass': {
            'units': 'kg m-2 (equivalent to mm of water)',
            'description': 'Specific mass per unit area of the snowcover or snow water equivalent',  # noqa
            'ipysnobal_var': 'm_s'
        },
        'liquid_water': {
            'units': 'kg m-2 (equivalent to mm of water)',
            'description': 'Mass per unit area of liquid water in the snowcover',  # noqa
            'ipysnobal_var': 'h2o'
        },
        'temperature_surface': {
            'units': 'C',
            'description': 'Temperature of the surface layer',
            'ipysnobal_var': 'T_s_0'
        },
        'temperature_lower': {
            'units': 'C',
            'description': 'Temperature of the lower layer',
            'ipysnobal_var': 'T_s_l'
        },
        'temperature_snowcover': {
            'units': 'C',
            'description': 'Temperature of the snowcover',
            'ipysnobal_var': 'T_s'
        },
        'thickness_lower': {
            'units': 'm',
            'description': 'Thickness of the lower layer',
            'ipysnobal_var': 'z_s_l'
        },
        'water_saturation': {
            'units': 'percent',
            'description': 'Percentage of liquid water saturation of the snowcover',  # noqa
            'ipysnobal_var': 'h2o_sat'
        }
    }

    def __init__(self, output_file_name, output_path, myawsm):

        self._logger = logging.getLogger(__name__)

        self.output_file_name = output_file_name + '.nc'
        self.output_path = output_path
        self.output_filename = os.path.join(
            self.output_path, self.output_file_name)

        self.start_date = myawsm.start_date
        self.awsm = myawsm

        self.output_variables = self.awsm.pysnobal_output_vars
        self.precision = self.awsm.config['awsm system']['netcdf_output_precision'][0]  # noqa

        self._logger.info('PysnobalIO initialized')

    def create_output_files(self):
        """
        Create the ipysnobal output netCDF file
        """

        self._logger.info('Creating output iPysnobal file at {}'.format(
            self.output_file_name))

        fmt = '%Y-%m-%d %H:%M:%S'
        if os.path.isfile(self.output_filename):
            self._logger.warning(
                'Opening {}, data may be overwritten!'.format(
                    self.output_filename))
            em = nc.Dataset(self.output_filename, 'a')
            h = '[{}] Data added or updated'.format(
                datetime.now().strftime(fmt))
            setattr(em, 'last_modified', h)

            if 'projection' not in em.variables.keys():
                em = add_proj(em, None, self.awsm.topo.topoConfig['filename'])

        else:
            em = nc.Dataset(self.output_filename, 'w')

            dimensions = ('time', 'y', 'x')

            # create the dimensions
            em.createDimension('time', None)
            em.createDimension('y', len(self.awsm.topo.y))
            em.createDimension('x', len(self.awsm.topo.x))

            # create some variables
            # TODO what is the cell references, LL or center? #41
            em.createVariable('time', 'f', dimensions[0])
            em.createVariable('y', 'f', dimensions[1])
            em.createVariable('x', 'f', dimensions[2])

            setattr(em.variables['time'], 'units',
                    'hours since %s' % self.start_date.tz_localize(None))
            setattr(em.variables['time'], 'time_zone',
                    str(self.awsm.tzinfo).lower())
            setattr(em.variables['time'], 'calendar', 'standard')

            em.variables['x'][:] = self.awsm.topo.x
            em.variables['y'][:] = self.awsm.topo.y

            for var_name in self.output_variables:

                em.createVariable(
                    var_name,
                    self.precision,
                    dimensions[:3],
                    chunksizes=(6, 10, 10))
                setattr(em.variables[var_name],
                        'units',
                        self.OUTPUT_VARIABLES[var_name]['units'])
                setattr(em.variables[var_name],
                        'description',
                        self.OUTPUT_VARIABLES[var_name]['description'])

            # add projection info
            em = add_proj(em, None, self.awsm.topo.topoConfig['filename'])

        self.output_file = em

    def output_timestep(self, smrf_data, tstep):
        """
        Output the model results for the current time step

        Args:
            s:       dictionary of output variable numpy arrays
            tstep:   datetime time step

        """

        # preallocate
        output = {}

        # gather all the data together
        for key, att in self.OUTPUT_VARIABLES.items():
            output[key] = copy(smrf_data[att['ipysnobal_var']])

        # convert from K to C
        output['temperature_snowcover'] -= FREEZE
        output['temperature_surface'] -= FREEZE
        output['temperature_lower'] -= FREEZE

        # now find the correct index
        # the current time integer
        times = self.output_file.variables['time']

        # offset to match same convention as iSnobal
        tstep -= pd.to_timedelta(1, unit='h')
        t = nc.date2num(tstep.replace(tzinfo=None),
                        times.units,
                        times.calendar)

        try:
            index = nc.date2index(tstep, times, select='exact')
        except Exception:
            index = len(times)

        # insert the time
        self.output_file.variables['time'][index] = t

        # insert the data
        for key in self.output_variables:
            self.output_file.variables[key][index, :] = output[key]

        # sync to disk
        self.output_file.sync()