-
Notifications
You must be signed in to change notification settings - Fork 0
/
spatial_mean.py
41 lines (34 loc) · 1.17 KB
/
spatial_mean.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
import xarray as xr
import numpy as np
import utils as u
import sys
def main(argv):
inputfile = argv[0]
variable = argv[1]
n = int(argv[2])
output = argv[3]
outvar = 'mean_' + variable
confirmation = (f'Calculate spatial mean of {n} x {n} center pixels of '
f'variable {variable} in {inputfile} and save it '
f'as {outvar} in {output}? [y/n] ')
ok = input(confirmation)
if ok != 'y':
print('Aborting...')
exit()
print(f'Opening dataset')
ds = xr.open_dataset(inputfile)
center = u.crop_center(ds[variable], n)
print('Storing crop mask')
nonspatialcoords = [c for c in ds[variable].coords if c not in ['x', 'y']]
ds['cropped_area'] = xr.ones_like(
center.isel({c: 0 for c in nonspatialcoords}).drop(nonspatialcoords),
dtype=np.bool8
)
print(f'Calculating mean and std for variable {variable}')
ds['mean_' + variable] = center.mean(dim=['x', 'y'])
ds['std_' + variable] = center.std(dim=['x', 'y'])
ds = ds.drop(variable)
print(f'Saving result to {output}')
ds.to_netcdf(output)
if __name__ == '__main__':
main(sys.argv[1:])