11#!/usr/bin/env python
2+ import os
3+
24import numpy
35import xarray
4- import os
56
67from mpas_tools .io import write_netcdf
7- from mpas_tools .ocean .depth import compute_depth , compute_zmid , add_depth , \
8- add_zmid , write_time_varying_zmid
8+ from mpas_tools .ocean .depth import (
9+ add_depth ,
10+ add_zmid ,
11+ compute_depth ,
12+ compute_zmid ,
13+ write_time_varying_zmid ,
14+ )
915
1016
1117def create_3d_mesh ():
@@ -14,19 +20,25 @@ def create_3d_mesh():
1420 dsMesh = xarray .open_dataset (outFileName )
1521 else :
1622 dsMesh = xarray .open_dataset (
17- 'mesh_tools/mesh_conversion_tools/test/mesh.QU.1920km.151026.nc' )
23+ 'mesh_tools/mesh_conversion_tools/test/mesh.QU.1920km.151026.nc'
24+ )
1825 nCells = dsMesh .sizes ['nCells' ]
1926 nVertLevels = 10
20- zmax = 1000.
21- layerThickness = zmax / nVertLevels
22- dsMesh ['refBottomDepth' ] = \
23- ('nVertLevels' , numpy .linspace (layerThickness , zmax , nVertLevels ))
24- dsMesh ['maxLevelCell' ] = \
25- ('nCells' , nVertLevels * numpy .ones (nCells , dtype = int ))
26- dsMesh ['bottomDepth' ] = ('nCells' , zmax * numpy .ones (nCells ))
27- dsMesh ['layerThickness' ] = \
28- (('Time' , 'nCells' , 'nVertLevels' ),
29- layerThickness * numpy .ones ((1 , nCells , nVertLevels )))
27+ zmax = 1000.0
28+ layerThickness = zmax / nVertLevels
29+ dsMesh ['refBottomDepth' ] = (
30+ 'nVertLevels' ,
31+ numpy .linspace (layerThickness , zmax , nVertLevels ),
32+ )
33+ dsMesh ['maxLevelCell' ] = (
34+ 'nCells' ,
35+ nVertLevels * numpy .ones (nCells , dtype = int ),
36+ )
37+ dsMesh ['bottomDepth' ] = ('nCells' , zmax * numpy .ones (nCells ))
38+ dsMesh ['layerThickness' ] = (
39+ ('Time' , 'nCells' , 'nVertLevels' ),
40+ layerThickness * numpy .ones ((1 , nCells , nVertLevels )),
41+ )
3042 write_netcdf (dsMesh , 'test_depth_mesh.nc' )
3143
3244 return dsMesh
@@ -35,23 +47,30 @@ def create_3d_mesh():
3547def test_compute_depth ():
3648 dsMesh = create_3d_mesh ()
3749 depth , depth_bnds = compute_depth (dsMesh .refBottomDepth )
38- assert numpy .all (numpy .isclose (depth , numpy .linspace (50. , 950. , 10 )))
39- assert numpy .all (numpy .isclose (depth_bnds [:, 0 ],
40- numpy .linspace (0. , 900. , 10 )))
41- assert numpy .all (numpy .isclose (depth_bnds [:, 1 ],
42- numpy .linspace (100. , 1000. , 10 )))
50+ assert numpy .all (numpy .isclose (depth , numpy .linspace (50.0 , 950.0 , 10 )))
51+ assert numpy .all (
52+ numpy .isclose (depth_bnds [:, 0 ], numpy .linspace (0.0 , 900.0 , 10 ))
53+ )
54+ assert numpy .all (
55+ numpy .isclose (depth_bnds [:, 1 ], numpy .linspace (100.0 , 1000.0 , 10 ))
56+ )
4357
4458
4559def test_compute_zmid ():
4660 dsMesh = create_3d_mesh ()
47- zMid = compute_zmid (dsMesh .bottomDepth , dsMesh .maxLevelCell ,
48- dsMesh .layerThickness , depth_dim = 'nVertLevels' )
61+ zMid = compute_zmid (
62+ dsMesh .bottomDepth ,
63+ dsMesh .maxLevelCell ,
64+ dsMesh .layerThickness ,
65+ depth_dim = 'nVertLevels' ,
66+ )
4967
5068 assert zMid .dims == ('Time' , 'nCells' , 'nVertLevels' )
5169
5270 depth = zMid .isel (Time = 0 , nCells = 0 )
53- assert numpy .all (numpy .isclose (depth .values ,
54- numpy .linspace (- 50. , - 950. , 10 )))
71+ assert numpy .all (
72+ numpy .isclose (depth .values , numpy .linspace (- 50.0 , - 950.0 , 10 ))
73+ )
5574
5675
5776def test_add_depth ():
@@ -65,15 +84,18 @@ def test_add_depth():
6584
6685 # test adding depth coordinate once to the mesh and once to the input file,
6786 # with the mesh passed in separately
68- for in_filename , coord_filename in [(mesh_filename , None ),
69- ('test_depth_in.nc' , mesh_filename )]:
87+ for in_filename , coord_filename in [
88+ (mesh_filename , None ),
89+ ('test_depth_in.nc' , mesh_filename ),
90+ ]:
7091 add_depth (in_filename , out_filename , coordFileName = coord_filename )
7192 dsOut = xarray .open_dataset (out_filename )
7293 assert 'depth' in dsOut .dims
7394
7495 depth = dsOut .depth
75- assert numpy .all (numpy .isclose (depth .values ,
76- numpy .linspace (50. , 950. , 10 )))
96+ assert numpy .all (
97+ numpy .isclose (depth .values , numpy .linspace (50.0 , 950.0 , 10 ))
98+ )
7799
78100
79101def test_add_zmid ():
@@ -87,8 +109,10 @@ def test_add_zmid():
87109
88110 # test adding zMid once to the mesh and once to the input file, with the
89111 # mesh passed in separately
90- for in_filename , coord_filename in [(mesh_filename , None ),
91- ('test_depth_in.nc' , mesh_filename )]:
112+ for in_filename , coord_filename in [
113+ (mesh_filename , None ),
114+ ('test_depth_in.nc' , mesh_filename ),
115+ ]:
92116 add_zmid (in_filename , out_filename , coordFileName = coord_filename )
93117 dsOut = xarray .open_dataset (out_filename )
94118 assert 'depth' in dsOut .dims
@@ -97,54 +121,63 @@ def test_add_zmid():
97121 assert zMid .dims == ('nCells' , 'depth' )
98122
99123 depth = zMid .isel (nCells = 0 )
100- assert numpy .all (numpy .isclose (depth .values ,
101- numpy .linspace (- 50. , - 950. , 10 )))
124+ assert numpy .all (
125+ numpy .isclose (depth .values , numpy .linspace (- 50.0 , - 950.0 , 10 ))
126+ )
102127
103128
104129def test_write_time_varying_zmid ():
105-
106130 dsMesh = create_3d_mesh ()
107131 nCells = dsMesh .sizes ['nCells' ]
108132 nVertLevels = dsMesh .sizes ['nVertLevels' ]
109133 mesh_filename = 'test_depth_mesh.nc'
110134 in_filename = 'test_depth_in.nc'
111135 out_filename = 'test_depth_out.nc'
112136
113- layerThickness = 100.
137+ layerThickness = 100.0
114138
115139 # test adding zMid once to the mesh and once to the input file, with the
116140 # mesh passed in separately, each one without and once with a prefix
117- for coord_filename , prefix in [(None , '' ),
118- (mesh_filename , '' ),
119- (None , 'timeMonthly_avg_' ),
120- (mesh_filename , 'timeMonthly_avg_' )]:
141+ for coord_filename , prefix in [
142+ (None , '' ),
143+ (mesh_filename , '' ),
144+ (None , 'timeMonthly_avg_' ),
145+ (mesh_filename , 'timeMonthly_avg_' ),
146+ ]:
121147 print (coord_filename , prefix )
122148
123149 if coord_filename is None :
124150 dsIn = dsMesh .drop_vars ('layerThickness' )
125151 else :
126152 dsIn = xarray .Dataset ()
127- layerThicknessVar = '{}layerThickness' .format (prefix )
128- dsIn [layerThicknessVar ] = \
129- (('Time' , 'nCells' , 'nVertLevels' ),
130- layerThickness * numpy .ones ((2 , nCells , nVertLevels )))
131- dsIn ['{}temperature' .format (prefix )] = \
132- xarray .ones_like (dsIn [layerThicknessVar ])
153+ layerThicknessVar = f'{ prefix } layerThickness'
154+ dsIn [layerThicknessVar ] = (
155+ ('Time' , 'nCells' , 'nVertLevels' ),
156+ layerThickness * numpy .ones ((2 , nCells , nVertLevels )),
157+ )
158+ dsIn [f'{ prefix } temperature' ] = xarray .ones_like (
159+ dsIn [layerThicknessVar ]
160+ )
133161 write_netcdf (dsIn , in_filename )
134162 dsIn .close ()
135163
136- write_time_varying_zmid (in_filename , out_filename ,
137- coordFileName = coord_filename , prefix = prefix )
164+ write_time_varying_zmid (
165+ in_filename ,
166+ out_filename ,
167+ coordFileName = coord_filename ,
168+ prefix = prefix ,
169+ )
138170
139171 dsOut = xarray .open_dataset (out_filename )
140172 assert 'depth' in dsOut .dims
141173
142- zMid = dsOut ['{ }zMid'. format ( prefix ) ]
174+ zMid = dsOut [f' { prefix } zMid' ]
143175 assert zMid .dims == ('Time' , 'nCells' , 'depth' )
144176
145177 depth = zMid .isel (Time = 0 , nCells = 0 )
146- assert numpy .all (numpy .isclose (depth .values ,
147- numpy .linspace (- 50. , - 950. , 10 )))
178+ assert numpy .all (
179+ numpy .isclose (depth .values , numpy .linspace (- 50.0 , - 950.0 , 10 ))
180+ )
148181 dsOut .close ()
149182
150183
0 commit comments