shrecc.treatment
================

.. py:module:: shrecc.treatment


Functions
---------

.. autoapisummary::

   shrecc.treatment.add_missing_elements
   shrecc.treatment.calculate_Z_cons
   shrecc.treatment.calculate_results
   shrecc.treatment.concatenate_results
   shrecc.treatment.create_block_diagonal_matrix
   shrecc.treatment.data_processing
   shrecc.treatment.load_from_pickle
   shrecc.treatment.process_matrix
   shrecc.treatment.process_results_light
   shrecc.treatment.save_to_pickle
   shrecc.treatment.treating_data


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

.. py:function:: add_missing_elements(df, existing_elements, all_elements, axis=0, fill_value=0)

   Adds missing elements (rows or columns) to a dataframe.

   :param df: The dataframe to which elements will be added.
   :type df: pd.DataFrame
   :param existing_elements: The set of existing elements in the dataframe.
   :type existing_elements: set
   :param all_elements: The set of all possible elements.
   :type all_elements: set
   :param axis: The axis along which to add missing elements (0 for index, 1 for columns).
   :type axis: int
   :param fill_value: The value to fill for the missing elements.

   :returns: The dataframe with missing elements added.
   :rtype: pd.DataFrame


.. py:function:: calculate_Z_cons(filename, L_series, output, Z_indices)

   Calculates the consumption matrix Z_cons.

   :param year: The selected year.
   :type year: int
   :param L_series: The L series DataFrame.
   :type L_series: pd.DataFrame
   :param output: The output series.
   :type output: pd.Series
   :param Z_indices: The Z indices dictionary.
   :type Z_indices: dict

   :returns: The consumption matrix Z_cons.
   :rtype: pd.DataFrame


.. py:function:: calculate_results(year, n_c, n_p, t_index, Z_net, Z_load, data_dir)

   Calculates results by inverting matrices over time.

   :param year: The selected year.
   :type year: int
   :param n_c: The number of countries.
   :type n_c: int
   :param n_p: The number of production mix elements.
   :type n_p: int
   :param t_index: The time index.
   :type t_index: pd.Index
   :param Z_net: The net consumption dataframe.
   :type Z_net: pd.DataFrame
   :param Z_load: The load dataframe.
   :type Z_load: pd.DataFrame
   :param data_dir: location of the data.
   :type data_dir: Path

   :returns: The results dictionary.
   :rtype: dict


.. py:function:: concatenate_results(results, Z)

   Concatenates results into a single DataFrame.

   :param results_light: The light results dictionary.
   :type results_light: dict
   :param Z: The network dataframe.
   :type Z: pd.DataFrame

   :returns: The concatenated DataFrame.
   :rtype: pd.DataFrame


.. py:function:: create_block_diagonal_matrix(data_list)

   Creates a block diagonal matrix from a list of data arrays.

   :param data_list: A list of data arrays to be combined into a block diagonal matrix.
   :type data_list: list

   :returns: A block diagonal matrix containing the input data.
   :rtype: pd.DataFrame


.. py:function:: data_processing(data_df, year, path_to_data=None)

   Processes data, adds missing countries, and correctly divides them between consumption and demand.

   :param data_df: All of the downloaded data for the selected year (from `get_data`).
   :type data_df: pd.DataFrame
   :param year: The selected year, e.g., 2023.
   :type year: int
   :param path_to_data: location of the data.
   :type path_to_data: str or Path

   :returns: Doesn't return anything, but saves files to the directory 'data'.
   :rtype: None


.. py:function:: load_from_pickle(filename)

   Loads an object from a pickle file.

   :param filename: Path to the filename where the object will be saved.
   :type filename: Path

   :returns: The object loaded from the pickle file.
   :rtype: object


.. py:function:: process_matrix(df, operation, axis=1, **kwargs)

   Applies a specified operation to a dataframe or matrix.

   :param df: The dataframe or matrix to process.
   :type df: pd.DataFrame
   :param operation: The operation to perform (e.g., 'normalize', 'reorder_levels').
   :type operation: str
   :param axis: The axis along which the operation should be applied, if applicable.
   :type axis: int
   :param kwargs: Additional arguments for specific operations.

   :returns: The processed dataframe or matrix.
   :rtype: pd.DataFrame


.. py:function:: process_results_light(results, filename, n_c)

   Processes the results to a lighter format.

   :param results: The results dictionary.
   :type results: dict
   :param filename: Path to the filename where the object will be saved.
   :type filename: Path
   :param n_c: The number of countries.
   :type n_c: int

   :returns: The light results dictionary.
   :rtype: dict


.. py:function:: save_to_pickle(obj, filename)

   Saves an object to a pickle file.

   :param obj: The object to be saved.
   :param filename: Path to the filename where the object will be saved.
   :type filename: Path

   :returns: None


.. py:function:: treating_data(year, n_c, n_p, t_index, Z_net, data_dir)

   Function called from `data_processing`. This includes heavy operations, so it's advised to run this part on a server.

   :param year: The selected year, passed from `data_processing`.
   :type year: int
   :param n_c: The number of all countries (including missing ones) in the dataframe, passed from `data_processing`.
   :type n_c: int
   :param n_p: The number of production mix elements, passed from `data_processing`.
   :type n_p: int
   :param t_index: The time index, passed from `data_processing`.
   :type t_index: pd.Index
   :param Z_net: The net consumption dataframe, passed from `data_processing`.
   :type Z_net: pd.DataFrame
   :param data_dir: location of the data.
   :type data_dir: Path

   :returns: Doesn't return anything; everything gets saved to the directory 'data'.
   :rtype: None