Repository: ModellingWebLab/cellmlmanip

Issue:
Replace a variable's equation
Given a variable, we need a method to set/replace its RHS

From FC, I can read equations (with variables specified as ontology terms) and units (using an FC unit syntax).

In a two-way communication approach, I could then:

- ask the cellmlmanip model for the symbols for each term
- ask the cellmlmanip model for the unit objects corresponding to my units
- pass in an equation with symbols pointing to model variables and literals as dummies containing the model's unit vars

Buuuut, that would mean fc talking to cellmlmanip through cg, so would probably prefer a one-way approach

- give cellmlmanip a sympy equation with variables as ontology terms and units in some syntax that `pint` understands?
- cellmlmanip then replaces the ontology terms and units

So then we'd need to write some mini-spec for how cellmlmanip expects variable symbols and units?


Discussion:

Patch:
diff --git a/cellmlmanip/model.py b/cellmlmanip/model.py
index 5e014f6..cf611dd 100644
--- a/cellmlmanip/model.py
+++ b/cellmlmanip/model.py
@@ -112,9 +112,11 @@ class Model(object):
 
         :return: A :class:`VariableDummy` object.
         """
+        # Check for clashes
         if name in self._name_to_symbol:
             raise ValueError('Variable %s already exists.' % name)
 
+        # Add variable
         self._name_to_symbol[name] = var = VariableDummy(
             name=name,
             units=self.units.get_quantity(units),
@@ -125,6 +127,9 @@ class Model(object):
             cmeta_id=cmeta_id,
         )
 
+        # Invalidate cached graphs
+        self._invalidate_cache()
+
         return var
 
     def connect_variables(self, source_name: str, target_name: str):
@@ -141,45 +146,48 @@ class Model(object):
         source = self._name_to_symbol[source_name]
         target = self._name_to_symbol[target_name]
 
-        # If the source variable has already been assigned a final symbol
-        if source.assigned_to:
-
-            if target.assigned_to:
-                raise ValueError('Target already assigned to %s before assignment to %s' %
-                                 (target.assigned_to, source.assigned_to))
-
-            # If source/target variable is in the same unit
-            if source.units == target.units:
-                # Direct substitution is possible
-                target.assigned_to = source.assigned_to
-                # everywhere the target variable is used, replace with source variable
-                for index, equation in enumerate(self.equations):
-                    self.equations[index] = equation.xreplace({target: source.assigned_to})
-            # Otherwise, this connection requires a conversion
-            else:
-                # Get the scaling factor required to convert source units to target units
-                factor = self.units.convert_to(1 * source.units, target.units).magnitude
+        # If the source variable has not been assigned a symbol, we can't make this connection
+        if not source.assigned_to:
+            logger.info('The source variable has not been assigned to a symbol '
+                        '(i.e. expecting a connection): %s ⟶ %s',
+                        target.name, source.name)
+            return False
+
+        # If target is already assigned this is an error
+        if target.assigned_to:
+            raise ValueError('Target already assigned to %s before assignment to %s' %
+                             (target.assigned_to, source.assigned_to))
+
+        # If source/target variable is in the same unit
+        if source.units == target.units:
+            # Direct substitution is possible
+            target.assigned_to = source.assigned_to
+            # everywhere the target variable is used, replace with source variable
+            for index, equation in enumerate(self.equations):
+                self.equations[index] = equation.xreplace({target: source.assigned_to})
+
+        # Otherwise, this connection requires a conversion
+        else:
+            # Get the scaling factor required to convert source units to target units
+            factor = self.units.convert_to(1 * source.units, target.units).magnitude
 
-                # Dummy to represent this factor in equations, having units for conversion
-                factor_dummy = self.add_number(factor, str(target.units / source.units))
+            # Dummy to represent this factor in equations, having units for conversion
+            factor_dummy = self.add_number(factor, str(target.units / source.units))
 
-                # Add an equations making the connection with the required conversion
-                self.equations.append(sympy.Eq(target, source.assigned_to * factor_dummy))
+            # Add an equations making the connection with the required conversion
+            self.equations.append(sympy.Eq(target, source.assigned_to * factor_dummy))
 
-                logger.info('Connection req. unit conversion: %s', self.equations[-1])
+            logger.info('Connection req. unit conversion: %s', self.equations[-1])
 
-                # The assigned symbol for this variable is itself
-                target.assigned_to = target
+            # The assigned symbol for this variable is itself
+            target.assigned_to = target
 
-            logger.debug('Updated target: %s', target)
+        logger.debug('Updated target: %s', target)
 
-            return True
+        # Invalidate cached graphs
+        self._invalidate_cache()
 
-        # The source variable has not been assigned a symbol, so we can't make this connection
-        logger.info('The source variable has not been assigned to a symbol '
-                    '(i.e. expecting a connection): %s ⟶ %s',
-                    target.name, source.name)
-        return False
+        return True
 
     def add_rdf(self, rdf: str):
         """ Takes an RDF string and stores it in the model's RDF graph. """
@@ -569,36 +577,16 @@ class Model(object):
                 symbols |= self.find_symbols_and_derivatives(expr.args)
         return symbols
 
-    def set_equation(self, lhs, rhs):
+    def remove_equation(self, equation):
         """
-        Adds an equation defining the variable named in ``lhs``, or replaces an existing one.
-
-        As with :meth:`add_equation()` the LHS must be either a variable symbol or a derivative, and all numbers and
-        variable symbols used in ``lhs`` and ``rhs`` must have been obtained from this model, e.g. via
-        :meth:`add_number()`, :meth:`add_variable()`, or :meth:`get_symbol_by_ontology_term()`.
+        Removes an equation from the model.
 
-        :param lhs: An LHS expression (either a symbol or a derivative).
-        :param rhs: The new RHS expression for this variable.
+        :param equation: The equation to remove.
         """
-        # Get variable symbol named in the lhs
-        lhs_symbol = lhs
-        if lhs_symbol.is_Derivative:
-            lhs_symbol = lhs_symbol.free_symbols.pop()
-        assert isinstance(lhs_symbol, VariableDummy)
-
-        # Check if the variable named in the lhs already has an equation
-        i_existing = None
-        for i, eq in enumerate(self.equations):
-            symbol = eq.lhs.free_symbols.pop() if eq.lhs.is_Derivative else eq.lhs
-            if symbol == lhs_symbol:
-                i_existing = i
-                break
-
-        # Add or replace equation
-        if i_existing is None:
-            self.equations.append(sympy.Eq(lhs, rhs))
-        else:
-            self.equations[i_existing] = sympy.Eq(lhs, rhs)
+        try:
+            self.equations.remove(equation)
+        except ValueError:
+            raise KeyError('Equation not found in model ' + str(equation))
 
         # Invalidate cached equation graphs
         self._invalidate_cache()


Test patch:
diff --git a/tests/test_model_functions.py b/tests/test_model_functions.py
index 0cff687..3a98f46 100644
--- a/tests/test_model_functions.py
+++ b/tests/test_model_functions.py
@@ -359,35 +359,8 @@ class TestModelFunctions():
         assert eqn[2].lhs == symbol2
         assert eqn[2].rhs == sp.Add(symbol, symbol1)
 
-    def test_set_equation(self, local_model):
-        """ Tests the Model.set_equation method.
-        """
-        model = local_model
-        assert len(model.equations) == 1
-        # so we are adding
-        # newvar2 = newvar + newvar1
-        # but need to also add newvar1 = 2; newvar = 2 in order or the graph to resolve correctly
-        model.add_variable(name='newvar', units='mV')
-        symbol = model.get_symbol_by_name('newvar')
-        model.add_variable(name='newvar1', units='mV')
-        symbol1 = model.get_symbol_by_name('newvar1')
-        model.add_variable(name='newvar2', units='mV')
-        symbol2 = model.get_symbol_by_name('newvar2')
-        model.set_equation(symbol, 2.0)
-        model.set_equation(symbol1, 2.0)
-        model.set_equation(symbol2, sp.Add(symbol, symbol1))
-        assert len(model.equations) == 4
-        eqn = model.get_equations_for([symbol2])
-        assert len(eqn) == 3
-        assert eqn[0].lhs == symbol
-        assert eqn[0].rhs == 2.0
-        assert eqn[1].lhs == symbol1
-        assert eqn[1].rhs == 2.0
-        assert eqn[2].lhs == symbol2
-        assert eqn[2].rhs == sp.Add(symbol, symbol1)
-
-    def test_set_equation2(self, local_hh_model):
-        """ Tests replacing an equation in a model. """
+    def test_remove_equation(self, local_hh_model):
+        """ Tests the Model.remove_equation method. """
 
         model = local_hh_model
         # Get model, assert that V is a state variable
@@ -395,8 +368,11 @@ class TestModelFunctions():
         assert v.type == 'state'
 
         # Now clamp it to -80mV
-        rhs = model.add_number(-80, str(v.units))
-        model.set_equation(v, rhs)
+        t = model.get_symbol_by_ontology_term(shared.OXMETA, 'time')
+        equation = model.graph.nodes[sp.Derivative(v, t)]['equation']
+        model.remove_equation(equation)
+        equation = sp.Eq(v, model.add_number(-80, str(v.units)))
+        model.add_equation(equation)
 
         # Check that V is no longer a state
         v = model.get_symbol_by_ontology_term(shared.OXMETA, 'membrane_voltage')
@@ -406,24 +382,23 @@ class TestModelFunctions():
         # See: https://github.com/ModellingWebLab/cellmlmanip/issues/133
 
         # Now make V a state again
-        t = model.get_symbol_by_ontology_term(shared.OXMETA, 'time')
-        lhs = sp.Derivative(v, t)
         dvdt_units = 'unlikely_unit_name'
         model.add_unit(dvdt_units, [
             {'units': str(v.units)},
             {'units': str(t.units), 'exponent': -1},
         ])
-        rhs = model.add_number(0, dvdt_units)
-        model.set_equation(lhs, rhs)
+        model.remove_equation(equation)
+        equation = sp.Eq(sp.Derivative(v, t), model.add_number(0, dvdt_units))
+        model.add_equation(equation)
 
         # Check that V is a state again
         v = model.get_symbol_by_ontology_term(shared.OXMETA, 'membrane_voltage')
         assert v.type == 'state'
 
-        # Set equation for a newly created variable
-        lhs = model.add_variable(name='an_incredibly_unlikely_variable_name', units=str(v.units))
-        rhs = model.add_number(12, str(v.units))
-        model.set_equation(lhs, rhs)
+        # Test removing non-existing equation
+        equation = sp.Eq(sp.Derivative(v, t), model.add_number(5, dvdt_units))
+        with pytest.raises(KeyError, match='Equation not found'):
+            model.remove_equation(equation)
 
     def test_add_number(self, local_model):
         """ Tests the Model.add_number method. """