Skip to content

Commit e77a83f

Browse files
committed
Fix tests
1 parent 6423820 commit e77a83f

File tree

2 files changed

+106
-39
lines changed

2 files changed

+106
-39
lines changed
Lines changed: 105 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,85 +1,153 @@
11
#!/usr/bin/env python
22
import xarray as xr
33

4+
from mpas_tools.cime.constants import constants
45
from mpas_tools.mesh.conversion import convert, cull
56
from mpas_tools.planar_hex import make_planar_hex_mesh
67
from mpas_tools.translate import center
7-
from mpas_tools.cime.constants import constants
88
from 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

1316
def 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

2334
def 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

5088
def 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

75131
def _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():
88156
if __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()

conda_package/mpas_tools/viz/transect/horiz.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -552,8 +552,7 @@ def find_planar_transect_cells_and_weights(
552552
553553
subdivision_res : float, optional
554554
Resolution in m to use to subdivide the transect when looking for
555-
intersection candidates. Should be small enough that curvature is
556-
small.
555+
intersection candidates.
557556
558557
Returns
559558
-------

0 commit comments

Comments
 (0)