New Basic Operations

ndsl/stencils/__init__.py

@@ -1,28 +1,65 @@
 from .basic_operations import (
-    adjust_divide_stencil,
-    adjustmentfactor_stencil,
-    average_in,
+    add,
+    add_to_self,
     copy,
-    dim,
+    divide,
+    divide_self,
+    multiply,
+    multiply_to_self,
     select_k,
-    set_IJ_mask_value,
+    set_bool_value_2D,
     set_value,
-    set_value_2D,
+    set_value_2d,
     sign,
+    subtract,
+    subtract_from_self,
 )
+from .basic_operations_2d import (
+    add_2d,
+    add_to_self_2d,
+    copy_2d,
+    divide_2d,
+    divide_self_2d,
+    multiply_2d,
+    multiply_to_self_2d,
+    set_boolean_value_2d,
+    set_value_2d,
+    subtract_2d,
+    subtract_from_self_2d,
+)
+from .column_operations import column_max, column_max_ddim, column_min, column_min_ddim
 from .corners import FillCornersBGrid
 
 
 __all__ = [
     "FillCornersBGrid",
+    "add",
+    "add_to_self",
     "copy",
-    "adjustmentfactor_stencil",
-    "set_value",
-    "set_value_2D",
-    "set_IJ_mask_value",
-    "adjust_divide_stencil",
+    "divide",
+    "divide_self",
+    "multiply",
+    "multiply_to_self",
     "select_k",
-    "average_in",
+    "set_bool_value_2D",
+    "set_value",
+    "set_value_2d",
     "sign",
-    "dim",
+    "subtract",
+    "subtract_from_self",
+    "add_2d",
+    "add_to_self_2d",
+    "copy_2d",
+    "divide_2d",
+    "divide_self_2d",
+    "multiply_2d",
+    "multiply_to_self_2d",
+    "set_value_2d",
+    "subtract_2d",
+    "subtract_from_self_2d",
+    "set_boolean_value_2d",
+    "column_max",
+    "column_max_ddim",
+    "column_min",
+    "column_min_ddim",
 ]

ndsl/stencils/arithmetical_functions.py

@@ -0,0 +1,27 @@
+import typing
+
+from ndsl.dsl.gt4py import function
+
+
+@typing.no_type_check
+@function
+def sign(a, b):
+    """
+    Defines a_sign_b as the absolute value of a, and checks if b is positive or
+    negative, assigning the analogous sign value to a_sign_b. a_sign_b is returned.
+
+    Args:
+        a: A number
+        b: A number
+    """
+    a_sign_b = abs(a)
+    return a_sign_b if b > 0 else -a_sign_b
+
+
+@typing.no_type_check
+@function
+def dim(a, b):
+    """
+    Calculates a - b, camped to 0, i.e. max(a - b, 0).
+    """
+    return max(a - b, 0)

ndsl/stencils/basic_operations.py

@@ -1,95 +1,191 @@
-import typing
-
-from ndsl.dsl.gt4py import FORWARD, PARALLEL, computation, function, interval
+from ndsl.dsl.gt4py import FORWARD, PARALLEL, K, computation, interval
 from ndsl.dsl.typing import (
     Bool,
     BoolFieldIJ,
     Float,
     FloatField,
     FloatFieldIJ,
-    IntField,
     IntFieldIJ,
 )
 
 
-def copy(q_in: FloatField, q_out: FloatField) -> None:
+def copy(input: FloatField, output: FloatField) -> None:
     """
-    Copy q_in to q_out.
+    Copy one field into another.
 
     Args:
-        q_in: input field
-        q_out: output field
+        input: input field
+        output: output field
     """
     with computation(PARALLEL), interval(...):
-        q_out = q_in
+        output = input
 
 
-def adjustmentfactor_stencil(adjustment: FloatFieldIJ, q_out: FloatField) -> None:
+def add(summand_1: FloatField, summand_2: FloatField, sum: FloatField) -> None:
     """
-    Multiplies every element of q_out by every element of the adjustment field over the
-    interval, replacing the elements of q_out by the result of the multiplication.
+    Add two inputs together, sum to a new field.
 
     Args:
-        adjustment: adjustment field
-        q_out: output field
+        summand_1: input field
+        summand_2: input field
+        sum: output field
     """
     with computation(PARALLEL), interval(...):
-        q_out = q_out * adjustment
+        sum = summand_1 + summand_2
 
 
-def set_value(q_out: FloatField, value: Float) -> None:
+def add_to_self(field: FloatField, summand: FloatField) -> None:
     """
-    Sets every element of q_out to the value specified by value argument.
+    Add a summand to a field.
 
     Args:
-        q_out: output field
-        value: NDSL Float type
+        field: field to be modifid
+        summand: modification to be made
+    """
+    with computation(PARALLEL), interval(...):
+        field = field + summand
+
+
+def subtract(
+    minuend: FloatField, subtrahend: FloatField, difference: FloatField
+) -> None:
+    """
+    Subtract subtrahend from minuend, output to a new field.
+
+    Args:
+        minuend: input field
+        subtrahend: input field
+        difference: output field
+    """
+    with computation(PARALLEL), interval(...):
+        difference = minuend - subtrahend
+
+
+def subtract_from_self(field: FloatField, subtrahend: FloatField) -> None:
+    """
+    Subtract a subtrahend from a field.
+
+    Args:
+        field: field to be modifid
+        subtrahend: modification to be made
     """
     with computation(PARALLEL), interval(...):
-        q_out = value
+        field = field - subtrahend
 
 
-def set_value_2D(buffer: FloatFieldIJ, value: Float) -> None:
+def multiply(factor_1: FloatField, factor_2: FloatField, product: FloatField) -> None:
     """
-    Sets every element of buffer to the value specified by value argument.
+    Multiply two inputs together, output to a new field.
 
     Args:
-        buffer: output field
+        factor_1: input field
+        factor_2: input field
+        output: output field
+    """
+    with computation(PARALLEL), interval(...):
+        product = factor_1 * factor_2
+
+
+def multiply_to_self(field: FloatField, factor: FloatField) -> None:
+    """
+    Muultiply a field by a factor.
+
+    Args:
+        field: field to be modifid
+        factor: modification factor
+    """
+    with computation(PARALLEL), interval(...):
+        field = field * factor
+
+
+def divide(dividend: FloatField, divisor: FloatField, quotient: FloatField) -> None:
+    """
+    Divide dividend by divisor, output to a new field.
+
+    Args:
+        dividend: input field
+        divisor: input field
+        quotient: output field
+    """
+    with computation(PARALLEL), interval(...):
+        quotient = dividend / divisor
+
+
+def divide_self(field: FloatField, divisor: FloatField) -> None:
+    """
+    Muultiply a field by a factor - 2D variant.
+
+    Args:
+        field: field to be modifid
+        divisor: modification factor
+    """
+    with computation(PARALLEL), interval(...):
+        field = field / divisor
+
+
+def set_value(field: FloatField, value: Float) -> None:
+    """
+    Sets every element of a field to a single value.
+
+    Args:
+        field: output field
         value: value of Float type
     """
-    with computation(FORWARD), interval(0, 1):
-        buffer = value
+    with computation(PARALLEL), interval(...):
+        field = value
+
+
+def set_value_2d(field: FloatFieldIJ, value: Float) -> None:
+    """
+    Sets every element of a field to a single value - 2D variant.
+
+    Args:
+        field: output field
+        value: value of Float type
+    """
+    with computation(FORWARD), interval(...):
+        field = value
 
 
-def set_IJ_mask_value(mask_out: BoolFieldIJ, value: Bool) -> None:
+def set_bool_value_2D(field: BoolFieldIJ, value: Bool) -> None:
     """
-    Sets every element of buffer to the value specified by value argument.
+    Sets every element of buffer to either True or False.
 
     Args:
-        mask_out: output field
+        field: output field
         value: value of Bool type
     """
     with computation(FORWARD), interval(0, 1):
-        mask_out = value
+        field = value
 
 
-def adjust_divide_stencil(adjustment: FloatField, q_out: FloatField) -> None:
+def adjustmentfactor_stencil(adjustment: FloatFieldIJ, field: FloatField) -> None:
     """
-    Divides every element of q_out by every element of the adjustment field over the
-    interval, replacing the elements of q_out by the result of the division.
+    Multiplies a field by an adjustment factor, modifying the original field.
 
     Args:
-        adjustment: adjustment field
-        q_out: output field
+        adjustment: adjustment factor
+        field: field to be modified
+    """
+    with computation(PARALLEL), interval(...):
+        field = field * adjustment
+
+
+def adjust_divide_stencil(adjustment: FloatField, field: FloatField) -> None:
+    """
+    Divides a field by an adjustment factor, modifying the original field.
+
+    Args:
+        adjustment: adjustment factor
+        field: field to be modified
     """
     with computation(PARALLEL), interval(...):
-        q_out = q_out / adjustment
+        field = field / adjustment
 
 
 def select_k(
     in_field: FloatField,
     out_field: FloatFieldIJ,
-    k_mask: IntField,
     k_select: IntFieldIJ,
 ) -> None:
     """
@@ -99,17 +195,15 @@ def select_k(
     Args:
         in_field: A 3D array to select from
         out_field: A 2D field to save values in
-        k_mask: a field that lists each k-index
         k_select: the k-value to extract from in_field
     """
-    # TODO: refactor this using THIS_K instead of a mask
     with computation(FORWARD), interval(...):
-        if k_mask == k_select:
+        if K == k_select:
             out_field = in_field
 
 
-def average_in(
-    q_out: FloatField,
+def average_input(
+    field: FloatField,
     adjustment: FloatField,
 ) -> None:
     """
@@ -121,28 +215,4 @@ def average_in(
         q_out: output field
     """
     with computation(PARALLEL), interval(...):
-        q_out = (q_out + adjustment) * 0.5
-
-
-@typing.no_type_check
-@function
-def sign(a, b):
-    """
-    Defines a_sign_b as the absolute value of a, and checks if b is positive or
-    negative, assigning the analogous sign value to a_sign_b. a_sign_b is returned.
-
-    Args:
-        a: A number
-        b: A number
-    """
-    a_sign_b = abs(a)
-    return a_sign_b if b > 0 else -a_sign_b
-
-
-@typing.no_type_check
-@function
-def dim(a, b):
-    """
-    Calculates a - b, camped to 0, i.e. max(a - b, 0).
-    """
-    return max(a - b, 0)
+        field = (field + adjustment) * 0.5