io.py 5.38 KB
Newer Older
1
"""
Eva Börgens's avatar
Eva Börgens committed
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
import netCDF4

import numpy as np

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

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

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

    return grid, grid_std, lon, lat, time

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

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

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

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

66
    Args:
Eva Börgens's avatar
Eva Börgens committed
67
        filename: string
Eva Börgens's avatar
Eva Börgens committed
68
    """
Eva Börgens's avatar
Eva Börgens committed
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
    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")


85
86
def same_dist_elems(arr: np.ndarray) -> bool:
    """
Eva Börgens's avatar
Eva Börgens committed
87
    Tests if coordinates are evenly spaced
Eva Börgens's avatar
Eva Börgens committed
88

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

Eva Börgens's avatar
Eva Börgens committed
98
99
def save_results(outname: str,
                 inname: str,
100
101
102
                 results: Dict[str, np.ndarray],
                 region_coords:  np.ndarray,
                 timeseries:  np.ndarray,
Eva Börgens's avatar
Eva Börgens committed
103
104
105
                 flag_uncertainty: bool,
                 flag_matrix: bool,
                 flag_timeseries: bool):
106
    """
Eva Börgens's avatar
Eva Börgens committed
107
    Saving the results of the regional tws uncerty tools to a NetCDF file
Eva Börgens's avatar
Eva Börgens committed
108

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

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

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

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

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