Features/save flex bands (#326)

* Add flexibility_bands as property

* Add flexibility_bands in csv export

* Minor doc changes

* Change definition of flex band csv files

* Read flex bands from files

* Add tests

* Add option to choose which electromobility attributes to save

* Fix link to function

edisgo/edisgo.py

@@ -2070,6 +2070,12 @@ def save(
             To only store certain results provide a dictionary. See function docstring
             `parameters` parameter in :func:`~.network.results.Results.to_csv`
             for more information.
+        electromobility_attributes : None or list(str)
+            Specifies which electromobility attributes to store. By default this is set
+            to None, in which case all attributes are stored.
+            See function docstring `attributes` parameter in
+            :attr:`~.network.electromobility.Electromobility.to_csv` for more
+            information.
         archive : bool, optional
             Save disk storage capacity by archiving the csv files. The
             archiving takes place after the generation of the CSVs and
@@ -2102,7 +2108,10 @@ def save(
             )
 
         if save_electromobility:
-            self.electromobility.to_csv(os.path.join(directory, "electromobility"))
+            self.electromobility.to_csv(
+                os.path.join(directory, "electromobility"),
+                attributes=kwargs.get("electromobility_attributes", None),
+            )
 
         # save configs
         self.config.to_json(directory)

edisgo/network/electromobility.py

@@ -319,11 +319,47 @@ def eta_charging_points(self):
         except Exception:
             return None
 
+    @property
+    def flexibility_bands(self):
+        """
+        Dictionary with flexibility bands (lower and upper energy band as well as
+        upper power band).
+
+        Parameters
+        -----------
+        flex_dict : dict(str, :pandas:`pandas.DataFrame<DataFrame>`)
+            Keys are 'upper_power', 'lower_energy' and 'upper_energy'.
+            Values are dataframes containing the corresponding band per each charging
+            point. Columns of the dataframe are the charging point names as in
+            :attr:`~.network.topology.Topology.loads_df`. Index is a time index.
+
+        Returns
+        -------
+        dict(str, :pandas:`pandas.DataFrame<DataFrame>`)
+            See input parameter `flex_dict` for more information on the dictionary.
+
+        """
+        try:
+            return self._flexibility_bands
+        except Exception:
+            return {
+                "upper_power": pd.DataFrame(),
+                "lower_energy": pd.DataFrame(),
+                "upper_energy": pd.DataFrame(),
+            }
+
+    @flexibility_bands.setter
+    def flexibility_bands(self, flex_dict):
+        self._flexibility_bands = flex_dict
+
     def get_flexibility_bands(self, edisgo_obj, use_case):
         """
         Method to determine flexibility bands (lower and upper energy band as well as
         upper power band).
 
+        Besides being returned by this function, flexibility bands are written to
+        :attr:`flexibility_bands`.
+
         Parameters
         -----------
         edisgo_obj : :class:`~.EDisGo`
@@ -442,17 +478,20 @@ def _shorten_and_set_index(band):
         upper_power = _shorten_and_set_index(upper_power)
         lower_energy = _shorten_and_set_index(lower_energy)
         upper_energy = _shorten_and_set_index(upper_energy)
-        return {
+
+        flex_band_dict = {
             "upper_power": upper_power,
             "lower_energy": lower_energy,
             "upper_energy": upper_energy,
         }
+        self.flexibility_bands = flex_band_dict
+        return flex_band_dict
 
-    def to_csv(self, directory):
+    def to_csv(self, directory, attributes=None):
         """
-        Exports electromobility to csv files.
+        Exports electromobility data to csv files.
 
-        The following attributes are exported:
+        The following attributes can be exported:
 
         * 'charging_processes_df' : Attribute :py:attr:`~charging_processes_df`
           is saved to `charging_processes.csv`.
@@ -464,23 +503,39 @@ def to_csv(self, directory):
           `integrated_charging_parks.csv`.
         * 'simbev_config_df' : Attribute :py:attr:`~simbev_config_df` is
           saved to `simbev_config.csv`.
+        * 'flexibility_bands' : The three flexibility bands in attribute
+          :py:attr:`~flexibility_bands` are saved to
+          `flexibility_band_upper_power.csv`, `flexibility_band_lower_energy.csv`, and
+          `flexibility_band_upper_energy.csv`.
 
         Parameters
         ----------
         directory : str
-            Path to save electromobility to.
+            Path to save electromobility data to.
+        attributes : list(str) or None
+            List of attributes to export. See above for attributes that can be exported.
+            If None, all specified attributes are exported. Default: None.
 
         """
         os.makedirs(directory, exist_ok=True)
 
-        attrs = _get_matching_dict_of_attributes_and_file_names()
+        attrs_file_names = _get_matching_dict_of_attributes_and_file_names()
 
-        for attr, file in attrs.items():
-            df = getattr(self, attr)
+        if attributes is None:
+            attributes = list(attrs_file_names.keys())
 
-            if not df.empty:
-                path = os.path.join(directory, file)
-                df.to_csv(path)
+        for attr in attributes:
+            file = attrs_file_names[attr]
+            df = getattr(self, attr)
+            if attr == "flexibility_bands":
+                for band in file.keys():
+                    if band in df.keys() and not df[band].empty:
+                        path = os.path.join(directory, file[band])
+                        df[band].to_csv(path)
+            else:
+                if not df.empty:
+                    path = os.path.join(directory, file)
+                    df.to_csv(path)
 
     def from_csv(self, data_path, edisgo_obj, from_zip_archive=False):
         """
@@ -506,23 +561,47 @@ def from_csv(self, data_path, edisgo_obj, from_zip_archive=False):
             files = zip.namelist()
 
             # add directory and .csv to files to match zip archive
-            attrs = {k: f"electromobility/{v}" for k, v in attrs.items()}
+            attrs = {
+                k: (
+                    f"electromobility/{v}"
+                    if isinstance(v, str)
+                    else {k2: f"electromobility/{v2}" for k2, v2 in v.items()}
+                )
+                for k, v in attrs.items()
+            }
 
         else:
             # read from directory
             # check files within the directory
             files = os.listdir(data_path)
 
-        attrs_to_read = {k: v for k, v in attrs.items() if v in files}
+        attrs_to_read = {
+            k: v
+            for k, v in attrs.items()
+            if (isinstance(v, str) and v in files)
+            or (isinstance(v, dict) and any([_ in files for _ in v.values()]))
+        }
 
         for attr, file in attrs_to_read.items():
-            if from_zip_archive:
-                # open zip file to make it readable for pandas
-                with zip.open(file) as f:
-                    df = pd.read_csv(f, index_col=0)
+            if attr == "flexibility_bands":
+                df = {}
+                for band, file_name in file.items():
+                    if file_name in files:
+                        if from_zip_archive:
+                            # open zip file to make it readable for pandas
+                            with zip.open(file_name) as f:
+                                df[band] = pd.read_csv(f, index_col=0, parse_dates=True)
+                        else:
+                            path = os.path.join(data_path, file_name)
+                            df[band] = pd.read_csv(path, index_col=0, parse_dates=True)
             else:
-                path = os.path.join(data_path, file)
-                df = pd.read_csv(path, index_col=0)
+                if from_zip_archive:
+                    # open zip file to make it readable for pandas
+                    with zip.open(file) as f:
+                        df = pd.read_csv(f, index_col=0)
+                else:
+                    path = os.path.join(data_path, file)
+                    df = pd.read_csv(path, index_col=0)
 
             if attr == "potential_charging_parks_gdf":
                 epsg = edisgo_obj.topology.grid_district["srid"]
@@ -635,7 +714,8 @@ def _get_matching_dict_of_attributes_and_file_names():
     restore and maps them to the file name.
 
     Is used in functions
-    :attr:`~.network.electromobility.Electromobility.from_csv`.
+    :attr:`~.network.electromobility.Electromobility.from_csv` and
+    attr:`~.network.electromobility.Electromobility.to_csv`.
 
     Returns
     -------
@@ -649,6 +729,11 @@ def _get_matching_dict_of_attributes_and_file_names():
         "potential_charging_parks_gdf": "potential_charging_parks.csv",
         "integrated_charging_parks_df": "integrated_charging_parks.csv",
         "simbev_config_df": "metadata_simbev_run.csv",
+        "flexibility_bands": {
+            "upper_power": "flexibility_band_upper_power.csv",
+            "lower_energy": "flexibility_band_lower_energy.csv",
+            "upper_energy": "flexibility_band_upper_energy.csv",
+        },
     }
 
     return emob_dict

tests/network/test_electromobility.py

@@ -5,6 +5,8 @@
 import pandas as pd
 import pytest
 
+from pandas.util.testing import assert_frame_equal
+
 from edisgo.edisgo import EDisGo
 from edisgo.io.electromobility_import import (
     import_electromobility,
@@ -61,11 +63,32 @@ def test_eta_charging_points(self):
     def test_to_csv(self):
         """Test for method to_csv."""
         dir = os.path.join(os.getcwd(), "electromobility")
+        timeindex = pd.date_range("1/1/1970", periods=2, freq="H")
+        flex_bands = {
+            "upper_energy": pd.DataFrame({"cp_1": [1, 2]}, index=timeindex),
+            "upper_power": pd.DataFrame({"cp_1": [1, 2]}, index=timeindex),
+            "lower_energy": pd.DataFrame({"cp_1": [1, 2]}, index=timeindex),
+        }
+        self.edisgo_obj.electromobility.flexibility_bands = flex_bands
+
+        # ############ test with default values #####################
         self.edisgo_obj.electromobility.to_csv(dir)
 
         saved_files = os.listdir(dir)
-        assert len(saved_files) == 3
+        assert len(saved_files) == 6
         assert "charging_processes.csv" in saved_files
+        assert "flexibility_band_upper_power.csv" in saved_files
+
+        shutil.rmtree(dir)
+
+        # ############ test specifying attributes #####################
+        self.edisgo_obj.electromobility.to_csv(
+            dir, attributes=["potential_charging_parks_gdf"]
+        )
+
+        saved_files = os.listdir(dir)
+        assert len(saved_files) == 1
+        assert "potential_charging_parks.csv" in saved_files
 
         shutil.rmtree(dir)
 
@@ -75,9 +98,16 @@ def test_from_csv(self):
 
         """
         dir = os.path.join(os.getcwd(), "electromobility")
+        timeindex = pd.date_range("1/1/1970", periods=2, freq="H")
+        flex_bands = {
+            "upper_energy": pd.DataFrame({"cp_1": [1, 2]}, index=timeindex),
+            "upper_power": pd.DataFrame({"cp_1": [1, 2]}, index=timeindex),
+            "lower_energy": pd.DataFrame({"cp_1": [1, 2]}, index=timeindex),
+        }
+        self.edisgo_obj.electromobility.flexibility_bands = flex_bands
         self.edisgo_obj.electromobility.to_csv(dir)
 
-        # reset self.topology
+        # reset Electromobility
         self.edisgo_obj.electromobility = Electromobility()
 
         self.edisgo_obj.electromobility.from_csv(dir, self.edisgo_obj)
@@ -86,4 +116,20 @@ def test_from_csv(self):
         assert len(self.edisgo_obj.electromobility.potential_charging_parks_gdf) == 1621
         assert self.edisgo_obj.electromobility.integrated_charging_parks_df.empty
 
+        assert_frame_equal(
+            self.edisgo_obj.electromobility.flexibility_bands["upper_energy"],
+            flex_bands["upper_energy"],
+            check_freq=False,
+        )
+        assert_frame_equal(
+            self.edisgo_obj.electromobility.flexibility_bands["lower_energy"],
+            flex_bands["lower_energy"],
+            check_freq=False,
+        )
+        assert_frame_equal(
+            self.edisgo_obj.electromobility.flexibility_bands["upper_power"],
+            flex_bands["upper_power"],
+            check_freq=False,
+        )
+
         shutil.rmtree(dir)