Merge pull request #391 from xylar/add_depth_coord_testing

xylar · web-flow · commit 026d6968fb6d · 2021-01-12T12:12:15.000+01:00
Fix further bugs in ocean depth module
diff --git a/azure-pipelines.yml b/azure-pipelines.yml
@@ -24,6 +24,8 @@ jobs:
         python.version: '3.7'
       Python38:
         python.version: '3.8'
+      Python39:
+        python.version: '3.9'
 
   steps:
   - bash: echo "##vso[task.prependpath]$CONDA/bin"
@@ -141,6 +143,8 @@ jobs:
         python.version: '3.7'
       Python38:
         python.version: '3.8'
+      Python39:
+        python.version: '3.9'
 
   steps:
   - bash: echo "##vso[task.prependpath]$CONDA/bin"
diff --git a/conda_package/ci/linux_python3.6.yaml b/conda_package/ci/linux_python3.6.yaml
@@ -3,7 +3,7 @@ channel_sources:
 cxx_compiler:
 - gxx
 cxx_compiler_version:
-- '7'
+- '9'
 hdf5:
 - 1.10.6
 libnetcdf:
diff --git a/conda_package/ci/linux_python3.7.yaml b/conda_package/ci/linux_python3.7.yaml
@@ -3,7 +3,7 @@ channel_sources:
 cxx_compiler:
 - gxx
 cxx_compiler_version:
-- '7'
+- '9'
 hdf5:
 - 1.10.6
 libnetcdf:
diff --git a/conda_package/ci/linux_python3.8.yaml b/conda_package/ci/linux_python3.8.yaml
@@ -3,7 +3,7 @@ channel_sources:
 cxx_compiler:
 - gxx
 cxx_compiler_version:
-- '7'
+- '9'
 hdf5:
 - 1.10.6
 pin_run_as_build:
diff --git a/conda_package/ci/linux_python3.9.yaml b/conda_package/ci/linux_python3.9.yaml
@@ -0,0 +1,14 @@
+channel_sources:
+- conda-forge,defaults
+cxx_compiler:
+- gxx
+cxx_compiler_version:
+- '9'
+hdf5:
+- 1.10.6
+pin_run_as_build:
+  python:
+    min_pin: x.x
+    max_pin: x.x
+python:
+- 3.9.* *_cpython
diff --git a/conda_package/ci/osx_python3.6.yaml b/conda_package/ci/osx_python3.6.yaml
@@ -3,7 +3,7 @@ channel_sources:
 cxx_compiler:
 - clangxx
 cxx_compiler_version:
-- '10'
+- '11'
 hdf5:
 - 1.10.6
 libnetcdf:
diff --git a/conda_package/ci/osx_python3.7.yaml b/conda_package/ci/osx_python3.7.yaml
@@ -3,7 +3,7 @@ channel_sources:
 cxx_compiler:
 - clangxx
 cxx_compiler_version:
-- '10'
+- '11'
 hdf5:
 - 1.10.6
 libnetcdf:
diff --git a/conda_package/ci/osx_python3.8.yaml b/conda_package/ci/osx_python3.8.yaml
@@ -3,7 +3,7 @@ channel_sources:
 cxx_compiler:
 - clangxx
 cxx_compiler_version:
-- '10'
+- '11'
 hdf5:
 - 1.10.6
 libnetcdf:
diff --git a/conda_package/ci/osx_python3.9.yaml b/conda_package/ci/osx_python3.9.yaml
@@ -0,0 +1,16 @@
+channel_sources:
+- conda-forge,defaults
+cxx_compiler:
+- clangxx
+cxx_compiler_version:
+- '11'
+hdf5:
+- 1.10.6
+libnetcdf:
+- 4.7.4
+pin_run_as_build:
+  python:
+    min_pin: x.x
+    max_pin: x.x
+python:
+- 3.9.* *_cpython
diff --git a/conda_package/mpas_tools/ocean/depth.py b/conda_package/mpas_tools/ocean/depth.py
@@ -89,7 +89,10 @@ def compute_zmid(bottomDepth, maxLevelCell, layerThickness,
     zMid = zLayerBot + 0.5*layerThickness
 
     zMid = zMid.where(vertIndex < maxLevelCell)
-    zMid = zMid.transpose('Time', 'nCells', depth_dim)
+    if 'Time' in zMid.dims:
+        zMid = zMid.transpose('Time', 'nCells', depth_dim)
+    else:
+        zMid = zMid.transpose('nCells', depth_dim)
 
     return zMid
 
@@ -202,6 +205,8 @@ def add_zmid(inFileName, outFileName, coordFileName=None):
         # dsCoord doesn't have masking disabled because we want it for zMid
         dsCoord = xarray.open_dataset(coordFileName)
         dsCoord = dsCoord.rename({'nVertLevels': 'depth'})
+        if 'Time' in dsCoord.dims:
+            dsCoord = dsCoord.isel(Time=0)
 
         ds.coords['zMid'] = compute_zmid(dsCoord.bottomDepth,
                                          dsCoord.maxLevelCell,
@@ -254,7 +259,7 @@ def main_add_zmid():
 
 
 def write_time_varying_zmid(inFileName, outFileName, coordFileName=None,
-                            prefix=None):
+                            prefix=''):
     """
     Add a 3D, time-independent depth coordinate to an MPAS-Ocean file.
 
diff --git a/conda_package/mpas_tools/tests/define_base_mesh.py b/conda_package/mpas_tools/tests/define_base_mesh.py
diff --git a/conda_package/mpas_tools/tests/test_conversion.py b/conda_package/mpas_tools/tests/test_conversion.py
@@ -19,9 +19,6 @@ def test_conversion():
     dsCulled = cull(dsIn=dsMesh, dsMask=dsMask)
     write_netcdf(dsCulled, 'culled_mesh.nc')
 
-    dsMask = xarray.open_dataset(
-        'mesh_tools/mesh_conversion_tools/test/land_mask_final.nc')
-
     fcMask = read_feature_collection(
         'mesh_tools/mesh_conversion_tools/test/Arctic_Ocean.geojson')
     dsMask = mask(dsMesh=dsMesh, fcMask=fcMask)
diff --git a/conda_package/mpas_tools/tests/test_depth.py b/conda_package/mpas_tools/tests/test_depth.py
@@ -0,0 +1,156 @@
+#!/usr/bin/env python
+import numpy
+import xarray
+import os
+
+from mpas_tools.io import write_netcdf
+from mpas_tools.ocean.depth import compute_depth, compute_zmid, add_depth, \
+    add_zmid, write_time_varying_zmid
+
+
+def create_3d_mesh():
+    outFileName = 'test_depth_mesh.nc'
+    if os.path.exists(outFileName):
+        dsMesh = xarray.open_dataset(outFileName)
+    else:
+        dsMesh = xarray.open_dataset(
+            'mesh_tools/mesh_conversion_tools/test/mesh.QU.1920km.151026.nc')
+        nCells = dsMesh.sizes['nCells']
+        nVertLevels = 10
+        zmax = 1000.
+        layerThickness = zmax/nVertLevels
+        dsMesh['refBottomDepth'] = \
+            ('nVertLevels', numpy.linspace(layerThickness, zmax, nVertLevels))
+        dsMesh['maxLevelCell'] = \
+            ('nCells', nVertLevels*numpy.ones(nCells, dtype=int))
+        dsMesh['bottomDepth'] = ('nCells', zmax*numpy.ones(nCells))
+        dsMesh['layerThickness'] = \
+            (('Time', 'nCells', 'nVertLevels'),
+             layerThickness*numpy.ones((1, nCells, nVertLevels)))
+        write_netcdf(dsMesh, 'test_depth_mesh.nc')
+
+    return dsMesh
+
+
+def test_compute_depth():
+    dsMesh = create_3d_mesh()
+    depth, depth_bnds = compute_depth(dsMesh.refBottomDepth)
+    assert numpy.all(numpy.isclose(depth, numpy.linspace(50., 950., 10)))
+    assert numpy.all(numpy.isclose(depth_bnds[:, 0],
+                                   numpy.linspace(0., 900., 10)))
+    assert numpy.all(numpy.isclose(depth_bnds[:, 1],
+                                   numpy.linspace(100., 1000., 10)))
+
+
+def test_compute_zmid():
+    dsMesh = create_3d_mesh()
+    zMid = compute_zmid(dsMesh.bottomDepth, dsMesh.maxLevelCell,
+                        dsMesh.layerThickness, depth_dim='nVertLevels')
+
+    assert zMid.dims == ('Time', 'nCells', 'nVertLevels')
+
+    depth = zMid.isel(Time=0, nCells=0)
+    assert numpy.all(numpy.isclose(depth.values,
+                                   numpy.linspace(-50., -950., 10)))
+
+
+def test_add_depth():
+    dsMesh = create_3d_mesh()
+    mesh_filename = 'test_depth_mesh.nc'
+    out_filename = 'test_depth_out.nc'
+
+    dsIn = xarray.Dataset()
+    dsIn['temperature'] = xarray.ones_like(dsMesh.layerThickness)
+    write_netcdf(dsIn, 'test_depth_in.nc')
+
+    # test adding depth coordinate once to the mesh and once to the input file,
+    # with the mesh passed in separately
+    for in_filename, coord_filename in [(mesh_filename, None),
+                                        ('test_depth_in.nc', mesh_filename)]:
+        add_depth(in_filename, out_filename, coordFileName=coord_filename)
+        dsOut = xarray.open_dataset(out_filename)
+        assert 'depth' in dsOut.dims
+
+        depth = dsOut.depth
+        assert numpy.all(numpy.isclose(depth.values,
+                                       numpy.linspace(50., 950., 10)))
+
+
+def test_add_zmid():
+    dsMesh = create_3d_mesh()
+    mesh_filename = 'test_depth_mesh.nc'
+    out_filename = 'test_depth_out.nc'
+
+    dsIn = xarray.Dataset()
+    dsIn['temperature'] = xarray.ones_like(dsMesh.layerThickness)
+    write_netcdf(dsIn, 'test_depth_in.nc')
+
+    # test adding zMid once to the mesh and once to the input file, with the
+    # mesh passed in separately
+    for in_filename, coord_filename in [(mesh_filename, None),
+                                        ('test_depth_in.nc', mesh_filename)]:
+        add_zmid(in_filename, out_filename, coordFileName=coord_filename)
+        dsOut = xarray.open_dataset(out_filename)
+        assert 'depth' in dsOut.dims
+
+        zMid = dsOut.zMid
+        assert zMid.dims == ('nCells', 'depth')
+
+        depth = zMid.isel(nCells=0)
+        assert numpy.all(numpy.isclose(depth.values,
+                                       numpy.linspace(-50., -950., 10)))
+
+
+def test_write_time_varying_zmid():
+
+    dsMesh = create_3d_mesh()
+    nCells = dsMesh.sizes['nCells']
+    nVertLevels = dsMesh.sizes['nVertLevels']
+    mesh_filename = 'test_depth_mesh.nc'
+    in_filename = 'test_depth_in.nc'
+    out_filename = 'test_depth_out.nc'
+
+    layerThickness = 100.
+
+    # test adding zMid once to the mesh and once to the input file, with the
+    # mesh passed in separately, each one without and once with a prefix
+    for coord_filename, prefix in [(None, ''),
+                                   (mesh_filename, ''),
+                                   (None, 'timeMonthly_avg_'),
+                                   (mesh_filename, 'timeMonthly_avg_')]:
+        print(coord_filename, prefix)
+
+        if coord_filename is None:
+            dsIn = dsMesh.drop_vars('layerThickness')
+        else:
+            dsIn = xarray.Dataset()
+        layerThicknessVar = '{}layerThickness'.format(prefix)
+        dsIn[layerThicknessVar] = \
+            (('Time', 'nCells', 'nVertLevels'),
+             layerThickness*numpy.ones((2, nCells, nVertLevels)))
+        dsIn['{}temperature'.format(prefix)] = \
+            xarray.ones_like(dsIn[layerThicknessVar])
+        write_netcdf(dsIn, in_filename)
+        dsIn.close()
+
+        write_time_varying_zmid(in_filename, out_filename,
+                                coordFileName=coord_filename, prefix=prefix)
+
+        dsOut = xarray.open_dataset(out_filename)
+        assert 'depth' in dsOut.dims
+
+        zMid = dsOut['{}zMid'.format(prefix)]
+        assert zMid.dims == ('Time', 'nCells', 'depth')
+
+        depth = zMid.isel(Time=0, nCells=0)
+        assert numpy.all(numpy.isclose(depth.values,
+                                       numpy.linspace(-50., -950., 10)))
+        dsOut.close()
+
+
+if __name__ == '__main__':
+    test_compute_depth()
+    test_compute_zmid()
+    test_add_depth()
+    test_add_zmid()
+    test_write_time_varying_zmid()
diff --git a/conda_package/recipe/conda_build_config.yaml b/conda_package/recipe/conda_build_config.yaml
@@ -8,3 +8,4 @@ python:
    - 3.6
    - 3.7
    - 3.8
+   - 3.9
