io.py 5.43 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
import netCDF4

import numpy as np
import numpy.typing as npt

def read_netcdf(filename: str) -> Dict[str, np.ndarray]:
    """"
Eva Börgens's avatar
Eva Börgens committed
16
17
    Reads a TWS netcdf file as provided from GravIS

18
    Args:
Eva Börgens's avatar
Eva Börgens committed
19
        filename: string
Eva Börgens's avatar
Eva Börgens committed
20

Eva Börgens's avatar
Eva Börgens committed
21
    Returns:
Eva Börgens's avatar
Eva Börgens committed
22
        Dict[str, np.ndarray]
Eva Börgens's avatar
Eva Börgens committed
23
    """
Eva Börgens's avatar
Eva Börgens committed
24
    try:
Eva Börgens's avatar
Eva Börgens committed
25
26
27
28
29
30
        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
31
32
        grid_std = np.array([np.nanmean(g) for g in grid_std])
    except Exception:
Eva Börgens's avatar
Eva Börgens committed
33
34
        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
35
36
37
38
        sys.exit()

    return grid, grid_std, lon, lat, time

Eva Börgens's avatar
Eva Börgens committed
39
40
41
def read_ascii(filename: str) -> np.ndarray:
    """"
    Reads an Ascii file with coordinates
Eva Börgens's avatar
Eva Börgens committed
42

43
    Args:
Eva Börgens's avatar
Eva Börgens committed
44
        filename: string
Eva Börgens's avatar
Eva Börgens committed
45

Eva Börgens's avatar
Eva Börgens committed
46
    Returns:
Eva Börgens's avatar
Eva Börgens committed
47
        np.array
Eva Börgens's avatar
Eva Börgens committed
48
    """
Eva Börgens's avatar
Eva Börgens committed
49
50
51
    try:
        with open(filename) as fid:
            coords = []
Eva Börgens's avatar
Eva Börgens committed
52
53
            for l_i in fid:
                if l_i[0]=='#':
Eva Börgens's avatar
Eva Börgens committed
54
                    continue
Eva Börgens's avatar
Eva Börgens committed
55
                line = [float(i) for i in re.split(r'\s+', l_i.strip())]
Eva Börgens's avatar
Eva Börgens committed
56
                coords.append([line[0], line[1]])
Eva Börgens's avatar
Eva Börgens committed
57
58
59
    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
60
61
62
        sys.exit()
    return coords

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

67
    Args:
Eva Börgens's avatar
Eva Börgens committed
68
        filename: string
Eva Börgens's avatar
Eva Börgens committed
69
    """
Eva Börgens's avatar
Eva Börgens committed
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
    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
86
87
88
def same_dist_elems(arr: npt.ArrayLike) -> bool:
    """"
    Tests if coordinates are evenly spaced
Eva Börgens's avatar
Eva Börgens committed
89

90
    Args:
Eva Börgens's avatar
Eva Börgens committed
91
92
        arr: numpy array
    """
Eva Börgens's avatar
Eva Börgens committed
93
    diff = arr[1] - arr[0]
Eva Börgens's avatar
Eva Börgens committed
94
95
    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
96
97
98
            return False
    return True

Eva Börgens's avatar
Eva Börgens committed
99
100
101
102
103
104
105
106
107
108
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
Eva Börgens's avatar
Eva Börgens committed
109

110
    Args:
Eva Börgens's avatar
Eva Börgens committed
111
112
113
114
115
116
117
118
119
120
        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
121
122
123
    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
124
125
126
127
    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
128

Eva Börgens's avatar
Eva Börgens committed
129
        tws_unit = sin.variables['tws'].units
Eva Börgens's avatar
Eva Börgens committed
130
131

    try:
Eva Börgens's avatar
Eva Börgens committed
132
133
134
        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
135
136
137
138
            time_out.units = time_unit
            time_out.calendar = time_calendar
            time_out[:] = time

Eva Börgens's avatar
Eva Börgens committed
139
140
            sout.createDimension('pointID', len(region_coords))
            coords_id = sout.createVariable('pointID', 'i4', 'pointID')
Eva Börgens's avatar
Eva Börgens committed
141
142
            coords_id[:] = np.arange(0, len(region_coords))

Eva Börgens's avatar
Eva Börgens committed
143
            coords_lon = sout.createVariable('lon', 'f4', ('pointID'))
Eva Börgens's avatar
Eva Börgens committed
144
145
146
147
            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
148
            coords_lat = sout.createVariable('lat', 'f4', ('pointID'))
Eva Börgens's avatar
Eva Börgens committed
149
150
151
152
153
            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
154
                uncertainty = sout.createVariable('std', 'f4', 'time')
Eva Börgens's avatar
Eva Börgens committed
155
156
157
158
                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
159
                cov_matrix = sout.createVariable('cov_matrix', 'f4', ('time','pointID', 'pointID'))
Eva Börgens's avatar
Eva Börgens committed
160
161
162
163
                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
164
                timeseries_out = sout.createVariable('tws_timeseries', 'f4', 'time')
Eva Börgens's avatar
Eva Börgens committed
165
166
167
168
169
170
                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()
171