rtc-tools 2.5.2rc4__py3-none-any.whl → 2.6.0__py3-none-any.whl

rtctools/data/netcdf.py

@@ -17,7 +17,6 @@ import numpy as np
 
 
 class Stations:
-
     def __init__(self, dataset: Dataset, station_variable: Variable):
         self.__station_variable = station_variable
 
@@ -43,14 +42,17 @@ class Stations:
     @property
     def station_ids(self) -> Iterable:
         """
-        :return: An ordered iterable of the station ids (location ids) for which station data is available.
+        :return: An ordered iterable of the station ids (location ids) for which
+            station data is available.
+
         """
         return self.__attributes.keys()
 
     @property
     def attributes(self) -> OrderedDict:
         """
-        :return: An OrderedDict containing dicts containing the values for all station attributes of the input dataset.
+        :return: An OrderedDict containing dicts containing the values for all
+            station attributes of the input dataset.
         """
         return self.__attributes
 
@@ -77,26 +79,28 @@ class ImportDataset:
         self.__ensemble_size = 1
 
         # Load the content of a NetCDF file into a Dataset.
-        self.__filename = os.path.join(
-            folder,
-            basename + ".nc"
-        )
+        self.__filename = os.path.join(folder, basename + ".nc")
         self.__dataset = Dataset(self.__filename)
 
         # Find the number of ensemble members and the time and station id variables
         self.__time_variable = self.__find_time_variable()
         if self.__time_variable is None:
-            raise Exception('No time variable found in file ' + self.__filename + '. '
-                            'Please ensure the file contains a time variable with standard_name "time" and axis "T".')
+            raise Exception(
+                "No time variable found in file " + self.__filename + ". "
+                "Please ensure the file contains a time variable with standard_name "
+                '"time" and axis "T".'
+            )
 
         self.__ensemble_member_variable = self.__find_ensemble_member_variable()
         if self.__ensemble_member_variable:
-            self.__ensemble_size = self.__dataset.dimensions['realization'].size
+            self.__ensemble_size = self.__dataset.dimensions["realization"].size
 
         self.__station_variable = self.__find_station_variable()
         if self.__station_variable is None:
-            raise Exception('No station variable found in file ' + self.__filename + '. '
-                            'Please ensure the file contains a variable with cf_role "timeseries_id".')
+            raise Exception(
+                "No station variable found in file " + self.__filename + ". "
+                'Please ensure the file contains a variable with cf_role "timeseries_id".'
+            )
 
     def __str__(self):
         return self.__filename
@@ -109,21 +113,26 @@ class ImportDataset:
         :return: a netCDF4.Variable object of the time variable (or None if none found)
         """
         for variable in self.__dataset.variables.values():
-            if ('standard_name' in variable.ncattrs() and 'axis' in variable.ncattrs()
-                    and variable.standard_name == 'time' and variable.axis == 'T'):
+            if (
+                "standard_name" in variable.ncattrs()
+                and "axis" in variable.ncattrs()
+                and variable.standard_name == "time"
+                and variable.axis == "T"
+            ):
                 return variable
 
         return None
 
     def __find_ensemble_member_variable(self) -> Union[Variable, None]:
         """
-                Find the variable containing the ensemble member index in the given Dataset.
+        Find the variable containing the ensemble member index in the given Dataset.
 
-                :param dataset: The Dataset to be searched.
-                :return: a netCDF4.Variable object of the ensemble member index variable (or None if none found)
+        :param dataset: The Dataset to be searched.
+        :return: a netCDF4.Variable object of the ensemble member index variable (or None
+                 if none found)
         """
         for variable in self.__dataset.variables.values():
-            if 'standard_name' in variable.ncattrs() and variable.standard_name == "realization":
+            if "standard_name" in variable.ncattrs() and variable.standard_name == "realization":
                 return variable
 
         return None
@@ -136,7 +145,7 @@ class ImportDataset:
         :return: a netCDF4.Variable object of the station id variable (or None if none found)
         """
         for variable in self.__dataset.variables.values():
-            if 'cf_role' in variable.ncattrs() and variable.cf_role == 'timeseries_id':
+            if "cf_role" in variable.ncattrs() and variable.cf_role == "timeseries_id":
                 return variable
 
         return None
@@ -153,7 +162,7 @@ class ImportDataset:
         try:
             time_calendar = self.__time_variable.calendar
         except AttributeError:
-            time_calendar = u'gregorian'
+            time_calendar = "gregorian"
 
         return num2date(time_values, units=time_unit, calendar=time_calendar)
 
@@ -162,9 +171,9 @@ class ImportDataset:
 
     def find_timeseries_variables(self) -> List[str]:
         """
-        Find the keys of all 2-D or 3-D variables with dimensions {station, time} or {station, time, realization}
-        where station is the dimension of the station_variable, time the dimension of the time_variable and
-        realization the dimension for ensemble_member_index.
+        Find the keys of all 2-D or 3-D variables with dimensions {station, time} or {station, time,
+        realization} where station is the dimension of the station_variable, time the dimension of
+        the time_variable and realization the dimension for ensemble_member_index.
 
         :param dataset:           The Dataset to be searched.
         :param station_variable:  The station id variable.
@@ -176,10 +185,13 @@ class ImportDataset:
         if self.__ensemble_member_variable is not None:
             ensemble_dim = self.__ensemble_member_variable.dimensions[0]
             expected_dims = [
-                (time_dim, station_dim, ensemble_dim), (time_dim, ensemble_dim, station_dim),
-                (station_dim, time_dim, ensemble_dim), (station_dim, ensemble_dim, time_dim),
-                (ensemble_dim, time_dim, station_dim), (ensemble_dim, station_dim, time_dim)] \
-                + [(station_dim, time_dim), (time_dim, station_dim)]
+                (time_dim, station_dim, ensemble_dim),
+                (time_dim, ensemble_dim, station_dim),
+                (station_dim, time_dim, ensemble_dim),
+                (station_dim, ensemble_dim, time_dim),
+                (ensemble_dim, time_dim, station_dim),
+                (ensemble_dim, station_dim, time_dim),
+            ] + [(station_dim, time_dim), (time_dim, station_dim)]
         else:
             expected_dims = [(station_dim, time_dim), (time_dim, station_dim)]
 
@@ -190,7 +202,9 @@ class ImportDataset:
 
         return timeseries_variables
 
-    def read_timeseries_values(self, station_index: int, variable_name: str, ensemble_member: int = 0) -> np.ndarray:
+    def read_timeseries_values(
+        self, station_index: int, variable_name: str, ensemble_member: int = 0
+    ) -> np.ndarray:
         """
         Reads the specified timeseries from the input file.
 
@@ -206,7 +220,10 @@ class ImportDataset:
         # assert set(timeseries_variable.dimensions)==set(('time', 'station')) \
         #        or set(timeseries_variable.dimensions)==set(('time', 'station', 'realization'))
 
-        if self.__ensemble_member_variable is not None and 'realization' in timeseries_variable.dimensions:
+        if (
+            self.__ensemble_member_variable is not None
+            and "realization" in timeseries_variable.dimensions
+        ):
             ensemble_member_dim = self.__ensemble_member_variable.dimensions[0]
             for i in range(3):
                 if timeseries_variable.dimensions[i] == station_dim:
@@ -215,7 +232,7 @@ class ImportDataset:
                     ensemble_arg_Index = i
             time_arg_Index = set(range(3)) - {station_arg_Index, ensemble_arg_Index}
             time_arg_Index = time_arg_Index.pop()
-            argument = [None]*3
+            argument = [None] * 3
             argument[station_arg_Index] = station_index
             argument[ensemble_arg_Index] = ensemble_member
             argument[time_arg_Index] = slice(None)
@@ -255,11 +272,11 @@ class ImportDataset:
 
 class ExportDataset:
     """
-    A class used to write data to a NetCDF file.
-    Creates a new file or overwrites an old file. The file metadata will be written upon initialization. Data such
-    as times, station data and timeseries data should be presented to the ExportDataset through the various methods.
-    When all data has been written, the close method must be called to flush the changes from local memory to the
-    actual file on disk.
+    A class used to write data to a NetCDF file. Creates a new file or overwrites an old file. The
+    file metadata will be written upon initialization. Data such as times, station data and
+    timeseries data should be presented to the ExportDataset through the various methods. When all
+    data has been written, the close method must be called to flush the changes from local memory
+    to the actual file on disk.
     """
 
     def __init__(self, folder: str, basename: str):
@@ -268,22 +285,20 @@ class ExportDataset:
         :param basename: Basename of the file, extension ".nc" will be appended to this
         """
         # Create the file and open a Dataset to access it
-        self.__filename = os.path.join(
-            folder,
-            basename + ".nc"
-        )
+        self.__filename = os.path.join(folder, basename + ".nc")
         # use same write format as FEWS
-        self.__dataset = Dataset(self.__filename, mode='w', format='NETCDF3_CLASSIC')
+        self.__dataset = Dataset(self.__filename, mode="w", format="NETCDF3_CLASSIC")
 
         # write metadata to the file
-        self.__dataset.title = 'RTC-Tools Output Data'
-        self.__dataset.institution = 'Deltares'
-        self.__dataset.source = 'RTC-Tools'
-        self.__dataset.history = 'Generated on {}'.format(datetime.now())
-        self.__dataset.Conventions = 'CF-1.6'
-        self.__dataset.featureType = 'timeseries'
-
-        # dimensions are created when writing times and station data, must be created before writing variables
+        self.__dataset.title = "RTC-Tools Output Data"
+        self.__dataset.institution = "Deltares"
+        self.__dataset.source = "RTC-Tools"
+        self.__dataset.history = "Generated on {}".format(datetime.now())
+        self.__dataset.Conventions = "CF-1.6"
+        self.__dataset.featureType = "timeseries"
+
+        # dimensions are created when writing times and station data, must be created before
+        # writing variables
         self.__time_dim = None
         self.__station_dim = None
         self.__station_id_to_index_mapping = None
@@ -298,35 +313,41 @@ class ExportDataset:
         """
         Writes a time variable to the given dataset.
 
-        :param dataset:        The NetCDF4.Dataset object that the times will be written to (must have write permission)
+        :param dataset:        The NetCDF4.Dataset object that the times will be written to
+            (must have write permission)
         :param times:          The times that are to be written in seconds.
         :param forecast_time:  The forecast time in seconds corresponding to the forecast date
-        :param forecast_date:  The datetime corresponding with time in seconds at the forecast index.
+        :param forecast_date:  The datetime corresponding with time in seconds at the forecast
+                               index.
         """
 
         # in a NetCDF file times are written with respect to a reference date
-        # the written values for the times may never be negative, so use the earliest time as the reference date
+        # the written values for the times may never be negative, so use the earliest time as the
+        # reference date
         reference_date = forecast_date
         minimum_time = np.min(times)
         if minimum_time < 0:
             times = times - minimum_time
             reference_date = reference_date - timedelta(seconds=forecast_time - minimum_time)
 
-        self.__time_dim = self.__dataset.createDimension('time', None)
+        self.__time_dim = self.__dataset.createDimension("time", None)
 
-        time_var = self.__dataset.createVariable('time', 'f8', ('time',))
-        time_var.standard_name = 'time'
-        time_var.units = 'seconds since {}'.format(reference_date)
-        time_var.axis = 'T'
+        time_var = self.__dataset.createVariable("time", "f8", ("time",))
+        time_var.standard_name = "time"
+        time_var.units = "seconds since {}".format(reference_date)
+        time_var.axis = "T"
         time_var[:] = times
 
     def write_ensemble_data(self, ensemble_size):
-
         if ensemble_size > 1:
-            self.__ensemble_member_dim = self.__dataset.createDimension('realization', ensemble_size)
-            ensemble_member_var = self.__dataset.createVariable('realization', 'i', ('realization',))
-            ensemble_member_var.standard_name = 'realization'
-            ensemble_member_var.long_name = 'Index of an ensemble member within an ensemble'
+            self.__ensemble_member_dim = self.__dataset.createDimension(
+                "realization", ensemble_size
+            )
+            ensemble_member_var = self.__dataset.createVariable(
+                "realization", "i", ("realization",)
+            )
+            ensemble_member_var.standard_name = "realization"
+            ensemble_member_var.long_name = "Index of an ensemble member within an ensemble"
             ensemble_member_var.units = 1
 
     def write_station_data(self, stations: Stations, output_station_ids: List[str]) -> None:
@@ -334,20 +355,24 @@ class ExportDataset:
         Writes the station ids and additional station information to the given dataset.
 
         :param stations:           The stations data read from the input file.
-        :param output_station_ids: The set of station ids for which output will be written. Must be unique.
+        :param output_station_ids: The set of station ids for which output will be written. Must be
+            unique.
         """
         assert len(set(output_station_ids)) == len(output_station_ids)
 
-        self.__station_dim = self.__dataset.createDimension('station', len(output_station_ids))
+        self.__station_dim = self.__dataset.createDimension("station", len(output_station_ids))
 
         # first write the ids
         max_id_length = max(len(id) for id in output_station_ids)
-        self.__dataset.createDimension('char_leng_id', max_id_length)
-        station_id_var = self.__dataset.createVariable('station_id', 'c', ('station', 'char_leng_id'))
-        station_id_var.long_name = 'station identification code'
-        station_id_var.cf_role = 'timeseries_id'
+        self.__dataset.createDimension("char_leng_id", max_id_length)
+        station_id_var = self.__dataset.createVariable(
+            "station_id", "c", ("station", "char_leng_id")
+        )
+        station_id_var.long_name = "station identification code"
+        station_id_var.cf_role = "timeseries_id"
 
-        # we must store the index we use for each station id, to be able to write the data at the correct index later
+        # we must store the index we use for each station id, to be able to write the data at the
+        # correct index later
         self.__station_id_to_index_mapping = {}
         for i, id in enumerate(output_station_ids):
             station_id_var[i, :] = list(id)
@@ -355,7 +380,7 @@ class ExportDataset:
 
         # now write the stored attributes
         for var_name, attr_var in stations.attribute_variables.items():
-            variable = self.__dataset.createVariable(var_name, attr_var.datatype, ('station',))
+            variable = self.__dataset.createVariable(var_name, attr_var.datatype, ("station",))
             # copy all attributes from the original input variable
             variable.setncatts(attr_var.__dict__)
 
@@ -367,50 +392,69 @@ class ExportDataset:
     def create_variables(self, variable_names: List[str], ensemble_size: int) -> None:
         """
         Creates variables in the dataset for each of the provided parameter ids.
-        The write_times and write_station_data methods must be called first, to ensure the necessary dimensions have
-        already been created in the output NetCDF file.
+        The write_times and write_station_data methods must be called first, to ensure the necessary
+        dimensions have already been created in the output NetCDF file.
 
-        :param variable_names: The parameter ids for which variables must be created. Must be unique.
+        :param variable_names: The parameter ids for which variables must be created. Must be
+            unique.
         :param ensemble_size: the number of members in the ensemble
         """
         assert len(set(variable_names)) == len(variable_names)
 
-        assert self.__time_dim is not None, 'First call write_times to ensure the time dimension has been created.'
-        assert self.__station_dim is not None, 'First call write_station_data to ensure ' \
-                                               'the station dimension has been created'
-        assert self.__station_id_to_index_mapping is not None  # should also be created in write_station_data
+        assert (
+            self.__time_dim is not None
+        ), "First call write_times to ensure the time dimension has been created."
+        assert (
+            self.__station_dim is not None
+        ), "First call write_station_data to ensure the station dimension has been created"
+        assert (
+            self.__station_id_to_index_mapping is not None
+        )  # should also be created in write_station_data
 
         if ensemble_size > 1:
-            assert self.__ensemble_member_dim is not None, \
-                'First call write_ensemble_data to ensure the realization dimension has been created.'
+            assert (
+                self.__ensemble_member_dim is not None
+            ), "First call write_ensemble_data to ensure the realization dimension has been created"
 
             for variable_name in variable_names:
-                self.__dataset.createVariable(variable_name,
-                                              'f8', ('time', 'station', 'realization'), fill_value=np.nan)
+                self.__dataset.createVariable(
+                    variable_name, "f8", ("time", "station", "realization"), fill_value=np.nan
+                )
         else:
             for variable_name in variable_names:
-                self.__dataset.createVariable(variable_name,
-                                              'f8', ('time', 'station'), fill_value=np.nan)
-
-    def write_output_values(self, station_id: str, variable_name: str,
-                            ensemble_member_index: int, values: np.ndarray,
-                            ensemble_size: int) -> None:
-        """
-        Writes the given data to the dataset. The variable must have already been created through the
-        create_variables method. After all calls to write_output_values, the close method must be called to flush all
-        changes.
+                self.__dataset.createVariable(
+                    variable_name, "f8", ("time", "station"), fill_value=np.nan
+                )
+
+    def write_output_values(
+        self,
+        station_id: str,
+        variable_name: str,
+        ensemble_member_index: int,
+        values: np.ndarray,
+        ensemble_size: int,
+    ) -> None:
+        """
+        Writes the given data to the dataset. The variable must have already been created through
+        the create_variables method. After all calls to write_output_values, the close method must
+        be called to flush all changes.
 
         :param station_id: The id of the station the data is written for.
-        :param variable_name: The name of the variable the data is written to (must have already been created).
+        :param variable_name: The name of the variable the data is written to (must have already
+            been created).
         :param ensemble_member_index: The index associated to the ensemble member
         :param values:        The values that are to be written to the file
         :param ensemble_size: the number of members in the ensemble
         """
-        assert self.__station_id_to_index_mapping is not None, 'First call write_station_data and create_variables.'
+        assert (
+            self.__station_id_to_index_mapping is not None
+        ), "First call write_station_data and create_variables."
 
         station_index = self.__station_id_to_index_mapping[station_id]
         if ensemble_size > 1:
-            self.__dataset.variables[variable_name][:, station_index, ensemble_member_index] = values
+            self.__dataset.variables[variable_name][
+                :, station_index, ensemble_member_index
+            ] = values
         else:
             self.__dataset.variables[variable_name][:, station_index] = values