Adjusted atol in shallow_water_states_tests.

PiperOrigin-RevId: 677822010
google-research · Sep 23, 2024 · ce7209b · ce7209b
1 parent 67c6869
commit ce7209b
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/dinosaur/shallow_water_states_test.py b/dinosaur/shallow_water_states_test.py
@@ -65,21 +65,25 @@ def testZonalVelocity(self, seed):
     latitude = np.linspace(-np.pi / 2, np.pi / 2, 101)
     zonal_velocity = shallow_water_states.get_zonal_velocity(
         latitude, parameters)
+    atol = np.finfo(np.float32).eps
 
     with self.subTest('IsZeroOutsideJet'):
       outside_jet = ((latitude > parameters.jet_northern_lat)
                      + (latitude < parameters.jet_southern_lat))
       np.testing.assert_allclose(outside_jet * zonal_velocity, 0)
 
     with self.subTest('IsBoundedByMaxVelocity'):
-      assert_array_less_equal(np.abs(zonal_velocity),
-                              parameters.jet_max_velocity)
+      assert_array_less_equal(
+          np.abs(zonal_velocity), parameters.jet_max_velocity + atol
+      )
 
     with self.subTest('AttainsMaxVelocityAtJetCenter'):
       center = (parameters.jet_northern_lat + parameters.jet_southern_lat) / 2
       center_velocity = shallow_water_states.get_zonal_velocity(
           center, parameters)
-      np.testing.assert_allclose(center_velocity, parameters.jet_max_velocity)
+      np.testing.assert_allclose(
+          center_velocity, parameters.jet_max_velocity, atol=atol
+      )
 
   @parameterized.parameters(
       dict(seed=0),