Fixed leaf_nodes computation (#3645)

* Fixed leaf_nodes computation

* Refactored get_named_nodes_and_relations to make it consistent with theano naming

* Added test for #3643

* Added release notes

* Fixed float32 error

* Applied changes suggested by @rpgoldman

* Fixed ignored nodes bug

* Modify docstrings for numpy style.

Noticed @junpenglao's comment and modified for numpy formatting.

* Changed leaves to leaf_dict as per @junpeglao's suggestion

Co-authored-by: rpgoldman <rpgoldman@goldman-tribe.org>

RELEASE-NOTES.md

@@ -42,6 +42,8 @@
 - Added `theano.gof.graph.Constant` to type checks done in `_draw_value` (fixes issue [3595](https://github.com/pymc-devs/pymc3/issues/3595))
 - `HalfNormal` did not used to work properly in `draw_values`, `sample_prior_predictive`, or `sample_posterior_predictive` (fixes issue [3686](https://github.com/pymc-devs/pymc3/pull/3686))
 - Random variable transforms were inadvertently left out of the API documentation. Added them. (See PR [3690](https://github.com/pymc-devs/pymc3/pull/3690)).
+- Refactored `pymc3.model.get_named_nodes_and_relations` to use the ancestors and descendents, in a way that is consistent with `theano`'s naming convention.
+- Changed the way in which `pymc3.model.get_named_nodes_and_relations` computes nodes without ancestors to make it robust to changes in var_name orderings (issue [#3643](https://github.com/pymc-devs/pymc3/issues/3643))
 
 ## PyMC3 3.7 (May 29 2019)
 

pymc3/distributions/distribution.py

@@ -20,7 +20,7 @@
 import theano
 from ..memoize import memoize
 from ..model import (
-    Model, get_named_nodes_and_relations, FreeRV,
+    Model, build_named_node_tree, FreeRV,
     ObservedRV, MultiObservedRV, ContextMeta
 )
 from ..vartypes import string_types, theano_constant
@@ -569,31 +569,19 @@ def draw_values(params, point=None, size=None):
         # Distribution parameters may be nodes which have named node-inputs
         # specified in the point. Need to find the node-inputs, their
         # parents and children to replace them.
-        leaf_nodes = {}
-        named_nodes_parents = {}
-        named_nodes_children = {}
-        for _, param in symbolic_params:
-            if hasattr(param, 'name'):
-                # Get the named nodes under the `param` node
-                nn, nnp, nnc = get_named_nodes_and_relations(param)
-                leaf_nodes.update(nn)
-                # Update the discovered parental relationships
-                for k in nnp.keys():
-                    if k not in named_nodes_parents.keys():
-                        named_nodes_parents[k] = nnp[k]
-                    else:
-                        named_nodes_parents[k].update(nnp[k])
-                # Update the discovered child relationships
-                for k in nnc.keys():
-                    if k not in named_nodes_children.keys():
-                        named_nodes_children[k] = nnc[k]
-                    else:
-                        named_nodes_children[k].update(nnc[k])
+        leaf_nodes, named_nodes_descendents, named_nodes_ancestors = (
+            build_named_node_tree(
+                (
+                    param for _, param in symbolic_params
+                    if hasattr(param, "name")
+                )
+            )
+        )
 
         # Init givens and the stack of nodes to try to `_draw_value` from
         givens = {p.name: (p, v) for (p, size), v in drawn.items()
                   if getattr(p, 'name', None) is not None}
-        stack = list(leaf_nodes.values())  # A queue would be more appropriate
+        stack = list(leaf_nodes.values())
         while stack:
             next_ = stack.pop(0)
             if (next_, size) in drawn:
@@ -614,7 +602,7 @@ def draw_values(params, point=None, size=None):
                 # of TensorConstants or SharedVariables, we must add them
                 # to the stack or risk evaluating deterministics with the
                 # wrong values (issue #3354)
-                stack.extend([node for node in named_nodes_parents[next_]
+                stack.extend([node for node in named_nodes_descendents[next_]
                               if isinstance(node, (ObservedRV,
                                                    MultiObservedRV))
                               and (node, size) not in drawn])
@@ -623,7 +611,7 @@ def draw_values(params, point=None, size=None):
                 # If the node does not have a givens value, try to draw it.
                 # The named node's children givens values must also be taken
                 # into account.
-                children = named_nodes_children[next_]
+                children = named_nodes_ancestors[next_]
                 temp_givens = [givens[k] for k in givens if k in children]
                 try:
                     # This may fail for autotransformed RVs, which don't
@@ -638,7 +626,7 @@ def draw_values(params, point=None, size=None):
                     # The node failed, so we must add the node's parents to
                     # the stack of nodes to try to draw from. We exclude the
                     # nodes in the `params` list.
-                    stack.extend([node for node in named_nodes_parents[next_]
+                    stack.extend([node for node in named_nodes_descendents[next_]
                                   if node is not None and
                                   (node, size) not in drawn])
 
@@ -662,8 +650,8 @@ def draw_values(params, point=None, size=None):
                         # This may set values for certain nodes in the drawn
                         # dictionary, but they don't get added to the givens
                         # dictionary. Here, we try to fix that.
-                        if param in named_nodes_children:
-                            for node in named_nodes_children[param]:
+                        if param in named_nodes_ancestors:
+                            for node in named_nodes_ancestors[param]:
                                 if (
                                     node.name not in givens and
                                     (node, size) in drawn

pymc3/model.py

@@ -105,6 +105,7 @@ def incorporate_methods(source, destination, methods,
         else:
             setattr(destination, method, None)
 
+
 def get_named_nodes_and_relations(graph):
     """Get the named nodes in a theano graph (i.e., nodes whose name
     attribute is not None) along with their relationships (i.e., the
@@ -113,67 +114,128 @@ def get_named_nodes_and_relations(graph):
 
     Parameters
     ----------
-    graph - a theano node
+    graph: a theano node
 
     Returns:
-    leaf_nodes: A dictionary of name:node pairs, of the named nodes that
-        are also leafs of the graph
-    node_parents: A dictionary of node:set([parents]) pairs. Each key is
-        a theano named node, and the corresponding value is the set of
-        theano named nodes that are parents of the node. These parental
-        relations skip unnamed intermediate nodes.
-    node_children: A dictionary of node:set([children]) pairs. Each key
+    --------
+    leaf_dict: Dict[str, node]
+        A dictionary of name:node pairs, of the named nodes that
+        have no named ancestors in the provided theano graph.
+    descendents: Dict[node, Set[node]]
+        Each key is a theano named node, and the corresponding value
+        is the set of theano named nodes that are descendents with no
+        intervening named nodes in the supplied ``graph``.
+    ancestors: Dict[node, Set[node]]
+        A dictionary of node:set([ancestors]) pairs. Each key
         is a theano named node, and the corresponding value is the set
-        of theano named nodes that are children of the node. These child
-        relations skip unnamed intermediate nodes.
+        of theano named nodes that are ancestors with no intervening named
+        nodes in the supplied ``graph``.
 
     """
+    # We don't enforce distribution parameters to have a name but we may
+    # attempt to get_named_nodes_and_relations from them anyway in
+    # distributions.draw_values. This means that must take care only to add
+    # graph to the ancestors and descendents dictionaries if it has a name.
     if graph.name is not None:
-        node_parents = {graph: set()}
-        node_children = {graph: set()}
+        ancestors = {graph: set()}
+        descendents = {graph: set()}
     else:
-        node_parents = {}
-        node_children = {}
-    return _get_named_nodes_and_relations(graph, None, {}, node_parents, node_children)
-
-def _get_named_nodes_and_relations(graph, parent, leaf_nodes,
-                                        node_parents, node_children):
+        ancestors = {}
+        descendents = {}
+    descendents, ancestors = _get_named_nodes_and_relations(
+        graph, None, ancestors, descendents
+    )
+    leaf_dict = {
+        node.name: node for node, ancestor in ancestors.items()
+        if len(ancestor) == 0
+    }
+    return leaf_dict, descendents, ancestors
+
+
+def _get_named_nodes_and_relations(graph, descendent, descendents, ancestors):
     if getattr(graph, 'owner', None) is None:  # Leaf node
         if graph.name is not None:  # Named leaf node
-            leaf_nodes.update({graph.name: graph})
-            if parent is not None:  # Is None for the root node
+            if descendent is not None:  # Is None for the first node
                 try:
-                    node_parents[graph].add(parent)
+                    descendents[graph].add(descendent)
                 except KeyError:
-                    node_parents[graph] = {parent}
-                node_children[parent].add(graph)
+                    descendents[graph] = {descendent}
+                ancestors[descendent].add(graph)
             else:
-                node_parents[graph] = set()
+                descendents[graph] = set()
             # Flag that the leaf node has no children
-            node_children[graph] = set()
+            ancestors[graph] = set()
     else:  # Intermediate node
         if graph.name is not None:  # Intermediate named node
-            if parent is not None:  # Is only None for the root node
+            if descendent is not None:  # Is only None for the root node
                 try:
-                    node_parents[graph].add(parent)
+                    descendents[graph].add(descendent)
                 except KeyError:
-                    node_parents[graph] = {parent}
-                node_children[parent].add(graph)
+                    descendents[graph] = {descendent}
+                ancestors[descendent].add(graph)
             else:
-                node_parents[graph] = set()
-            # The current node will be set as the parent of the next
+                descendents[graph] = set()
+            # The current node will be set as the descendent of the next
             # nodes only if it is a named node
-            parent = graph
+            descendent = graph
             # Init the nodes children to an empty set
-            node_children[graph] = set()
+            ancestors[graph] = set()
         for i in graph.owner.inputs:
-            temp_nodes, temp_inter, temp_tree = \
-                _get_named_nodes_and_relations(i, parent, leaf_nodes,
-                                               node_parents, node_children)
-            leaf_nodes.update(temp_nodes)
-            node_parents.update(temp_inter)
-            node_children.update(temp_tree)
-    return leaf_nodes, node_parents, node_children
+            temp_desc, temp_ances = _get_named_nodes_and_relations(
+                i, descendent, descendents, ancestors
+            )
+            descendents.update(temp_desc)
+            ancestors.update(temp_ances)
+    return descendents, ancestors
+
+
+def build_named_node_tree(graphs):
+    """Build the combined descence/ancestry tree of named nodes (i.e., nodes
+    whose name attribute is not None) in a list (or iterable) of theano graphs.
+    The relationship tree does not include unnamed intermediate nodes present
+    in the supplied graphs.
+
+    Parameters
+    ----------
+    graphs - iterable of theano graphs
+
+    Returns:
+    --------
+    leaf_dict: Dict[str, node]
+        A dictionary of name:node pairs, of the named nodes that
+        have no named ancestors in the provided theano graphs.
+    descendents: Dict[node, Set[node]]
+        A dictionary of node:set([parents]) pairs. Each key is
+        a theano named node, and the corresponding value is the set of
+        theano named nodes that are descendents with no intervening named
+        nodes in the supplied ``graphs``.
+    ancestors: Dict[node, Set[node]]
+        A dictionary of node:set([ancestors]) pairs. Each key
+        is a theano named node, and the corresponding value is the set
+        of theano named nodes that are ancestors with no intervening named
+        nodes in the supplied ``graphs``.
+
+    """
+    leaf_dict = {}
+    named_nodes_descendents = {}
+    named_nodes_ancestors = {}
+    for graph in graphs:
+        # Get the named nodes under the `param` node
+        nn, nnd, nna = get_named_nodes_and_relations(graph)
+        leaf_dict.update(nn)
+        # Update the discovered parental relationships
+        for k in nnd.keys():
+            if k not in named_nodes_descendents.keys():
+                named_nodes_descendents[k] = nnd[k]
+            else:
+                named_nodes_descendents[k].update(nnd[k])
+        # Update the discovered child relationships
+        for k in nna.keys():
+            if k not in named_nodes_ancestors.keys():
+                named_nodes_ancestors[k] = nna[k]
+            else:
+                named_nodes_ancestors[k].update(nna[k])
+    return leaf_dict, named_nodes_descendents, named_nodes_ancestors
 
 T = TypeVar('T', bound='ContextMeta')
 

pymc3/tests/test_sampling.py

@@ -436,6 +436,32 @@ def test_deterministic_of_observed(self):
             rtol = 1e-5 if theano.config.floatX == "float64" else 1e-3
             assert np.allclose(ppc["in_1"] + ppc["in_2"], ppc["out"], rtol=rtol)
 
+    def test_var_name_order_invariance(self):
+        # Issue #3643 exposed a bug in sample_posterior_predictive, which made
+        # it return incorrect values depending on the order of the supplied
+        # var_names. This tests that sample_posterior_predictive is robust
+        # to different var_names orders.
+        obs_a = theano.shared(pm.theanof.floatX(np.array([10., 20., 30.])))
+        with pm.Model() as m:
+            pm.Normal('mu', 3, 5)
+            a = pm.Normal('a', 20, 10, observed=obs_a)
+            pm.Deterministic('b', a * 2)
+            trace = pm.sample(10)
+
+        np.random.seed(123)
+        var_names = ['b', 'a']
+        ppc1 = pm.sample_posterior_predictive(
+            trace, model=m, var_names=var_names
+        )
+        np.random.seed(123)
+        var_names = ['a', 'b']
+        ppc2 = pm.sample_posterior_predictive(
+            trace, model=m, var_names=var_names
+        )
+        assert np.all(ppc1["a"] == ppc2["a"])
+        assert np.all(ppc1["b"] == ppc2["b"])
+        assert np.allclose(ppc1["b"], (2 * ppc1["a"]))
+
     def test_deterministic_of_observed_modified_interface(self):
         meas_in_1 = pm.theanof.floatX(2 + 4 * np.random.randn(100))
         meas_in_2 = pm.theanof.floatX(5 + 4 * np.random.randn(100))