Add support for VECTOR SELECT

docs/structure/vensim_translation.rst

@@ -118,4 +118,3 @@ Planed New Functions and Features
 ---------------------------------
 - ALLOCATE BY PRIORITY
 - SHIFT IF TRUE
-- VECTOR SELECT

docs/tables/functions.tab

@@ -21,10 +21,10 @@ IF THEN ELSE	"IF THEN ELSE(A, B, C)"	if_then_else	"if_then_else(A, B, C)"	"CallS
 XIDZ	"XIDZ(A, B, X)"	safediv	"safediv(A, B, X)"	"CallStructure('xidz', (A, B, X))"	"pysd.functions.xidz(A, B, X)"
 ZIDZ	"ZIDZ(A, B)"	safediv	"safediv(A, B)"	"CallStructure('zidz', (A, B))"	"pysd.functions.zidz(A, B)"
 
-VMIN	VMIN(A)			"CallStructure('vmin', (A,))"	pysd.functions.vmin(A)
-VMAX	VMAX(A)			"CallStructure('vmax', (A,))"	pysd.functions.vmax(A)
-SUM	SUM(A)			"CallStructure('sum', (A,))"	pysd.functions.sum(A)
-PROD	PROD(A)			"CallStructure('prod', (A,))"	pysd.functions.prod(A)
+VMIN	VMIN(A[dim!])			"CallStructure('vmin', (A,))"	pysd.functions.vmin(A, ['dim!'])
+VMAX	VMAX(A[dim!])			"CallStructure('vmax', (A,))"	pysd.functions.vmax(A, ['dim!'])
+SUM	SUM(A[dim!])			"CallStructure('sum', (A,))"	pysd.functions.sum(A, ['dim!'])
+PROD	PROD(A[dim!])			"CallStructure('prod', (A,))"	pysd.functions.prod(A, ['dim!'])
 
 PULSE	PULSE(start, width)			"CallStructure('pulse', (start, width))"	pysd.functions.pulse(start, width=width)
 		pulse	pulse(magnitude, start)	"CallStructure('Xpulse', (start, magnitude))"	pysd.functions.pulse(start, magnitude=magnitude)		Not tested for Xmile!
@@ -34,6 +34,7 @@ RAMP	RAMP(slope, start_time, end_time)	ramp	ramp(slope, start_time, end_time)	"C
 		ramp	ramp(slope, start_time)	"CallStructure('ramp', (slope, start_time))"	pysd.functions.ramp(time, slope, start_time)		Not tested for Xmile!
 STEP	STEP(height, step_time)	step	step(height, step_time)	"CallStructure('step', (height, step_time))"	pysd.functions.step(time, height, step_time)		Not tested for Xmile!
 GET TIME VALUE	GET TIME VALUE(relativeto, offset, measure)			"CallStructure('get_time_value', (relativeto, offset, measure))"	pysd.functions.get_time_value(time, relativeto, offset, measure)	Not all the cases implemented!
+VECTOR SELECT	VECTOR SELECT(sel_array[dim!], exp_array[dim!], miss_val, n_action, e_action)			"CallStructure('vector_select', (sel_array, exp_array, miss_val, n_action, e_action))"	pysd.functions.vector_select(sel_array, exp_array, ['dim!'], miss_val, n_action, e_action)
 VECTOR RANK	VECTOR RANK(vec, direction)			"CallStructure('vector_rank', (vec, direction))"	vector_rank(vec, direction)
 VECTOR REORDER	VECTOR REORDER(vec, svec)			"CallStructure('vector_reorder', (vec, svec))"	vector_reorder(vec, svec)
 VECTOR SORT ORDER	VECTOR SORT ORDER(vec, direction)			"CallStructure('vector_sort_order', (vec, direction))"	vector_sort_order(vec, direction)

docs/whats_new.rst

@@ -2,12 +2,13 @@
 What's New
 ==========
 
-v3.2.0 (unreleased)
+v3.2.0 (2022/06/10)
 -------------------
 
 New Features
 ~~~~~~~~~~~~
 - Add support for Vensim's `GET TIME VALUE <https://www.vensim.com/documentation/fn_get_time_value.html>`_ (:func:`pysd.py_backend.functions.get_time_value`) function (:issue:`332`). Not all cases have been implemented.
+- Add support for Vensim's `VECTOR SELECT <http://vensim.com/documentation/fn_vector_select.html>`_ (:func:`pysd.py_backend.functions.vector_select`) function (:issue:`266`).
 
 Breaking changes
 ~~~~~~~~~~~~~~~~

pysd/builders/python/python_expressions_builder.py

@@ -644,7 +644,14 @@ def build_function_call(self, arguments: dict) -> BuildAST:
         if "%(axis)s" in expression:
             # Vectorial expressions, compute the axis using dimensions
             # with ! operator
-            final_subscripts, arguments["axis"] = self._compute_axis(arguments)
+            if "%(1)s" in expression:
+                subs = self.reorder(arguments)
+                # NUMPY: following line may be avoided
+                [arguments[i].reshape(self.section.subscripts, subs, True)
+                 for i in ["0", "1"]]
+            else:
+                subs = arguments["0"].subscripts
+            final_subscripts, arguments["axis"] = self._compute_axis(subs)
 
         elif "%(size)s" in expression:
             # Random expressions, need to give the final size of the
@@ -713,14 +720,14 @@ def build_function_call(self, arguments: dict) -> BuildAST:
             subscripts=final_subscripts,
             order=0)
 
-    def _compute_axis(self, arguments: dict) -> tuple:
+    def _compute_axis(self, subscripts: dict) -> tuple:
         """
         Compute the axis to apply a vectorial function.
 
         Parameters
         ----------
-        arguments: dict
-            The dictionary of builded arguments.
+        subscripts: dict
+            The final_subscripts after reordering all the elements.
 
         Returns
         -------
@@ -731,7 +738,6 @@ def _compute_axis(self, arguments: dict) -> tuple:
             dimensions with "!" at the end.
 
         """
-        subscripts = arguments["0"].subscripts
         axis = []
         coords = {}
         for subs in subscripts:

pysd/builders/python/python_functions.py

@@ -33,6 +33,10 @@
     "sum": ("sum(%(0)s, dim=%(axis)s)", ("functions", "sum")),
     "vmax": ("vmax(%(0)s, dim=%(axis)s)", ("functions", "vmax")),
     "vmin": ("vmin(%(0)s, dim=%(axis)s)", ("functions", "vmin")),
+    "vector_select": (
+        "vector_select(%(0)s, %(1)s, %(axis)s, %(2)s, %(3)s, %(4)s)",
+        ("functions", "vector_select")
+    ),
 
     # functions defined in pysd.py_bakcend.functions
     "active_initial": (

pysd/py_backend/functions.py

@@ -476,6 +476,145 @@ def invert_matrix(mat):
     return xr.DataArray(np.linalg.inv(mat.values), mat.coords, mat.dims)
 
 
+def vector_select(selection_array, expression_array, dim,
+                  missing_vals, numerical_action, error_action):
+    """
+    Implements Vensim's VECTOR SELECT function.
+    http://vensim.com/documentation/fn_vector_select.html
+
+    Parameters
+    ----------
+    selection_array: xr.DataArray
+        This specifies a selection array with a mixture of zeroes and
+        non-zero values.
+    expression_array: xarray.DataArray
+        This is the expression that elements are being selected from
+        based on the selection array.
+    dim: list of strs
+        Dimensions to apply the function over.
+    missing_vals: float
+        The value to use in the case where there are only zeroes in the
+        selection array.
+    numerical_action: int
+        The action to take:
+            - 0 It will calculate the weighted sum.
+            - 1 When values in the selection array are non-zero, this
+              will calculate the product of the
+              selection_array * expression_array.
+            - 2 The weighted minimum, for non zero values of the
+              selection array, this is minimum of
+              selection_array * expression_array.
+            - 3 The weighted maximum, for non zero values of the
+              selection array, this is maximum of
+              selection_array * expression_array.
+            - 4 For non zero values of the selection array, this is
+              the average of selection_array * expression_array.
+            - 5 When values in the selection array are non-zero,
+              this will calculate the product of the
+              expression_array ^ selection_array.
+            - 6 When values in the selection array are non-zero,
+              this will calculate the sum of the expression_array.
+              The same as the SUM function for non-zero values in
+              the selection array.
+            - 7 When values in the selection array are non-zero,
+              this will calculate the product of the expression_array.
+              The same as the PROD function for non-zero values in
+              the selection array.
+            - 8 The unweighted minimum, for non zero values of the
+              selection array, this is minimum of the expression_array.
+              The same as the VMIN function for non-zero values in
+              the selection array.
+            - 9 The unweighted maximum, for non zero values of the
+              selection array, this is maximum of expression_array.
+              The same as the VMAX function for non-zero values in
+              the selection array.
+            - 10 For non zero values of the selection array,
+              this is the average of expression_array.
+    error_action: int
+        Indicates how to treat too many or too few entries in the selection:
+            - 0 No error is raised.
+            - 1 Raise a floating point error is selection array only
+              contains zeros.
+            - 2 Raise an error if the selection array contains more
+              than one non-zero value.
+            - 3 Raise an error if all elements in selection array are
+              zero, or more than one element is non-zero
+              (this is a combination of error_action = 1 and error_action = 2).
+
+    Returns
+    -------
+    result: xarray.DataArray or float
+        The output of the numerical action.
+
+    """
+    zeros = (selection_array == 0).all(dim=dim)
+    non_zeros = (selection_array != 0).sum(dim=dim)
+
+    # Manage error actions
+    if np.any(zeros) and error_action in (1, 3):
+        raise FloatingPointError(
+            "All the values of selection_array are 0...")
+
+    if np.any(non_zeros > 1) and error_action in (2, 3):
+        raise FloatingPointError(
+            "More than one non-zero values in selection_array...")
+
+    # Manage numeric actions (array to operate)
+    # NUMPY: replace by np.where
+    if numerical_action in range(5):
+        array = xr.where(
+            selection_array == 0,
+            np.nan,
+            selection_array * expression_array
+        )
+    elif numerical_action == 5:
+        warnings.warn(
+            "Vensim's help says that numerical_action=5 computes the "
+            "product of selection_array ^ expression_array. But, in fact,"
+            " Vensim is computing the product of expression_array ^ "
+            " selection array. The output of this function behaves as "
+            "Vensim, expression_array ^ selection_array."
+        )
+        array = xr.where(
+            selection_array == 0,
+            np.nan,
+            expression_array ** selection_array
+        )
+    elif numerical_action in range(6, 11):
+        array = xr.where(
+            selection_array == 0,
+            np.nan,
+            expression_array
+        )
+    else:
+        raise ValueError(
+            f"Invalid argument value 'numerical_action={numerical_action}'. "
+            "'numerical_action' must be 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 or 10.")
+
+    # Manage numeric actions (operation)
+    # NUMPY: use the axis
+    if numerical_action in (0, 6):
+        out = array.sum(dim=dim, skipna=True)
+    elif numerical_action in (1, 5, 7):
+        out = array.prod(dim=dim, skipna=True)
+    elif numerical_action in (2, 8):
+        out = array.min(dim=dim, skipna=True)
+    elif numerical_action in (3, 9):
+        out = array.max(dim=dim, skipna=True)
+    elif numerical_action in (4, 10):
+        out = array.mean(dim=dim, skipna=True)
+
+    # Replace missin vals
+    if len(out.shape) == 0 and np.all(zeros):
+        return missing_vals
+    elif len(out.shape) == 0:
+        return float(out)
+    elif np.any(zeros):
+        out.values[zeros.values] = missing_vals
+
+    return out
+
+
 def vector_sort_order(vector, direction):
     """
     Implements Vensim's VECTOR SORT ORDER function. Sorting is done on

tests/pytest_integration/pytest_integration_vensim_pathway.py

@@ -517,6 +517,10 @@
         "folder": "vector_order",
         "file": "test_vector_order.mdl"
     },
+    "vector_select": {
+        "folder": "vector_select",
+        "file": "test_vector_select.mdl"
+    },
     "xidz_zidz": {
         "folder": "xidz_zidz",
         "file": "xidz_zidz.mdl"

tests/pytest_pysd/pytest_functions.py

@@ -7,7 +7,7 @@
 from pysd.py_backend.components import Time
 from pysd.py_backend.functions import\
     ramp, step, pulse, xidz, zidz, if_then_else, sum, prod, vmin, vmax,\
-    invert_matrix, get_time_value
+    invert_matrix, get_time_value, vector_select
 
 
 class TestInputFunctions():
@@ -439,6 +439,90 @@ def test_get_time_value_errors(self, measure, relativeto,
             get_time_value(
                 lambda: 0, relativeto, np.random.randint(-100, 100), measure)
 
+    def test_vector_select(self):
+        warning_message =\
+            r"Vensim's help says that numerical_action=5 computes the "\
+            r"product of selection_array \^ expression_array\. But, in fact,"\
+            r" Vensim is computing the product of expression_array \^ "\
+            r" selection array\. The output of this function behaves as "\
+            r"Vensim, expression_array \^ selection_array\."
+
+        array = xr.DataArray([3, 10, 2], {'dim': ["A", "B", "C"]})
+        sarray = xr.DataArray([1, 0, 2], {'dim': ["A", "B", "C"]})
+
+        with pytest.warns(UserWarning, match=warning_message):
+            assert vector_select(sarray, array, ["dim"], np.nan, 5, 1)\
+                == 12
+
+        sarray = xr.DataArray([0, 0, 0], {'dim': ["A", "B", "C"]})
+        assert vector_select(sarray, array, ["dim"], 123, 0, 2) == 123
+
+    @pytest.mark.parametrize(
+        "selection_array,expression_array,dim,numerical_action,"
+        "error_action,raise_type,error_message",
+        [
+            (  # error_action=1
+                xr.DataArray([0, 0], {'dim': ["A", "B"]}),
+                xr.DataArray([1, 2], {'dim': ["A", "B"]}),
+                ["dim"],
+                0,
+                1,
+                FloatingPointError,
+                r"All the values of selection_array are 0\.\.\."
+            ),
+            (  # error_action=2
+                xr.DataArray([1, 1], {'dim': ["A", "B"]}),
+                xr.DataArray([1, 2], {'dim': ["A", "B"]}),
+                ["dim"],
+                0,
+                2,
+                FloatingPointError,
+                r"More than one non-zero values in selection_array\.\.\."
+            ),
+            (  # error_action=3a
+                xr.DataArray([0, 0], {'dim': ["A", "B"]}),
+                xr.DataArray([1, 2], {'dim': ["A", "B"]}),
+                ["dim"],
+                0,
+                3,
+                FloatingPointError,
+                r"All the values of selection_array are 0\.\.\."
+            ),
+            (  # error_action=3b
+                xr.DataArray([1, 1], {'dim': ["A", "B"]}),
+                xr.DataArray([1, 2], {'dim': ["A", "B"]}),
+                ["dim"],
+                0,
+                3,
+                FloatingPointError,
+                r"More than one non-zero values in selection_array\.\.\."
+            ),
+            (  # numerical_action=11
+                xr.DataArray([1, 1], {'dim': ["A", "B"]}),
+                xr.DataArray([1, 2], {'dim': ["A", "B"]}),
+                ["dim"],
+                11,
+                0,
+                ValueError,
+                r"Invalid argument value 'numerical_action=11'\. "
+                r"'numerical_action' must be 0, 1, 2, 3, 4, 5, 6, "
+                r"7, 8, 9 or 10\."
+            ),
+        ],
+        ids=[
+            "error_action=1", "error_action=2", "error_action=3a",
+            "error_action=3b", "numerical_action=11"
+        ]
+    )
+    def test_vector_select_errors(self, selection_array, expression_array,
+                                  dim, numerical_action, error_action,
+                                  raise_type, error_message):
+
+        with pytest.raises(raise_type, match=error_message):
+            vector_select(
+                selection_array, expression_array, dim,  0,
+                numerical_action, error_action)
+
     def test_incomplete(self):
         from pysd.py_backend.functions import incomplete