|  | 
| 1 |  | -import numpy as np | 
|  | 1 | +from typing import Any, Optional | 
| 2 | 2 | from geos.hdf5_wrapper import wrapper as h5w | 
| 3 |  | -import matplotlib as mpl | 
| 4 | 3 | import matplotlib.pyplot as plt | 
| 5 | 4 | import os | 
| 6 |  | -import sys | 
| 7 | 5 | import argparse | 
| 8 | 6 | 
 | 
| 9 | 7 | import re | 
| 10 | 8 | 
 | 
| 11 | 9 | 
 | 
| 12 |  | -def isiterable( obj ): | 
|  | 10 | +def isiterable( obj: Any ) -> bool: | 
|  | 11 | +    """Check if input is iterable.""" | 
| 13 | 12 |     try: | 
| 14 |  | -        it = iter( obj ) | 
|  | 13 | +        it = iter( obj )  # noqa: F841 | 
| 15 | 14 |     except TypeError: | 
| 16 | 15 |         return False | 
| 17 | 16 |     return True | 
| 18 | 17 | 
 | 
| 19 | 18 | 
 | 
| 20 |  | -def getHistorySeries( database, variable, setname, indices=None, components=None ): | 
| 21 |  | -    """ | 
| 22 |  | -    Retrieve a series of time history structures suitable for plotting in addition to the specific set index and component for the time series | 
|  | 19 | +def getHistorySeries( database: h5w, | 
|  | 20 | +                      variable: str, | 
|  | 21 | +                      setname: str, | 
|  | 22 | +                      indices: Optional[ int | list[ int ] ] = None, | 
|  | 23 | +                      components: Optional[ int | list[ int ] ] = None ) -> Optional[ list[ tuple[ Any, ...] ] ]: | 
|  | 24 | +    """Retrieve a series of time history structures suitable for plotting in addition to the specific set index and component for the time series. | 
| 23 | 25 | 
 | 
| 24 | 26 |     Args: | 
| 25 | 27 |         database (geos.hdf5_wrapper.hdf5_wrapper): database to retrieve time history data from | 
| 26 | 28 |         variable (str): the name of the time history variable for which to retrieve time-series data | 
| 27 | 29 |         setname (str): the name of the index set as specified in the geosx input xml for which to query time-series data | 
| 28 |  | -        indices (int, list): the indices in the named set to query for, if None, defaults to all | 
| 29 |  | -        components (int, list): the components in the flattened data types to retrieve, defaults to all | 
| 30 |  | -     | 
|  | 30 | +        indices (Optional[int | list[ int ]]): the indices in the named set to query for, if None, defaults to all | 
|  | 31 | +        components (Optional[int | list[ int ]]): the components in the flattened data types to retrieve, defaults to all | 
|  | 32 | +
 | 
| 31 | 33 |     Returns: | 
| 32 |  | -        list: list of (time, data, idx, comp) timeseries tuples for each time history data component | 
|  | 34 | +        Optional[list[ tuple[ Any, ...] ]]: list of (time, data, idx, comp) timeseries tuples for each time history data component | 
| 33 | 35 |     """ | 
| 34 |  | - | 
| 35 | 36 |     set_regex = re.compile( variable + '(.*?)', re.IGNORECASE ) | 
| 36 | 37 |     if setname is not None: | 
| 37 |  | -        set_regex = re.compile( variable + '\s*' + str( setname ), re.IGNORECASE ) | 
|  | 38 | +        set_regex = re.compile( variable + r'\s*' + str( setname ), re.IGNORECASE ) | 
| 38 | 39 |     time_regex = re.compile( 'Time', re.IGNORECASE )  # need to make this per-set, thought that was in already? | 
| 39 | 40 | 
 | 
| 40 | 41 |     set_match = list( filter( set_regex.match, database.keys() ) ) | 
| @@ -62,39 +63,48 @@ def getHistorySeries( database, variable, setname, indices=None, components=None | 
| 62 | 63 |         print( | 
| 63 | 64 |             f"Error: The length of the time-series {time_match} and data-series {set_match} do not match: {time_series.shape} and {data_series.shape} !" | 
| 64 | 65 |         ) | 
| 65 |  | - | 
|  | 66 | +    indices1: list[ int ] = [] | 
| 66 | 67 |     if indices is not None: | 
| 67 | 68 |         if type( indices ) is int: | 
| 68 |  | -            indices = [ indices ] | 
| 69 |  | -        if isiterable( indices ): | 
| 70 |  | -            oob_idxs = list( filter( lambda idx: not 0 <= idx < data_series.shape[ 1 ], indices ) ) | 
|  | 69 | +            indices1 = [ indices ] | 
|  | 70 | +        elif isiterable( indices ): | 
|  | 71 | +            oob_idxs: list[ int ] = list( | 
|  | 72 | +                filter( | 
|  | 73 | +                    lambda idx: not 0 <= idx < data_series.shape[ 1 ],  # type: ignore[arg-type] | 
|  | 74 | +                    indices ) )  # type: ignore[arg-type] | 
| 71 | 75 |             if len( oob_idxs ) > 0: | 
| 72 |  | -                print( f"Error: The specified indices: ({', '.join(oob_idxs)}) " + "\n\t" + | 
| 73 |  | -                       f" are out of the dataset index range: [0,{data_series.shape[1]})" ) | 
| 74 |  | -            indices = list( set( indices ) - set( oob_idxs ) ) | 
|  | 76 | +                print( f"Error: The specified indices: ({', '.join(map(str, oob_idxs))}) " + "\n\t" + | 
|  | 77 | +                       f" are out of the dataset index range: [0,{data_series.shape[1]})" )  # type: ignore[arg-type] | 
|  | 78 | +            indices1 = list( set( indices ) - set( oob_idxs ) )  # type: ignore[arg-type] | 
| 75 | 79 |         else: | 
| 76 | 80 |             print( f"Error: unsupported indices type: {type(indices)}" ) | 
| 77 | 81 |     else: | 
| 78 |  | -        indices = range( data_series.shape[ 1 ] ) | 
|  | 82 | +        indices1 = list( range( data_series.shape[ 1 ] ) ) | 
| 79 | 83 | 
 | 
|  | 84 | +    components1: list[ int ] = [] | 
| 80 | 85 |     if components is not None: | 
| 81 | 86 |         if type( components ) is int: | 
| 82 |  | -            components = [ components ] | 
| 83 |  | -        if isiterable( components ): | 
| 84 |  | -            oob_comps = list( filter( lambda comp: not 0 <= comp < data_series.shape[ 2 ], components ) ) | 
|  | 87 | +            components1 = [ components ] | 
|  | 88 | +        elif isiterable( components ): | 
|  | 89 | +            oob_comps: list[ int ] = list( | 
|  | 90 | +                filter( | 
|  | 91 | +                    lambda comp: not 0 <= comp < data_series.shape[ 2 ],  # type: ignore[arg-type] | 
|  | 92 | +                    components ) )  # type: ignore[arg-type] | 
| 85 | 93 |             if len( oob_comps ) > 0: | 
| 86 |  | -                print( f"Error: The specified components: ({', '.join(oob_comps)}) " + "\n\t" + | 
|  | 94 | +                print( f"Error: The specified components: ({', '.join(map(str, oob_comps))}) " + "\n\t" + | 
| 87 | 95 |                        " is out of the dataset component range: [0,{data_series.shape[1]})" ) | 
| 88 |  | -            components = list( set( components ) - set( oob_comps ) ) | 
|  | 96 | +            components1 = list( set( components ) - set( oob_comps ) )  # type: ignore[arg-type] | 
| 89 | 97 |         else: | 
| 90 | 98 |             print( f"Error: unsupported components type: {type(components)}" ) | 
| 91 | 99 |     else: | 
| 92 |  | -        components = range( data_series.shape[ 2 ] ) | 
|  | 100 | +        components1 = list( range( data_series.shape[ 2 ] ) ) | 
| 93 | 101 | 
 | 
| 94 |  | -    return [ ( time_series[ :, 0 ], data_series[ :, idx, comp ], idx, comp ) for idx in indices for comp in components ] | 
|  | 102 | +    return [ ( time_series[ :, 0 ], data_series[ :, idx, comp ], idx, comp ) for idx in indices1 | 
|  | 103 | +             for comp in components1 ] | 
| 95 | 104 | 
 | 
| 96 | 105 | 
 | 
| 97 |  | -def commandLinePlotGen(): | 
|  | 106 | +def commandLinePlotGen() -> int: | 
|  | 107 | +    """Parse commande line.""" | 
| 98 | 108 |     parser = argparse.ArgumentParser( | 
| 99 | 109 |         description="A script that parses geosx HDF5 time-history files and produces time-history plots using matplotlib" | 
| 100 | 110 |     ) | 
|  | 
0 commit comments