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.make_xarray_diag
   credit.output.save_netcdf_increment
   credit.output.save_netcdf_clean
   credit.output.save_netcdf_diag


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

.. py:data:: logger

.. py:function:: load_metadata(conf)

   Load metadata attributes from yaml file in credit/metadata directory.

   If the configured path is a bare filename (no directory separators), it is
   resolved against the installed ``credit.metadata`` package directory so
   users don't need to hard-code absolute paths.

   By default the function will assume era5.yaml as the default metadata file.


.. 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
   :type conf: dict
   :param and variables.:

   :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:: make_xarray_diag(pred, forecast_datetime, lat, lon, conf)

.. 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


.. py:function:: save_netcdf_clean(darray_upper_air, darray_single_level, nc_filename, forecast_hour, meta_data, conf, use_logger=True)

   Save forecast data (upper-air and optionally single-level variables) to a NetCDF file.

   This function is similar to `save_netcdf_increment` but cleaned-up all the interpolations

   :param darray_upper_air: Upper-air forecast data with dimension "vars".
   :type darray_upper_air: xr.DataArray
   :param darray_single_level: Single-level forecast data with dimension "vars". If None, only upper-air
                               variables are included.
   :type darray_single_level: xr.DataArray or None
   :param nc_filename: Base filename (used as subdirectory and prefix for the saved file).
   :type nc_filename: str
   :param forecast_hour: Forecast lead time in hours.
   :type forecast_hour: int
   :param meta_data: Dictionary containing variable metadata attributes.
                     If `False`, no metadata is applied and default time encoding is used.
   :type meta_data: dict or bool
   :param conf: Configuration dictionary. Must include at least:
                - conf["predict"]["save_forecast"]: directory where forecasts are saved.
                - conf["predict"]["save_vars"]: (optional) list of variable names to keep.
   :type conf: dict
   :param use_logger: If True (default), configure logging and print progress messages.
   :type use_logger: bool, optional

   .. rubric:: Notes

   - Files are saved into:
       ``{conf['predict']['save_forecast']}/{nc_filename}/pred_{nc_filename}_{forecast_hour:03d}.nc``

   - If `meta_data` is provided, variable attributes are updated accordingly.
     Otherwise, a default gregorian calendar encoding is applied to the "time" variable.

   - If `conf['predict']['save_vars']` is non-empty, variables not listed there
     are dropped before saving.

   :returns: The function saves a NetCDF file to disk and does not return a value.
   :rtype: None


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