io.py 5.45 KB
Newer Older
Eva Börgens's avatar
Eva Börgens committed
1
2
3
4
""""
Module handling input and output
"""

Eva Börgens's avatar
Eva Börgens committed
5
import sys
Eva Börgens's avatar
Eva Börgens committed
6
7
import re
from typing import Dict
Eva Börgens's avatar
Eva Börgens committed
8

Eva Börgens's avatar
Eva Börgens committed
9
10
11
12
13
14
15
16
17
18
19
20
21
import netCDF4

import numpy as np
import numpy.typing as npt

def read_netcdf(filename: str) -> Dict[str, np.ndarray]:
    """"
    Reads a TWS netcdf file as provided from GraVIS
    Arg:
        filename: string
    Returns:
        dict with string as keys and numpy arrays as arguments
    """
Eva Börgens's avatar
Eva Börgens committed
22
    try:
Eva Börgens's avatar
Eva Börgens committed
23
24
25
26
27
28
        with netCDF4.Dataset(filename, 'r') as data:
            lon = data.variables['lon'][:]
            lat = data.variables['lat'][:]
            time = data.variables['time'][:]
            grid = data.variables['tws'][:]
            grid_std = data.variables['std_tws'][:]
Eva Börgens's avatar
Eva Börgens committed
29
30
        grid_std = np.array([np.nanmean(g) for g in grid_std])
    except Exception:
Eva Börgens's avatar
Eva Börgens committed
31
32
        print(" Failed to read %s. "
              "Variable names in file have to be: lon, lat, time, tws, std_tws" %filename)
Eva Börgens's avatar
Eva Börgens committed
33
34
35
36
        sys.exit()

    return grid, grid_std, lon, lat, time

Eva Börgens's avatar
Eva Börgens committed
37
38
39
40
41
42
43
44
def read_ascii(filename: str) -> np.ndarray:
    """"
    Reads an Ascii file with coordinates
    Arg:
        filename: string
    Returns:
        numpy array
    """
Eva Börgens's avatar
Eva Börgens committed
45
46
47
    try:
        with open(filename) as fid:
            coords = []
Eva Börgens's avatar
Eva Börgens committed
48
49
            for l_i in fid:
                if l_i[0]=='#':
Eva Börgens's avatar
Eva Börgens committed
50
                    continue
Eva Börgens's avatar
Eva Börgens committed
51
                line = [float(i) for i in re.split(r'\s+', l_i.strip())]
Eva Börgens's avatar
Eva Börgens committed
52
                coords.append([line[0], line[1]])
Eva Börgens's avatar
Eva Börgens committed
53
54
55
    except Exception:
        print("Could not read %s: Please ensure that every line contains two columns and "
              "comment lines are indicated with '#'" %filename)
Eva Börgens's avatar
Eva Börgens committed
56
57
58
        sys.exit()
    return coords

Eva Börgens's avatar
Eva Börgens committed
59
60
61
62
63
64
65
def test_coordiantes(filename: str):
    """"
    Tests if more than 2 coordinates are given in the file
    Arg:
        filename: string

    """
Eva Börgens's avatar
Eva Börgens committed
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
    coords = read_ascii(filename)

    lon = np.array([r[0] for r in coords])
    lat = np.array([r[1] for r in coords])
    if len(lon)<2:
        raise RuntimeError("Coordinate list of region have to contain at least 2 points")

    u_lon = np.unique(lon)
    u_lat = np.unique(lat)

    if not same_dist_elems(u_lon):
        raise RuntimeError("Longitudes of region not evenly spaced")
    if not same_dist_elems(u_lat):
        raise RuntimeError("Latitudes of region not evenly spaced")


Eva Börgens's avatar
Eva Börgens committed
82
83
84
85
86
87
def same_dist_elems(arr: npt.ArrayLike) -> bool:
    """"
    Tests if coordinates are evenly spaced
    Arg:
        arr: numpy array
    """
Eva Börgens's avatar
Eva Börgens committed
88
    diff = arr[1] - arr[0]
Eva Börgens's avatar
Eva Börgens committed
89
90
    for x_i in range(1, len(arr) - 1):
        if arr[x_i + 1] - arr[x_i] != diff:
Eva Börgens's avatar
Eva Börgens committed
91
92
93
            return False
    return True

Eva Börgens's avatar
Eva Börgens committed
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
def save_results(outname: str,
                 inname: str,
                 results: Dict[str, npt.ArrayLike],
                 region_coords:  npt.ArrayLike,
                 timeseries:  npt.ArrayLike,
                 flag_uncertainty: bool,
                 flag_matrix: bool,
                 flag_timeseries: bool):
    """"
    Saving the results of the regional tws uncerty tools to a NetCDF file
    Arg:
        outname: str
        inname: str
        results: Dict[str, np.ndarray]
        region_coords: np.array
        timeseries: np.array
        flag_uncertainty: bool
        flag_matrix: bool
        flag_timeseries: bool

    """
Eva Börgens's avatar
Eva Börgens committed
115
116
117
    lon = np.array([r[0] for r in region_coords])
    lat = np.array([r[1] for r in region_coords])

Eva Börgens's avatar
Eva Börgens committed
118
119
120
121
    with netCDF4.Dataset(inname, format='NETCDF4') as sin:
        time = sin.variables['time'][:]
        time_unit = sin.variables['time'].units
        time_calendar = sin.variables['time'].calendar
Eva Börgens's avatar
Eva Börgens committed
122

Eva Börgens's avatar
Eva Börgens committed
123
        tws_unit = sin.variables['tws'].units
Eva Börgens's avatar
Eva Börgens committed
124
125

    try:
Eva Börgens's avatar
Eva Börgens committed
126
127
128
        with netCDF4.Dataset(outname, 'w', format='NETCDF4') as sout:
            sout.createDimension('time', len(time))
            time_out = sout.createVariable('time', 'i4', 'time')
Eva Börgens's avatar
Eva Börgens committed
129
130
131
132
            time_out.units = time_unit
            time_out.calendar = time_calendar
            time_out[:] = time

Eva Börgens's avatar
Eva Börgens committed
133
134
            sout.createDimension('pointID', len(region_coords))
            coords_id = sout.createVariable('pointID', 'i4', 'pointID')
Eva Börgens's avatar
Eva Börgens committed
135
136
            coords_id[:] = np.arange(0, len(region_coords))

Eva Börgens's avatar
Eva Börgens committed
137
            coords_lon = sout.createVariable('lon', 'f4', ('pointID'))
Eva Börgens's avatar
Eva Börgens committed
138
139
140
141
            coords_lon[:] = lon
            coords_lon.long_name = "Lon coordinates tuple of the grid points of the region"
            coords_lon.units = "degrees_north"

Eva Börgens's avatar
Eva Börgens committed
142
            coords_lat = sout.createVariable('lat', 'f4', ('pointID'))
Eva Börgens's avatar
Eva Börgens committed
143
144
145
146
147
            coords_lat[:] = lat
            coords_lat.long_name = "Lat coordinates tuple of the grid points of the region"
            coords_lat.units = "degrees_east"

            if flag_uncertainty:
Eva Börgens's avatar
Eva Börgens committed
148
                uncertainty = sout.createVariable('std', 'f4', 'time')
Eva Börgens's avatar
Eva Börgens committed
149
150
151
152
                uncertainty.units = tws_unit
                uncertainty.long_name = "standard deviation of regions mean tws time series"
                uncertainty[:] = results['uncertainty']
            if flag_matrix:
Eva Börgens's avatar
Eva Börgens committed
153
                cov_matrix = sout.createVariable('cov_matrix', 'f4', ('time','pointID', 'pointID'))
Eva Börgens's avatar
Eva Börgens committed
154
155
156
157
                cov_matrix.long_name = "Covariance matrix between all points of the region"
                cov_matrix.units = tws_unit + '*' + tws_unit
                cov_matrix[:] = results['matrix']
            if flag_timeseries:
Eva Börgens's avatar
Eva Börgens committed
158
                timeseries_out = sout.createVariable('tws_timeseries', 'f4', 'time')
Eva Börgens's avatar
Eva Börgens committed
159
160
161
162
163
164
                timeseries_out.units = tws_unit
                timeseries_out.long_name = "Regions mean tws time series"
                timeseries_out[:] = timeseries
    except PermissionError:
        print("Could not open %s for writing results" %outname)
        sys.exit()