Merge pull request #59 from mnlevy1981/update_tests

Get CircleCI tests to pass:

1. increase timeout limit
2. provide workaround for failing upstream test

Also introduced new tests for the SOR fill algorithm and refactored `test_fill.py` to reduce duplicate code when initializing fields.

.circleci/config.yml

@@ -40,6 +40,7 @@ shared: &shared
         name: Running Tests
         command: |
           pytest --junitxml=test-reports/junit.xml --cov=./ --verbose
+        no_output_timeout: 45m
     - run:
         name: Uploading code coverage report
         command: |

ci/environment-upstream-dev.yml

@@ -6,6 +6,7 @@ dependencies:
   - bottleneck
   - codecov
   - docrep
+  - netcdf4
   - numba
   - pip
   - pooch

ci/environment.yml

@@ -8,6 +8,7 @@ dependencies:
   - dask
   - docrep
   - nbsphinx
+  - netcdf4
   - numba
   - pip
   - pooch

pop_tools/data_registry.txt

@@ -6,6 +6,8 @@ cesm_pop_monthly.T62_g17.nc dca4d607f82348cda83a35b2187c522672095aa775334115dda3
 lateral_fill_np_array_filled_ref.npz 0e7b16ce8df50d0c904e36aac22868ead3ba8f2a00d5eaeb857dc84ecffb4de8
 lateral_fill_np_array_tripole_filled_ref.npz 96e10e924525a6c3b14be4b3f97f41a4c9622938d140585c1f9067fc0f5a2d10
 lateral_fill_np_array_tripole_filled_ref.20200818.npz a6e5ac70da6762d8423dda3675fe03230ff8a55dc47bc15d6f47c1c53ca9e2d9
+lateral_fill_np_array_filled_SOR_ref.20200820.npz b7c1d573cb0f503b6a4bde31cf8a22b7cd930db6939d00e13e46eba08aeac5e8
+lateral_fill_np_array_tripole_filled_SOR_ref.20200820.npz f6d803bd05cd6a16857683e08f7ccf2e2f9fcfefc6f50fde2fc3c3d8784da7f4
 POP_gx3v7.nc 7586ba68ee5b4d8de8ffe5a11974bff81a412c62af9de230088575c8016e4b84
 g.e20.G.TL319_t13.control.001_hfreq.nc 439eb1abf14737341ead088bfd9a3c1b795dc1f79bc9052c09db08fef9ab83e5
 g.e20.G.TL319_t13.control.001_hfreq-coarsen.nc 145659813daf1a607d0857c10699c721693333e39aeb270e116103185b7236ae

pop_tools/eos.py

@@ -100,6 +100,8 @@ def eos(salt, temp, return_coefs=False, **kwargs):
         salt = xr.where(salt > smax, smax, salt)
 
         salt, temp, pressure = xr.broadcast(salt, temp, pressure)
+        if isinstance(salt.data, dask.array.Array):
+            pressure = pressure.chunk(salt.chunks)
 
         if return_coefs:
             RHO, dRHOdS, dRHOdT = _compute_eos_coeffs(salt, temp, pressure)

tests/test_fill.py

@@ -9,23 +9,10 @@
 
 def test_lateral_fill_np_array():
 
-    # generate psuedo-data
-    dx, dy = 0.05, 0.05
-    y, x = np.mgrid[slice(1, 3 + dy, dy), slice(1, 5 + dx, dx)]
-    z_orig = np.sin(x) ** 10 + np.cos(10 + y * x) * np.cos(x)
-
-    # construct mask and apply mask
-    valid_points = np.ones(z_orig.shape, dtype=np.bool)
-    valid_points = np.where(y < 0.5 * np.sin(5 * x) + 1.5, False, valid_points)
-    z_orig = np.where(~valid_points, np.nan, z_orig)
-
-    # add missing values
-    z_miss = z_orig.copy()
-    z_miss[:20, 62:] = np.nan
-    z_miss[15:18, 0:2] = 10.0
+    valid_points, z = _generate_2D_test_dataset()
 
     # compute lateral fill
-    z_fill = pop_tools.lateral_fill_np_array(z_miss, valid_points)
+    z_fill = pop_tools.lateral_fill_np_array(z, valid_points)
 
     # load reference data
     ref_data_file = DATASETS.fetch('lateral_fill_np_array_filled_ref.npz')
@@ -38,28 +25,45 @@ def test_lateral_fill_np_array():
 
 def test_lateral_fill_np_array_ltripole():
 
-    # generate psuedo-data
-    dx, dy = 0.05, 0.05
-    y, x = np.mgrid[slice(1 - dy, 3 + dy, dy), slice(1 - dx, 5 + dx, dx)]
-    z_orig = np.sin(x) ** 10 + np.cos(10 + y * x) * np.cos(x)
+    valid_points, z = _generate_2D_test_dataset(tripole=True)
 
-    # construct mask and apply mask
-    valid_points = np.ones(z_orig.shape, dtype=np.bool)
-    valid_points = np.where(y < 0.5 * np.sin(5 * x) + 1.5, False, valid_points)
-    z_orig = np.where(~valid_points, np.nan, z_orig)
+    # compute lateral fill
+    z_fill = pop_tools.lateral_fill_np_array(z, valid_points, ltripole=True)
 
-    # add missing values
-    z_miss = z_orig.copy()
-    z_miss[:20, 62:] = np.nan
-    z_miss[35:, 55:70] = np.nan
-    z_miss[15:18, 0:2] = 10.0
-    z_miss[-2:, 12:20] = 10.0
+    # load reference data
+    ref_data_file = DATASETS.fetch('lateral_fill_np_array_tripole_filled_ref.20200818.npz')
+    with np.load(ref_data_file) as data:
+        z_fill_ref = data['arr_0']
+
+    # assert that we match the reference solution
+    np.testing.assert_allclose(z_fill, z_fill_ref, atol=1e-14, equal_nan=True, verbose=True)
+
+
+def test_lateral_fill_np_array_SOR():
+
+    valid_points, z = _generate_2D_test_dataset()
 
     # compute lateral fill
-    z_fill = pop_tools.lateral_fill_np_array(z_miss, valid_points, ltripole=True)
+    z_fill = pop_tools.lateral_fill_np_array(z, valid_points, use_sor=True)
 
     # load reference data
-    ref_data_file = DATASETS.fetch('lateral_fill_np_array_tripole_filled_ref.20200818.npz')
+    ref_data_file = DATASETS.fetch('lateral_fill_np_array_filled_SOR_ref.20200820.npz')
+    with np.load(ref_data_file) as data:
+        z_fill_ref = data['arr_0']
+
+    # assert that we match the reference solution
+    np.testing.assert_allclose(z_fill, z_fill_ref, atol=1e-14, equal_nan=True, verbose=True)
+
+
+def test_lateral_fill_np_array_ltripole_SOR():
+
+    valid_points, z = _generate_2D_test_dataset(tripole=True)
+
+    # compute lateral fill
+    z_fill = pop_tools.lateral_fill_np_array(z, valid_points, ltripole=True, use_sor=True)
+
+    # load reference data
+    ref_data_file = DATASETS.fetch('lateral_fill_np_array_tripole_filled_SOR_ref.20200820.npz')
     with np.load(ref_data_file) as data:
         z_fill_ref = data['arr_0']
 
@@ -167,3 +171,28 @@ def test_lateral_fill_4D_3Dmask():
             np.testing.assert_array_equal(arr_0, arr_i)
 
     assert da_out.attrs == attrs
+
+
+def _generate_2D_test_dataset(tripole=False):
+    # generate psuedo-data
+    dx, dy = 0.05, 0.05
+    if tripole:
+        y, x = np.mgrid[slice(1 - dy, 3 + dy, dy), slice(1 - dx, 5 + dx, dx)]
+    else:
+        y, x = np.mgrid[slice(1, 3 + dy, dy), slice(1, 5 + dx, dx)]
+    z_orig = np.sin(x) ** 10 + np.cos(10 + y * x) * np.cos(x)
+
+    # construct mask and apply mask
+    valid_points = np.ones(z_orig.shape, dtype=np.bool)
+    valid_points = np.where(y < 0.5 * np.sin(5 * x) + 1.5, False, valid_points)
+    z_orig = np.where(~valid_points, np.nan, z_orig)
+
+    # add missing values
+    z_miss = z_orig.copy()
+    z_miss[:20, 62:] = np.nan
+    z_miss[15:18, 0:2] = 10.0
+    if tripole:
+        z_miss[35:, 55:70] = np.nan
+        z_miss[-2:, 12:20] = 10.0
+
+    return valid_points, z_miss