credit.output
=============

.. py:module:: credit.output

.. autoapi-nested-parse::

   output.py
   -------------------------------------------------------
   Content:
       - load_metadata()
       - make_xarray()
       - save_netcdf_increment()



Attributes
----------

.. autoapisummary::

   credit.output.logger


Functions
---------

.. autoapisummary::

   credit.output.load_metadata
   credit.output.split_and_reshape
   credit.output.make_xarray
   credit.output.save_netcdf_increment


Module Contents
---------------

.. py:data:: logger

.. py:function:: load_metadata(conf)

   Load metadata attributes from yaml file in credit/metadata directory


.. py:function:: split_and_reshape(tensor, conf)

   Split the output tensor of the model to upper air variables and diagnostics/surface variables.

   Upperair level arrangement: top-of-atmosphere--> near-surface --> single layer
   An example: U (top-of-atmosphere) --> U (near-surface) --> V (top-of-atmosphere) --> V (near-surface)
   The shape of the output tensor is (variables, latitude, longitude)

   :param tensor: PyTorch Tensor containing output of the AI NWP model
   :param conf: config file for the model


.. py:function:: make_xarray(pred, forecast_datetime, lat, lon, conf)

   Create two xarray.DataArray objects for upper air and surface variables.

   :param pred: Prediction tensor containing both upper air and surface variables.
   :type pred: torch.Tensor or np.ndarray
   :param forecast_datetime: The forecast initialization datetime.
   :type forecast_datetime: datetime
   :param lat: Latitude values.
   :type lat: np.ndarray or list
   :param lon: Longitude values.
   :type lon: np.ndarray or list
   :param conf: Configuration dictionary containing details about the data structure
                and variables.
   :type conf: dict

   :returns:

             DataArray containing upper air variables with dimensions
                 [time, vars, level, latitude, longitude].
             darray_single_level (xarray.DataArray): DataArray containing surface variables with dimensions
                 [time, vars, latitude, longitude].
   :rtype: darray_upper_air (xarray.DataArray)


.. py:function:: save_netcdf_increment(darray_upper_air: xarray.DataArray, darray_single_level: xarray.DataArray, nc_filename: str, forecast_hour: int, meta_data: dict, conf: dict)

   Save CREDIT model prediction output to netCDF file. Also performs pressure level
   interpolation on the output if you wish.

   :param darray_upper_air: upper air variable predictions
   :type darray_upper_air: xr.DataArray
   :param darray_single_level: surface variable predictions
   :type darray_single_level: xr.DataArray
   :param nc_filename: file description to go into output filenames
   :type nc_filename: str
   :param forecast_hour: how many hours since the initialization of the model.
   :type forecast_hour: int
   :param meta_data: metadata dictionary for output variables
   :type meta_data: dict
   :param conf: configuration dictionary for training and/or rollout
   :type conf: dict