Fix some issues

pysd/py_backend/builder.py

@@ -565,6 +565,117 @@ def add_n_delay(identifier, delay_input, delay_time, initial_value, order,
     return "%s()" % stateful['py_name'], new_structure
 
 
+def add_sample_if_true(identifier, condition, actual_value, initial_value,
+                subs, subscript_dict):
+    """
+    Creates code to instantiate a stateful 'SampleIfTrue' object,
+    and provides reference to that object's output.
+
+    Parameters
+    ----------
+    identifier: basestring
+        the python-safe name of the stock
+
+    condition: <string>
+        Reference to another model element that is the condition to the 
+        'sample if true' function
+
+    actual_value: <string>
+        Can be a number (in string format) or a reference to another model
+        element which is calculated throughout the simulation at runtime.
+
+    initial_value: <string>
+        This is used to initialize the state of the sample if true function.
+
+    subs: list of strings
+        List of strings of subscript indices that correspond to the
+        list of expressions, and collectively define the shape of the output
+
+    subscript_dict: dictionary
+        Dictionary describing the possible dimensions of the stock's subscripts
+
+    Returns
+    -------
+    reference: basestring
+        reference to the sample if true object `__call__` method, 
+        which will return the output of the sample if true process
+
+    new_structure: list
+        list of element construction dictionaries for the builder to assemble
+
+    """
+    import_modules['functions'].add("SampleIfTrue")
+
+    new_structure = []
+
+    if len(subs) == 0:
+        stateful_py_expr = 'SampleIfTrue(lambda: %s, lambda: %s,'\
+                           'lambda: %s)' % (condition, actual_value, initial_value)
+
+    else:
+        stateful_py_expr = 'SampleIfTrue(lambda: _condition_%s(),'\
+                           'lambda: _input_%s(), lambda: _init_%s(),)' % (
+                               identifier, identifier, identifier)
+        # following elements not specified in the model file, but must exist
+        # create the sample if true initialization element
+        new_structure.append({
+            'py_name': '_init_%s' % identifier,
+            'real_name': 'Implicit',
+            'kind': 'setup',  # not specified in the model file, but must exist
+            'py_expr': initial_value,
+            'subs': subs,
+            'doc': 'Provides initial value for %s function' % identifier,
+            'unit': 'See docs for %s' % identifier,
+            'lims': 'None',
+            'eqn': 'None',
+            'arguments': ''
+        })
+
+        new_structure.append({
+            'py_name': '_condition_%s' % identifier,
+            'real_name': 'Implicit',
+            'kind': 'component',
+            'doc': 'Provides condition for %s function' % identifier,
+            'subs': subs,
+            'unit': 'See docs for %s' % identifier,
+            'lims': 'None',
+            'eqn': 'None',
+            'py_expr': condition,
+            'arguments': ''
+        })
+
+        new_structure.append({
+            'py_name': '_input_%s' % identifier,
+            'real_name': 'Implicit',
+            'kind': 'component',
+            'doc': 'Provides input for %s function' % identifier,
+            'subs': subs,
+            'unit': 'See docs for %s' % identifier,
+            'lims': 'None',
+            'eqn': 'None',
+            'py_expr': actual_value,
+            'arguments': ''
+        })
+
+    # describe the stateful object
+    stateful = {
+        'py_name': '_sample_if_true_%s' % identifier,
+        'real_name': 'Sample if true of %s' % identifier,
+        'doc': 'Initial value: %s \n  Input: %s \n Condition: %s' % (
+            initial_value, actual_value, condition),
+        'py_expr': stateful_py_expr,
+        'unit': 'None',
+        'lims': 'None',
+        'eqn': 'None',
+        'subs': '',
+        'kind': 'stateful',
+        'arguments': ''
+    }
+    new_structure.append(stateful)
+
+    return "%s()" % stateful['py_name'], new_structure
+
+
 def add_n_smooth(identifier, smooth_input, smooth_time, initial_value, order,
                  subs, subscript_dict):
     """

pysd/py_backend/functions.py

@@ -170,6 +170,33 @@ def ddt(self):
             inflows[0] = self.input_func()
         return inflows - outflows
 
+class SampleIfTrue(Stateful):
+    def __init__(self, condition, actual_value, initial_value):
+        """
+
+        Parameters
+        ----------
+        condition: function
+        actual_value: function
+        initial_value: function
+        """
+        super().__init__()
+        self.condition = condition
+        self.actual_value = actual_value
+        self.initial_value = initial_value
+
+    def initialize(self):
+        self.state = self.initial_value()
+        if isinstance(self.state, xr.DataArray):
+            self.shape_info = {'dims': self.state.dims,
+                               'coords': self.state.coords}
+
+    def __call__(self):
+        self.state = sample_if_true(self.condition(), self.actual_value(), self.state)
+        return self.state
+
+    def ddt(self):
+        return 0
 
 class Smooth(Stateful):
     def __init__(self, smooth_input, smooth_time, initial_value, order):
@@ -1091,6 +1118,34 @@ def if_then_else(condition, val_if_true, val_if_false):
     return val_if_true() if condition else val_if_false()
 
 
+def sample_if_true(condition, actual_value, saved_value):
+    """
+    Implements Vensim's SAMPLE IF TRUE function.
+
+    Parameters
+    ----------
+    condition: bool or xarray.DataArray of bools
+    actual_value: int, float or xarray.DataArray
+        Value to return when condition is true.
+    saved_value: int, float or xarray.DataArray
+        Value to return when condition is false.
+
+    Returns
+    -------
+    The value depending on the condition.
+
+    """
+    if isinstance(condition, xr.DataArray):
+        if condition.all():
+            return actual_value
+        elif not condition.any():
+            return saved_value
+
+        return xr.where(condition, actual_value, saved_value)
+
+    return actual_value if condition else saved_value
+
+
 def xidz(numerator, denominator, value_if_denom_is_zero):
     """
     Implements Vensim's XIDZ function.

pysd/py_backend/vensim/vensim2py.py

@@ -268,7 +268,7 @@ def get_equation_components(equation_str, root_path=None):
     component_structure_grammar = _include_common_grammar(r"""
     entry = component / data_definition / test_definition / subscript_definition / lookup_definition
     component = name _ subscriptlist? _ "=" "="? _ expression
-    subscript_definition = name _ ":" _ (imported_subscript / literal_subscript)
+    subscript_definition = name _ ":" _ (imported_subscript / literal_subscript / numeric_range)
     data_definition = name _ subscriptlist? _ keyword? _ ":=" _ expression
     lookup_definition = name _ subscriptlist? &"(" _ expression  # uses lookahead assertion to capture whole group
     test_definition = name _ subscriptlist? _ &keyword _ expression
@@ -277,11 +277,15 @@ def get_equation_components(equation_str, root_path=None):
 
     literal_subscript = subscript _ ("," _ subscript _)*
     imported_subscript = imp_subs_func _ "(" _ (string _ ","? _)* ")"
+    numeric_range = _ (range / value) _ ("," _ (range / value) _)*
+    value = _ sequence_id _ 
+    range = "(" _ sequence_id _ "-" _ sequence_id _ ")"
     subscriptlist = '[' _ subscript _ ("," _ subscript _)* _ ']'
 
     expression = ~r".*"  # expression could be anything, at this point.
     keyword = ":" _ basic_id _ ":"
 
+    sequence_id = _ basic_id _
     subscript = basic_id / escape_group
     imp_subs_func = ~r"(%(imp_subs)s)"IU
     string = "\'" ( "\\\'" / ~r"[^\']"IU )* "\'"
@@ -328,6 +332,31 @@ def visit_imported_subscript(self, n, vc):
             self.subscripts +=\
                 external.ExtSubscript(*args, root=root_path).subscript
 
+        def visit_range(self, n, vc):
+            subs_start = vc[2].strip()
+            subs_end = vc[6].strip()
+            if(subs_start == subs_end): raise ValueError('Only different subscripts are valid in a numeric range, error in expression:\n\t %s\n' % (equation_str))
+
+            # get the common prefix and the starting and 
+            # ending number of the numeric range
+            subs_start = re.findall('\d+|\D+', subs_start)
+            subs_end = re.findall('\d+|\D+', subs_end)
+            prefix_start = ''.join(subs_start[:-1])
+            prefix_end = ''.join(subs_end[:-1])
+            num_start = int(subs_start[-1])
+            num_end = int(subs_end[-1])
+
+            if(not(prefix_start) or not(prefix_end)): raise ValueError('A numeric range must contain at least one letter, error in expression:\n\t %s\n' % (equation_str))
+            if(num_start>num_end): raise ValueError('The number of the first subscript value must be lower than the second subscript value in a subscript numeric range, error in expression:\n\t %s\n'% (equation_str))
+            if(prefix_start != prefix_end or subs_start[0].isdigit() or subs_end[0].isdigit()): raise ValueError('Only matching names ending in numbers are valid, error in expression:\n\t %s\n'% (equation_str))
+
+            for i in range(num_start, num_end+1):
+                s = prefix_start + str(i)
+                self.subscripts.append(s.strip())
+
+        def visit_value(self, n, vc):
+            self.subscripts.append(vc[1].strip())
+
         def visit_name(self, n, vc):
             (name,) = vc
             self.real_name = name.strip()
@@ -513,10 +542,7 @@ def parse_units(units_str):
                       "original_name": "INVERT MATRIX"},
     "get time value": {"name": "not_implemented_function",
                       "module": "functions",
-                      "original_name": "GET TIME VALUE"},
-    "sample if true": {"name": "not_implemented_function",
-                      "module": "functions",
-                      "original_name": "SAMPLE IF TRUE"}
+                      "original_name": "GET TIME VALUE"}
 }
 
 # list of fuctions that accept a dimension to apply over
@@ -625,7 +651,16 @@ def parse_units(units_str):
         subs=element['subs'],
         subscript_dict=subscript_dict
     ),
-
+
+    "sample if true": lambda element, subscript_dict, args: builder.add_sample_if_true(
+        identifier=element['py_name'],
+        condition=args[0],
+        actual_value=args[1],
+        initial_value=args[2],
+        subs=element['subs'],
+        subscript_dict=subscript_dict
+    ),
+
     "smooth": lambda element, subscript_dict, args: builder.add_n_smooth(
         identifier=element['py_name'],
         smooth_input=args[0],

tests/integration_test_vensim_pathway.py

@@ -288,9 +288,7 @@ def test_rounding(self):
         output, canon = runner('test-models/tests/rounding/test_rounding.mdl')
         assert_frames_close(output, canon, rtol=rtol)
 
-    @unittest.skip('working on it #217')
     def test_sample_if_true(self):
-        # issue https://github.com/JamesPHoughton/pysd/issues/217
         output, canon = runner('test-models/tests/sample_if_true/test_sample_if_true.mdl')
         assert_frames_close(output, canon, rtol=rtol)
 
@@ -379,6 +377,11 @@ def test_subscript_selection(self):
         output, canon = runner('test-models/tests/subscript_selection/subscript_selection.mdl')
         assert_frames_close(output, canon, rtol=rtol)
 
+    def test_subscript_numeric_range(self):
+        from .test_utils import runner, assert_frames_close
+        output, canon = runner('test-models/tests/subscript_numeric_range/test_subscript_numeric_range.mdl')
+        assert_frames_close(output, canon, rtol=rtol)
+
     def test_subscript_subranges(self):
         output, canon = runner('test-models/tests/subscript_subranges/test_subscript_subrange.mdl')
         assert_frames_close(output, canon, rtol=rtol)