Merge commit '9a20c4ae9414f61d9377ca7e54e8e3ec3d43f9aa' into feature/data_dimnesion_fields

Author: FlorianDeconinck

.github/workflows/fv3_translate_tests.yaml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   fv3_translate_tests:
-    uses: twicki/pyFV3/.github/workflows/translate.yaml@update/numpy_2x
+    uses: NOAA-GFDL/pyFV3/.github/workflows/translate.yaml@develop
     with:
       component_trigger: true
       component_name: NDSL

.github/workflows/pace_tests.yaml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   pace_main_tests:
-    uses: floriandeconinck/pace/.github/workflows/main_unit_tests.yaml@update/numpy_2x
+    uses: NOAA-GFDL/pace/.github/workflows/main_unit_tests.yaml@develop
     with:
       component_trigger: true
       component_name: NDSL

.github/workflows/shield_tests.yaml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   shield_translate_tests:
-    uses: floriandeconinck/pySHiELD/.github/workflows/translate.yaml@update/numpy_2x
+    uses: NOAA-GFDL/pySHiELD/.github/workflows/translate.yaml@develop
     with:
       component_trigger: true
       component_name: NDSL

external/gt4py

@@ -17,15 +17,16 @@ def column_max(field, start_index, end_index):
         Returns: [max value, index of max value]
     """
     max_index = start_index
+    max_value = field.at(K=max_index)
     level = start_index
     while level <= end_index:
-        new = field.at(K=level)
-        old = field.at(K=max_index)
-        if new > old:
+        value = field.at(K=level)
+        if value > max_value:
+            max_value = value
             max_index = level
         level += 1
 
-    return field.at(K=max_index), max_index
+    return max_value, max_index
 
 
 @typing.no_type_check
@@ -42,15 +43,16 @@ def column_max_ddim(field, ddim, start_index, end_index):
         Returns: [max value, index of max value]
     """
     max_index = start_index
+    max_value = field.at(K=max_index, ddim=[ddim])
     level = start_index
     while level <= end_index:
-        new = field.at(K=level, ddim=[ddim])
-        old = field.at(K=max_index, ddim=[ddim])
-        if new > old:
+        value = field.at(K=level, ddim=[ddim])
+        if value > max_value:
+            max_value = value
             max_index = level
         level += 1
 
-    return field.at(K=max_index, ddim=[ddim]), max_index
+    return max_value, max_index
 
 
 @typing.no_type_check
@@ -67,15 +69,16 @@ def column_min(field, start_index, end_index):
         Returns: [min value, index of min value]
     """
     min_index = start_index
+    min_value = field.at(K=min_index)
     level = start_index
     while level <= end_index:
-        new = field.at(K=level)
-        old = field.at(K=min_index)
-        if new < old:
+        value = field.at(K=level)
+        if value < min_value:
+            min_value = value
             min_index = level
         level += 1
 
-    return field.at(K=min_index), min_index
+    return min_value, min_index
 
 
 @typing.no_type_check
@@ -92,12 +95,13 @@ def column_min_ddim(field, ddim, start_index, end_index):
         Returns: [min value, index of min value]
     """
     min_index = start_index
+    min_value = field.at(K=min_index, ddim=[ddim])
     level = start_index
     while level <= end_index:
-        new = field.at(K=level, ddim=[ddim])
-        old = field.at(K=min_index, ddim=[ddim])
-        if new < old:
+        value = field.at(K=level, ddim=[ddim])
+        if value < min_value:
+            min_value = value
             min_index = level
         level += 1
 
-    return field.at(K=min_index, ddim=[ddim]), min_index
+    return min_value, min_index

ndsl/stencils/column_operations.py

@@ -66,10 +66,13 @@ def __init__(
         self.out_vars: dict[str, Any] = {}
         self.write_vars: list = []
         self.grid = grid
-        self.maxshape: tuple[int, ...] = grid.domain_shape_full(add=(1, 1, 1))
         self.ordered_input_vars = None
         self.ignore_near_zero_errors: dict[str, Any] = {}
         self.skip_test = skip_test
+        if self.stencil_factory.backend.is_fortran_aligned():
+            self.maxshape = self.grid.domain_shape_full()
+        else:
+            self.maxshape = self.grid.domain_shape_full(add=(1, 1, 1))
 
     def extra_data_load(self, data_loader: DataLoader):
         pass

ndsl/stencils/testing/translate.py

@@ -27,8 +27,8 @@ def __init__(
         reference_values: np.ndarray,
         computed_values: np.ndarray,
     ):
-        self.references = np.atleast_1d(reference_values)
-        self.computed = np.atleast_1d(computed_values)
+        self.references = np.squeeze(np.atleast_1d(reference_values))
+        self.computed = np.squeeze(np.atleast_1d(computed_values))
         self.check = False
 
     @abstractmethod

ndsl/testing/comparison.py

@@ -2,7 +2,6 @@
 
 import numpy as np
 import numpy.typing as npt
-from numpy._typing import _SupportsDType
 
 from ndsl.config import Backend
 from ndsl.dsl.typing import Float
@@ -14,12 +13,6 @@
 
 # Taking a page from cupy's playbook to have tuple & ndarray
 _ShapeLike = SupportsIndex | Sequence[SupportsIndex]
-_DTypeLikeFloat32 = (
-    np.dtype[np.float32] | _SupportsDType[np.dtype[np.float32]] | type[np.float32]
-)
-_DTypeLikeFloat64 = (
-    np.dtype[np.float64] | _SupportsDType[np.dtype[np.float64]] | type[np.float64]
-)
 
 
 def zeros(
@@ -55,7 +48,7 @@ def empty(
 def full(
     shape: _ShapeLike,
     backend: Backend,
-    value: np.ScalarType,
+    value: npt.DTypeLike,
     dtype: npt.DTypeLike = Float,
 ) -> np.ndarray | cp.ndarray:
     if backend.is_gpu_backend():
@@ -66,7 +59,7 @@ def full(
 def random(
     shape: _ShapeLike,
     backend: Backend,
-    dtype: _DTypeLikeFloat32 | _DTypeLikeFloat64 = Float,  # type: ignore [valid-type]
+    dtype: np.floating = Float,
 ) -> np.ndarray | cp.ndarray:
     if backend.is_gpu_backend():
         gen = cp.random.default_rng()