11#!/usr/bin/env python
22import xarray as xr
33
4+ from mpas_tools .cime .constants import constants
45from mpas_tools .mesh .conversion import convert , cull
56from mpas_tools .planar_hex import make_planar_hex_mesh
67from mpas_tools .translate import center
7- from mpas_tools .cime .constants import constants
88from mpas_tools .viz .mesh_to_triangles import mesh_to_triangles
9- from mpas_tools .viz .transects import make_triangle_tree , \
10- find_transect_cells_and_weights , find_planar_transect_cells_and_weights
9+ from mpas_tools .viz .transect import (
10+ find_planar_transect_cells_and_weights ,
11+ find_spherical_transect_cells_and_weights ,
12+ make_triangle_tree ,
13+ )
1114
1215
1316def test_mesh_to_triangles ():
1417 _ , ds_tris = _get_triangles ()
1518 for dim in ['nTriangles' , 'nNodes' , 'nInterp' ]:
1619 assert dim in ds_tris .dims
1720
18- for var in ['triCellIndices' , 'nodeCellIndices' , 'nodeCellWeights' ,
19- 'xNode' , 'yNode' , 'zNode' , 'latNode' , 'lonNode' ]:
21+ for var in [
22+ 'triCellIndices' ,
23+ 'nodeCellIndices' ,
24+ 'nodeCellWeights' ,
25+ 'xNode' ,
26+ 'yNode' ,
27+ 'zNode' ,
28+ 'latNode' ,
29+ 'lonNode' ,
30+ ]:
2031 assert var in ds_tris .data_vars
2132
2233
2334def test_make_triangle_tree ():
2435 _ , ds_tris = _get_triangles ()
2536 tree = make_triangle_tree (ds_tris )
26- assert tree .n == ds_tris .sizes ['nTriangles' ]* ds_tris .sizes ['nNodes' ]
37+ assert tree .n == ds_tris .sizes ['nTriangles' ] * ds_tris .sizes ['nNodes' ]
2738
2839
29- def test_find_transect_cells_and_weights ():
40+ def test_find_spherical_transect_cells_and_weights ():
3041 ds_mesh , ds_tris = _get_triangles ()
3142 tree = make_triangle_tree (ds_tris )
32- lon_transect = xr .DataArray ([- 10. , 0. , 10. ], dims = ('nPoints' ,))
33- lat_transect = xr .DataArray ([- 10. , 0. , 10. ], dims = ('nPoints' ,))
34- ds_transect = find_transect_cells_and_weights (
35- lon_transect , lat_transect , ds_tris , ds_mesh , tree , degrees = True ,
36- subdivisionRes = 10e3 )
37- for dim in ['nSegments' , 'nBounds' , 'nHorizWeights' , 'nPoints' ]:
43+ lon_transect = xr .DataArray ([- 10.0 , 0.0 , 10.0 ], dims = ('nPoints' ,))
44+ lat_transect = xr .DataArray ([- 10.0 , 0.0 , 10.0 ], dims = ('nPoints' ,))
45+ ds_transect = find_spherical_transect_cells_and_weights (
46+ lon_transect ,
47+ lat_transect ,
48+ ds_tris ,
49+ ds_mesh ,
50+ tree ,
51+ degrees = True ,
52+ subdivision_res = 10e3 ,
53+ )
54+
55+ for dim in [
56+ 'nSegments' ,
57+ 'nNodes' ,
58+ 'nHorizBounds' ,
59+ 'nHorizWeights' ,
60+ 'nPoints' ,
61+ ]:
3862 assert dim in ds_transect .dims
3963
40- for var in ['xCartNode' , 'yCartNode' , 'zCartNode' , 'dNode' , 'lonNode' ,
41- 'latNode' , 'horizTriangleIndices' , 'horizCellIndices' ,
42- 'horizTriangleNodeIndices' , 'interpHorizCellIndices' ,
43- 'interpHorizCellWeights' , 'lonTransect' , 'latTransect' ,
44- 'xCartTransect' , 'yCartTransect' , 'zCartTransect' ,
45- 'dTransect' , 'transectIndicesOnHorizNode' ,
46- 'transectWeightsOnHorizNode' ]:
64+ for var in [
65+ 'xCartNode' ,
66+ 'yCartNode' ,
67+ 'zCartNode' ,
68+ 'dNode' ,
69+ 'lonNode' ,
70+ 'latNode' ,
71+ 'horizTriangleIndices' ,
72+ 'horizCellIndices' ,
73+ 'horizTriangleNodeIndices' ,
74+ 'interpHorizCellIndices' ,
75+ 'interpHorizCellWeights' ,
76+ 'lonTransect' ,
77+ 'latTransect' ,
78+ 'xCartTransect' ,
79+ 'yCartTransect' ,
80+ 'zCartTransect' ,
81+ 'dTransect' ,
82+ 'transectIndicesOnHorizNode' ,
83+ 'transectWeightsOnHorizNode' ,
84+ ]:
4785 assert var in ds_transect .data_vars
4886
4987
5088def test_find_planar_transect_cells_and_weights ():
51- ds_mesh = make_planar_hex_mesh (nx = 102 , ny = 52 , dc = 4e3 , nonperiodic_x = True ,
52- nonperiodic_y = True )
89+ ds_mesh = make_planar_hex_mesh (
90+ nx = 102 , ny = 52 , dc = 4e3 , nonperiodic_x = True , nonperiodic_y = True
91+ )
5392 ds_mesh = cull (ds_mesh )
5493 center (ds_mesh )
5594 ds_tris = mesh_to_triangles (ds_mesh )
5695 tree = make_triangle_tree (ds_tris )
5796
58- x_transect = 1e3 * xr .DataArray ([- 10. , 0. , 10. ], dims = ('nPoints' ,))
59- y_transect = 1e3 * xr .DataArray ([- 10. , 0. , 10. ], dims = ('nPoints' ,))
97+ x_transect = 1e3 * xr .DataArray ([- 10.0 , 0.0 , 10.0 ], dims = ('nPoints' ,))
98+ y_transect = 1e3 * xr .DataArray ([- 10.0 , 0.0 , 10.0 ], dims = ('nPoints' ,))
6099
61100 ds_transect = find_planar_transect_cells_and_weights (
62- x_transect , y_transect , ds_tris , ds_mesh , tree , subdivisionRes = 1e3 )
63-
64- for dim in ['nSegments' , 'nBounds' , 'nHorizWeights' , 'nPoints' ]:
101+ x_transect , y_transect , ds_tris , ds_mesh , tree , subdivision_res = 1e3
102+ )
103+
104+ for dim in [
105+ 'nSegments' ,
106+ 'nNodes' ,
107+ 'nHorizBounds' ,
108+ 'nHorizWeights' ,
109+ 'nPoints' ,
110+ ]:
65111 assert dim in ds_transect .dims
66112
67- for var in ['xNode' , 'yNode' , 'dNode' , 'horizTriangleIndices' ,
68- 'horizCellIndices' , 'horizTriangleNodeIndices' ,
69- 'interpHorizCellIndices' , 'interpHorizCellWeights' ,
70- 'xTransect' , 'yTransect' , 'dTransect' ,
71- 'transectIndicesOnHorizNode' , 'transectWeightsOnHorizNode' ]:
113+ for var in [
114+ 'xNode' ,
115+ 'yNode' ,
116+ 'dNode' ,
117+ 'horizTriangleIndices' ,
118+ 'horizCellIndices' ,
119+ 'horizTriangleNodeIndices' ,
120+ 'interpHorizCellIndices' ,
121+ 'interpHorizCellWeights' ,
122+ 'xTransect' ,
123+ 'yTransect' ,
124+ 'dTransect' ,
125+ 'transectIndicesOnHorizNode' ,
126+ 'transectWeightsOnHorizNode' ,
127+ ]:
72128 assert var in ds_transect .data_vars
73129
74130
75131def _get_triangles ():
76132 ds_mesh = xr .open_dataset (
77- 'mesh_tools/mesh_conversion_tools/test/mesh.QU.1920km.151026.nc' )
133+ 'mesh_tools/mesh_conversion_tools/test/mesh.QU.1920km.151026.nc'
134+ )
78135 earth_radius = constants ['SHR_CONST_REARTH' ]
79136 ds_mesh .attrs ['sphere_radius' ] = earth_radius
80- for coord in ['xCell' , 'yCell' , 'zCell' , 'xVertex' , 'yVertex' , 'zVertex' ,
81- 'xEdge' , 'yEdge' , 'zEdge' , 'dcEdge' , 'dvEdge' ]:
82- ds_mesh [coord ] = earth_radius * ds_mesh [coord ]
137+ for coord in [
138+ 'xCell' ,
139+ 'yCell' ,
140+ 'zCell' ,
141+ 'xVertex' ,
142+ 'yVertex' ,
143+ 'zVertex' ,
144+ 'xEdge' ,
145+ 'yEdge' ,
146+ 'zEdge' ,
147+ 'dcEdge' ,
148+ 'dvEdge' ,
149+ ]:
150+ ds_mesh [coord ] = earth_radius * ds_mesh [coord ]
83151 ds_mesh = convert (ds_mesh )
84152 ds_tris = mesh_to_triangles (ds_mesh )
85153 return ds_mesh , ds_tris
@@ -88,5 +156,5 @@ def _get_triangles():
88156if __name__ == '__main__' :
89157 test_mesh_to_triangles ()
90158 test_make_triangle_tree ()
91- test_find_transect_cells_and_weights ()
159+ test_find_spherical_transect_cells_and_weights ()
92160 test_find_planar_transect_cells_and_weights ()
0 commit comments