diff --git a/src/cyclebane/graph.py b/src/cyclebane/graph.py index 78d5da0..d76cafe 100644 --- a/src/cyclebane/graph.py +++ b/src/cyclebane/graph.py @@ -20,6 +20,24 @@ def _get_unique_sink(graph: nx.DiGraph) -> Hashable: return sink_nodes[0] +def _labeled_key(graph: nx.DiGraph, key: Hashable, match_index: Hashable) -> Hashable: + """Find the node for ``key`` in a labeled graph, by original name if mapped.""" + if key in graph: + return key + matches = [ + node + for node in graph.nodes + if isinstance(node, MappedNode) + and node.name == key + and match_index in node.indices + ] + if len(matches) == 0: + raise KeyError(f"Node '{key}' does not exist in the graph.") + if len(matches) > 1: + raise KeyError(f"Node '{key}' is ambiguous. Found {matches}.") + return matches[0] + + def _get_new_node_name(graph: nx.DiGraph) -> str: while True: name = str(uuid4()) @@ -106,10 +124,16 @@ class MappedNode: indices: tuple[IndexName, ...] -def _node_with_indices(node: Hashable, indices: tuple[IndexName, ...]) -> MappedNode: - if isinstance(node, MappedNode): - return MappedNode(name=node.name, indices=indices + node.indices) - return MappedNode(name=node, indices=indices) +@dataclass(frozen=True, slots=True) +class _ReduceSpec: + """Records what a reduce node consumes; applied when deriving node indices. + + The index names to drop are resolved when ``reduce`` is called, so that + indices added by later ``map`` calls flow through the reduce node. + """ + + drop: frozenset[IndexName] + extra_index_name: None | IndexName def _node_name(node: Hashable) -> Hashable: @@ -124,17 +148,6 @@ def _node_indices(node: Hashable) -> tuple[IndexName, ...]: return () -def _find_successors( - graph: nx.DiGraph, *, root_nodes: tuple[Hashable] -) -> set[Hashable]: - successors = set() - for root in root_nodes: - if graph.in_degree(root) > 0: - raise ValueError(f"Mapped node '{root}' is not a source node") - successors.update(nx.descendants(graph, source=root) | {root}) - return successors - - def _rename_successors( graph: nx.DiGraph, *, successors: Iterable[Hashable], index: IndexValues ) -> nx.DiGraph: @@ -175,7 +188,11 @@ def __getitem__(self, key: int | slice) -> Graph: for name, col in self.graph._node_values.items() } ) - return Graph(self.graph.graph, node_values=node_values) + return Graph( + self.graph.graph, + node_values=node_values, + reductions=self.graph._reductions, + ) MappingToArrayLike = Any # dict[str, Numpy|DataArray], DataFrame, etc. @@ -207,7 +224,13 @@ class Graph: objects at nodes with multiple predecessors. """ - def __init__(self, graph: nx.DiGraph, *, node_values: NodeValues | None = None): + def __init__( + self, + graph: nx.DiGraph, + *, + node_values: NodeValues | None = None, + reductions: dict[Hashable, _ReduceSpec] | None = None, + ): """ Initialize a graph from a directed NetworkX graph. @@ -219,12 +242,19 @@ def __init__(self, graph: nx.DiGraph, *, node_values: NodeValues | None = None): A mapping from source node names to array-like objects. The implementation assumes that the graph has been setup correctly. Do not use this argument unless you know what you are doing. + reductions: + A mapping from reduce-node names to reduce specs. Internal, do not use. """ self.graph = graph self._node_values = node_values or NodeValues({}) + self._reductions = dict(reductions or {}) def copy(self) -> Graph: - return Graph(self.graph.copy(), node_values=self._node_values) + return Graph( + self.graph.copy(), + node_values=self._node_values, + reductions=self._reductions, + ) @property def index_names(self) -> tuple[IndexName, ...]: @@ -240,9 +270,10 @@ def map(self, node_values: MappingToArrayLike) -> Graph: """ Map the graph over the given values by associating source nodes with values. - All successors of the mapped source nodes are replaced with new nodes, one for - each index value. The value is set as an attribute on the new source nodes - (but not their successors). + The mapped source nodes and their successors gain an index (dimension). + This only records the values and indices; nodes are spelled out into one + copy per index value in :py:meth:`to_networkx`, with values set as an + attribute on the source-node copies. Parameters ---------- @@ -262,18 +293,97 @@ def map(self, node_values: MappingToArrayLike) -> Graph: graph = self.graph.copy() graph.add_nodes_from(new_values) - successors = _find_successors(graph, root_nodes=new_values) - name_mapping: dict[Hashable, MappedNode] = {} - for node in successors: - name_mapping[node] = _node_with_indices(node, tuple(new_values.indices)) + for root in new_values: + if graph.in_degree(root) > 0: + raise ValueError(f"Mapped node '{root}' is not a source node") + # Note that the graph is not relabeled: which nodes carry which indices is + # derived from the mapped roots and reachability in _derive_indices, at the + # time it is needed. This keeps node names stable and makes `map` commute + # with adding branches to the graph. return Graph( - nx.relabel_nodes(graph, name_mapping), + graph, node_values=self._node_values.merge(new_values), + reductions=self._reductions, ) + def _derive_indices(self) -> dict[Hashable, tuple[IndexName, ...]]: + """Derive the indices carried by each node from mapped roots and reduces. + + A node carries the indices of the mapped roots that reach it, minus what + reduce nodes on the way consume. The per-node index order matches what + incremental relabeling used to produce: indices of later ``map`` calls + first, order within one call preserved; groupby's extra index appended + last. Nodes carrying no indices are absent from the result. + """ + if not self._node_values and not self._reductions: + # Fast path; also keeps graphs with cycles (which some users allow + # until compute time) working as long as nothing is mapped. + return {} + graph = self.graph + root_blocks = { + root: arr.index_names + for root, arr in self._node_values.items() + if arr.get_grouping() is None and root in graph + } + priority: dict[IndexName, tuple[int, int]] = {} + for i, block in enumerate(root_blocks.values()): + for pos, name in enumerate(block): + priority.setdefault(name, (-i, pos)) + + names: dict[Hashable, set[IndexName]] = {} + extras: dict[Hashable, tuple[IndexName, ...]] = {} + result: dict[Hashable, tuple[IndexName, ...]] = {} + for node in nx.topological_sort(graph): + current: set[IndexName] = set() + extra: tuple[IndexName, ...] = () + for pred in graph.predecessors(node): + current |= names[pred] + for name in extras[pred]: + if name not in extra: + extra = (*extra, name) + if (root_block := root_blocks.get(node)) is not None: + current |= set(root_block) + if (spec := self._reductions.get(node)) is not None: + current -= spec.drop + extra = tuple(name for name in extra if name not in spec.drop) + if spec.extra_index_name is not None: + extra = (*extra, spec.extra_index_name) + names[node] = current + extras[node] = extra + if full := tuple(sorted(current, key=priority.__getitem__)) + extra: + result[node] = full + return result + + def _labeled_graph(self) -> nx.DiGraph: + """Return a copy of the graph with index-carrying nodes relabeled as + :py:class:`MappedNode`, the representation :py:meth:`to_networkx` works on.""" + mapping = { + node: MappedNode(name=node, indices=indices) + for node, indices in self._derive_indices().items() + } + return nx.relabel_nodes(self.graph, mapping, copy=True) + + def node_indices(self, key: Hashable) -> tuple[IndexName, ...]: + """Return the index names carried by the given node, () if unmapped.""" + return self._derive_indices().get(key, ()) + + @property + def value_keys(self) -> tuple[Hashable, ...]: + """Names of the nodes that have associated values. + + Contains the mapped source nodes, and the reduce nodes of groupby + operations (which store the grouping). + """ + return tuple(self._node_values) + def groupby(self, node: Hashable) -> GroupbyGraph: - return GroupbyGraph(self.graph, node_values=self._node_values, node=node) + return GroupbyGraph( + self.graph, + node_values=self._node_values, + node=node, + reductions=self._reductions, + ) def reduce( self, @@ -293,8 +403,8 @@ def reduce( Parameters ---------- key: - The name of the source node to reduce. This is the original name prior to - mapping. If not given, tries to find a unique sink node. + The name of the node to reduce. If not given, tries to find a unique + sink node. index: The name of the index to reduce over. Only one of index and axis can be given. @@ -311,55 +421,33 @@ def reduce( attrs = attrs or {} if index is not None and axis is not None: raise ValueError('Only one of index and axis can be given') - key = self._from_orig_key(key) - indices = _node_indices(key) - if index is not None and index not in indices: - raise ValueError(f"Node '{key}' does not have index '{index}'.") - # TODO We can support indexing from the back in the future. - if axis is not None and (axis < 0 or axis >= len(indices)): - raise ValueError(f"Node '{key}' does not have axis '{axis}'.") + if key not in self.graph: + raise KeyError(f"Node '{key}' does not exist in the graph.") + # Resolve what is reduced into concrete index names now; indices added + # by later `map` calls are unaffected and flow through the reduce node. + indices = self.node_indices(key) if index is not None: - new_index = tuple(value for value in indices if value != index) + if index not in indices: + raise ValueError(f"Node '{key}' does not have index '{index}'.") + drop = frozenset({index}) elif axis is not None: - # TODO Should axis refer to axes of graph, or the node? - new_index = tuple(value for i, value in enumerate(indices) if i != axis) + # TODO We can support indexing from the back in the future. + if axis < 0 or axis >= len(indices): + raise ValueError(f"Node '{key}' does not have axis '{axis}'.") + drop = frozenset({indices[axis]}) else: - new_index = None - if _extra_index_name is not None: - if new_index is None: - new_index = (_extra_index_name,) - else: - new_index = (*new_index, _extra_index_name) + drop = frozenset(indices) + if name in self.graph: raise ValueError(f"Node '{name}' already exists in the graph.") graph = self.graph.copy() - name = MappedNode(name=name, indices=new_index) if new_index else name graph.add_node(name, **attrs) graph.add_edge(key, name) - return Graph(graph, node_values=self._node_values) - - def _from_orig_key( - self, key: Hashable, match_index: None | Hashable = None - ) -> Hashable: - # Graph.map relabels nodes to include index names, which can be inconvenient - # for the user. Is this convenience of finding the node by its original name - # worth the complexity and a good idea? - if key not in self.graph: - matches = [ - node - for node in self.graph.nodes - if isinstance(node, MappedNode) and node.name == key - ] - if match_index is not None: - matches = [node for node in matches if match_index in node.indices] - if len(matches) == 0: - raise KeyError(f"Node '{key}' does not exist in the graph.") - if len(matches) > 1: - raise KeyError(f"Node '{key}' is ambiguous. Found {matches}.") - return matches[0] - return key + reductions = dict(self._reductions) + reductions[name] = _ReduceSpec(drop=drop, extra_index_name=_extra_index_name) + return Graph(graph, node_values=self._node_values, reductions=reductions) def by_position(self, index_name: IndexName) -> PositionalIndexer: return PositionalIndexer(self, index_name) @@ -374,7 +462,9 @@ def to_networkx(self, value_attr: str = 'value') -> nx.DiGraph: value_attr: The name of the attribute on nodes that holds the array-like object. """ - graph = self.graph.copy() + # Nodes carrying indices are stored under their plain names; relabel them + # as MappedNode based on the derived indices before spelling out. + graph = self._labeled_graph() # Maintain a list of actual node values, without groupings, since we only want # to set the former (user-provided) on (input) nodes. @@ -389,7 +479,7 @@ def to_networkx(self, value_attr: str = 'value') -> nx.DiGraph: for key, values in self._node_values.items(): if (grouping := values.get_grouping()) is not None: del node_values[key] - key = self._from_orig_key(key, match_index=grouping.group_index_name) + key = _labeled_key(graph, key, match_index=grouping.group_index_name) # Note there should be only a single predecessor for the grouping node. groupby_graph = graph.subgraph([*graph.predecessors(key), key]).copy() # Remove edges, or the loop for the regular map/reduce will add @@ -449,15 +539,16 @@ def __getitem__(self, key: Hashable | slice) -> Graph: """ if isinstance(key, slice): raise NotImplementedError('Only single nodes are supported ') - key = self._from_orig_key(key) + if key not in self.graph: + raise KeyError(f"Node '{key}' does not exist in the graph.") ancestors = nx.ancestors(self.graph, key) ancestors.add(key) # Drop all node values that are not in the branch - mapped = {a.name for a in ancestors if isinstance(a, MappedNode)} - keep_values = [key for key in self._node_values.keys() if key in mapped] + keep_values = [key for key in self._node_values.keys() if key in ancestors] return Graph( self.graph.subgraph(ancestors), node_values=self._node_values.get_columns(keep_values), + reductions={k: v for k, v in self._reductions.items() if k in ancestors}, ) def __delitem__(self, key: Hashable | slice) -> None: @@ -466,15 +557,16 @@ def __delitem__(self, key: Hashable | slice) -> None: """ if isinstance(key, slice): raise NotImplementedError('Only single nodes are supported ') - key = self._from_orig_key(key) - if isinstance(key, MappedNode): - # Not clear what to do in this case, as it would leave a lot of MappedNodes - # without a source that could provide data. + if key not in self.graph: + raise KeyError(f"Node '{key}' does not exist in the graph.") + if self.node_indices(key): + # Not clear what to do in this case, as it would leave a lot of mapped + # nodes without a source that could provide data. raise ValueError('Cannot delete mapped node.') graph = _remove_ancestors(self.graph, key) - mapped = {node.name for node in graph if isinstance(node, MappedNode)} - keep_values = [key for key in self._node_values.keys() if key in mapped] + keep_values = [key for key in self._node_values.keys() if key in graph] self._node_values = self._node_values.get_columns(keep_values) + self._reductions = {k: v for k, v in self._reductions.items() if k in graph} self.graph = graph def __setitem__(self, branch: Hashable | slice, other: Graph) -> None: @@ -492,17 +584,27 @@ def __setitem__(self, branch: Hashable | slice, other: Graph) -> None: raise TypeError(f'Expected {Graph}, got {type(other)}') new_branch = other.graph sink = _get_unique_sink(new_branch) - try: - # When setting at a MappedNode, allow using underlying node name for - # convenience. - branch = self._from_orig_key(branch) - except KeyError: - pass - if isinstance(sink, MappedNode) != isinstance(branch, MappedNode): + # Replacing an existing branch must not change whether it is mapped, as + # this would silently change the indices of its dependents. Setting a + # new branch is fine either way; its indices follow from derivation. + if branch in self.graph and bool(other.node_indices(sink)) != bool( + self.node_indices(branch) + ): raise NotImplementedError( 'Trying to set mapped node on non-mapped node (or vice versa) is not ' 'possible in __setitem__' ) + if branch in new_branch and branch != sink: + # Renaming the sink to the branch name would silently merge it with + # the like-named node inside the new branch. This typically means + # the new branch computes the branch node from itself (e.g. a + # reduction of a mapped branch assigned back to the same name); + # rename the node inside the new branch to make this well-defined. + raise ValueError( + f"Cannot set branch '{branch}': the new branch already contains " + f"a node of that name. Use a distinct name for the node inside " + "the new branch." + ) new_branch = nx.relabel_nodes(new_branch, {sink: branch}) if branch in self.graph: graph = _remove_ancestors(self.graph, branch) @@ -525,14 +627,22 @@ def __setitem__(self, branch: Hashable | slice, other: Graph) -> None: # Delay setting graph until we know no step fails self._node_values = self._node_values.merge(other._node_values) + reductions = dict(self._reductions) + # A reduce spec of a replaced branch node must not survive replacement. + reductions.pop(branch, None) + reductions.update(other._reductions) + if sink in reductions and sink != branch: + # The sink was renamed to the branch name above. + reductions[branch] = reductions.pop(sink) + self._reductions = reductions # Ensure we preserve the node values of the branch, if it exists. This step is # necessary since __setitem__ effectively renames the sink node of the input # graph to the branch name. - if _node_name(sink) in self._node_values: - node_values = self._node_values[_node_name(sink)] - del self._node_values[_node_name(sink)] - self._node_values[_node_name(branch)] = node_values + if sink in self._node_values: + node_values = self._node_values[sink] + del self._node_values[sink] + self._node_values[branch] = node_values self.graph = graph @@ -547,9 +657,16 @@ class GroupbyGraph: """ # TODO Should we support a custom new dim name here, instead of using `node`? - def __init__(self, graph: nx.DiGraph, node_values: NodeValues, node: Hashable): + def __init__( + self, + graph: nx.DiGraph, + node_values: NodeValues, + node: Hashable, + reductions: dict[Hashable, _ReduceSpec] | None = None, + ): self._graph = graph self._node_values = node_values + self._reductions = dict(reductions or {}) values_to_group_by = node_values[node] self._group_index_name = node self._index_name = values_to_group_by.index_names[0] @@ -569,8 +686,8 @@ def reduce( Parameters ---------- key: - The name of the source node to reduce. This is the original name prior to - mapping. If not given, tries to find a unique sink node. + The name of the node to reduce. If not given, tries to find a unique + sink node. name: The name of the new node. If not given, a unique name is generated. attrs: @@ -582,7 +699,7 @@ def reduce( # e.g., __getitem__, which needs to decided what subset of node values to keep # when returning a subgraph. node_values = self._node_values.merge({name: self._groups}) - graph = Graph(self._graph, node_values=node_values) + graph = Graph(self._graph, node_values=node_values, reductions=self._reductions) return graph.reduce( key=key, index=self._index_name, diff --git a/tests/graph_test.py b/tests/graph_test.py index b25448c..f8eb3b9 100644 --- a/tests/graph_test.py +++ b/tests/graph_test.py @@ -1102,3 +1102,38 @@ def test_node_attrs_are_preserved_in_map() -> None: assert result.nodes[idx('b', 0)] == {'attr': 22} assert result.nodes[idx('b', 1)] == {'attr': 22} assert result.nodes[idx('b', 2)] == {'attr': 22} + + +def test_map_after_axis_reduce_leaves_new_index_on_reduce_node() -> None: + g = nx.DiGraph() + g.add_edge('a', 'c') + g.add_edge('b', 'c') + graph = cb.Graph(g).map({'a': [1, 2]}) + graph = graph.reduce('c', axis=0, name='r') + assert graph.node_indices('r') == () + # The reduce consumed a's index; b's index, added later, flows through. + graph = graph.map({'b': [10, 20, 30]}) + assert graph.node_indices('r') == ('dim_1',) + + +def test_map_after_full_reduce_leaves_new_index_on_reduce_node() -> None: + g = nx.DiGraph() + g.add_edge('a', 'c') + g.add_edge('b', 'c') + graph = cb.Graph(g).map({'a': [1, 2]}) + graph = graph.reduce('c', name='r') + graph = graph.map({'b': [10, 20, 30]}) + assert graph.node_indices('r') == ('dim_1',) + + +def test_setitem_replacing_reduce_node_discards_its_reduction() -> None: + g = nx.DiGraph() + g.add_edges_from([('a', 'c'), ('b', 'c')]) + graph = cb.Graph(g) + graph = graph.map(pd.DataFrame({'a': [1, 2]}).rename_axis('x')) + graph = graph.map(pd.DataFrame({'b': [3, 4]}).rename_axis('y')) + graph = graph.reduce('c', index='x', name='r') + assert graph.node_indices('r') == ('y',) + # Replace the reduce node by the plain mapped branch; the reduction is gone. + graph['r'] = graph['c'] + assert graph.node_indices('r') == ('y', 'x') diff --git a/tests/groupby_test.py b/tests/groupby_test.py index c1d8bc3..fbe24cc 100644 --- a/tests/groupby_test.py +++ b/tests/groupby_test.py @@ -441,8 +441,9 @@ def test_groupby_with_branch_operations(self) -> None: graph1 = cb.Graph(g1).map(df) graph2 = cb.Graph(g2).map(df) - # Combine graphs - graph1['c'] = graph2['d'] + # Combine graphs. Note that using 'c' as the target name would raise, + # since the branch contains its mapped source node 'c'. + graph1['d'] = graph2['d'] # Groupby on combined graph grouped = graph1.groupby('param').reduce('b', name='reduced')