Skip to content

Commit

Permalink
Implement more node iteration stuff.
Browse files Browse the repository at this point in the history
  • Loading branch information
@kevinhartman
kevinhartman committed May 31, 2024
1 parent a53a0a7 commit 13d18fd
Showing 1 changed file with 155 additions and 86 deletions.
241 changes: 155 additions & 86 deletions crates/circuit/src/dag_circuit.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,35 +39,43 @@ use std::f64::consts::PI;
use std::ffi::c_double;
use std::hash::{Hash, Hasher};

trait IntoUniqueNodes {
trait IntoUnique {
type Output;
fn into_unique(self) -> Self::Output;
fn unique(self) -> Self::Output;
}

struct UniqueNeighbors<'a, E: 'a, Ix: 'a> {
neighbors: Neighbors<'a, E, Ix>,
seen: HashSet<NodeIndex>,
struct UniqueIterator<I, N: Hash + Eq> {
neighbors: I,
seen: HashSet<N>,
}

impl<'a, E: 'a, Ix: 'a> IntoUniqueNodes for Neighbors<'a, E, Ix> {
type Output = UniqueNeighbors<'a, E, Ix>;
impl<I> IntoUnique for I
where
I: Iterator,
I::Item: Hash + Eq + Clone,
{
type Output = UniqueIterator<I, I::Item>;

fn into_unique(self) -> Self::Output {
UniqueNeighbors {
fn unique(self) -> Self::Output {
UniqueIterator {
neighbors: self,
seen: HashSet::new(),
}
}
}

impl<'a, E> Iterator for UniqueNeighbors<'a, E, DefaultIx> {
type Item = NodeIndex;
impl<I> Iterator for UniqueIterator<I, I::Item>
where
I: Iterator,
I::Item: Hash + Eq + Clone,
{
type Item = I::Item;

fn next(&mut self) -> Option<NodeIndex> {
fn next(&mut self) -> Option<Self::Item> {
// First any outgoing edges
while let Some(node) = self.neighbors.next() {
if !self.seen.contains(&node) {
self.seen.insert(node);
self.seen.insert(node.clone());
return Some(node);
}
}
Expand Down Expand Up @@ -745,7 +753,7 @@ impl DAGCircuit {
let child = self
.dag
.neighbors_directed(input_node, Outgoing)
.into_unique()
.unique()
.next()
.ok_or_else(|| {
DAGCircuitError::new_err(format!(
Expand Down Expand Up @@ -1983,25 +1991,46 @@ impl DAGCircuit {
// for source, dest, edge in raw_nodes:
// yield (self._multi_graph[source], self._multi_graph[dest], edge)
//
// def op_nodes(self, op=None, include_directives=True):
// """Get the list of "op" nodes in the dag.
//
// Args:
// op (Type): :class:`qiskit.circuit.Operation` subclass op nodes to
// return. If None, return all op nodes.
// include_directives (bool): include `barrier`, `snapshot` etc.
//
// Returns:
// list[DAGOpNode]: the list of node ids containing the given op.
// """
// nodes = []
// for node in self._multi_graph.nodes():
// if isinstance(node, DAGOpNode):
// if not include_directives and getattr(node.op, "_directive", False):
// continue
// if op is None or isinstance(node.op, op):
// nodes.append(node)
// return nodes
/// Get the list of "op" nodes in the dag.
///
/// Args:
/// op (Type): :class:`qiskit.circuit.Operation` subclass op nodes to
/// return. If None, return all op nodes.
/// include_directives (bool): include `barrier`, `snapshot` etc.
///
/// Returns:
/// list[DAGOpNode]: the list of node ids containing the given op.
#[pyo3(signature=(op=None, include_directives=true))]
fn op_nodes(
&self,
py: Python,
op: Option<&Bound<PyAny>>,
include_directives: bool,
) -> PyResult<Vec<Py<PyAny>>> {
let mut nodes = Vec::new();
for (node, weight) in self.dag.node_references() {
if let NodeType::Operation(ref packed) = weight {
if !include_directives {
let directive = if packed.op.bind(py).hasattr(intern!(py, "_directive"))? {
packed
.op
.bind(py)
.getattr(intern!(py, "_directive"))?
.extract()?
} else {
false
};
if directive {
continue;
}
}
if op.is_none() || packed.op.bind(py).is_instance(op.unwrap())? {
nodes.push(node);
}
}
}
nodes.into_iter().map(|i| self.get_node(py, i)).collect()
}
//
// def gate_nodes(self):
// """Get the list of gate nodes in the dag.
Expand Down Expand Up @@ -2055,7 +2084,7 @@ impl DAGCircuit {
let successors: PyResult<Vec<_>> = self
.dag
.neighbors_directed(NodeIndex::new(node._node_id as usize), Outgoing)
.into_unique()
.unique()
.map(|i| self.get_node(py, i))
.collect();
Ok(PyTuple::new_bound(py, successors?)
Expand All @@ -2070,7 +2099,7 @@ impl DAGCircuit {
let predecessors: PyResult<Vec<_>> = self
.dag
.neighbors_directed(NodeIndex::new(node._node_id as usize), Incoming)
.into_unique()
.unique()
.map(|i| self.get_node(py, i))
.collect();
Ok(PyTuple::new_bound(py, predecessors?)
Expand Down Expand Up @@ -2099,57 +2128,97 @@ impl DAGCircuit {
)
.is_some()
}
//
// def quantum_predecessors(self, node):
// """Returns iterator of the predecessors of a node that are
// connected by a quantum edge as DAGOpNodes and DAGInNodes."""
// return iter(
// self._multi_graph.find_predecessors_by_edge(
// node._node_id, lambda edge_data: isinstance(edge_data, Qubit)
// )
// )
//
// def classical_predecessors(self, node):
// """Returns iterator of the predecessors of a node that are
// connected by a classical edge as DAGOpNodes and DAGInNodes."""
// return iter(
// self._multi_graph.find_predecessors_by_edge(
// node._node_id, lambda edge_data: isinstance(edge_data, Clbit)
// )
// )
//
// def ancestors(self, node):
// """Returns set of the ancestors of a node as DAGOpNodes and DAGInNodes."""
// return {self._multi_graph[x] for x in rx.ancestors(self._multi_graph, node._node_id)}
//
// def descendants(self, node):
// """Returns set of the descendants of a node as DAGOpNodes and DAGOutNodes."""
// return {self._multi_graph[x] for x in rx.descendants(self._multi_graph, node._node_id)}
//
// def bfs_successors(self, node):
// """
// Returns an iterator of tuples of (DAGNode, [DAGNodes]) where the DAGNode is the current node
// and [DAGNode] is its successors in BFS order.
// """
// return iter(rx.bfs_successors(self._multi_graph, node._node_id))
//
// def quantum_successors(self, node):
// """Returns iterator of the successors of a node that are
// connected by a quantum edge as Opnodes and DAGOutNodes."""
// return iter(
// self._multi_graph.find_successors_by_edge(
// node._node_id, lambda edge_data: isinstance(edge_data, Qubit)
// )
// )
//
// def classical_successors(self, node):
// """Returns iterator of the successors of a node that are
// connected by a classical edge as DAGOpNodes and DAGInNodes."""
// return iter(
// self._multi_graph.find_successors_by_edge(
// node._node_id, lambda edge_data: isinstance(edge_data, Clbit)
// )
// )

/// Returns iterator of the predecessors of a node that are
/// connected by a quantum edge as DAGOpNodes and DAGInNodes.
fn quantum_predecessors(&self, py: Python, node: &DAGNode) -> PyResult<Py<PyIterator>> {
let node = NodeIndex::new(node._node_id as usize);
let edges = self.dag.edges_directed(node, Incoming);
let filtered = edges.filter_map(|e| match e.weight() {
Wire::Qubit(_) => Some(e.source()),
_ => None,
});
let predecessors: PyResult<Vec<_>> =
filtered.unique().map(|i| self.get_node(py, i)).collect();
Ok(PyTuple::new_bound(py, predecessors?)
.into_any()
.iter()
.unwrap()
.unbind())
}

/// Returns iterator of the predecessors of a node that are
/// connected by a classical edge as DAGOpNodes and DAGInNodes.
fn classical_predecessors(&self, py: Python, node: &DAGNode) -> PyResult<Py<PyIterator>> {
let node = NodeIndex::new(node._node_id as usize);
let edges = self.dag.edges_directed(node, Incoming);
let filtered = edges.filter_map(|e| match e.weight() {
Wire::Clbit(_) => Some(e.source()),
_ => None,
});
let predecessors: PyResult<Vec<_>> =
filtered.unique().map(|i| self.get_node(py, i)).collect();
Ok(PyTuple::new_bound(py, predecessors?)
.into_any()
.iter()
.unwrap()
.unbind())
}

/// Returns set of the ancestors of a node as DAGOpNodes and DAGInNodes.
fn ancestors(&self, py: Python, node: &DAGNode) -> PyResult<Py<PySet>> {
// return {self._multi_graph[x] for x in rx.ancestors(self._multi_graph, node._node_id)}
todo!()
}

/// Returns set of the descendants of a node as DAGOpNodes and DAGOutNodes.
fn descendants(&self, py: Python, node: &DAGNode) -> PyResult<Py<PySet>> {
// return {self._multi_graph[x] for x in rx.descendants(self._multi_graph, node._node_id)}
todo!()
}

/// Returns an iterator of tuples of (DAGNode, [DAGNodes]) where the DAGNode is the current node
/// and [DAGNode] is its successors in BFS order.
fn bfs_successors(&self, py: Python, node: &DAGNode) -> PyResult<Py<PySet>> {
// return iter(rx.bfs_successors(self._multi_graph, node._node_id))
todo!()
}

/// Returns iterator of the successors of a node that are
/// connected by a quantum edge as DAGOpNodes and DAGOutNodes.
fn quantum_successors(&self, py: Python, node: &DAGNode) -> PyResult<Py<PyIterator>> {
let node = NodeIndex::new(node._node_id as usize);
let edges = self.dag.edges_directed(node, Outgoing);
let filtered = edges.filter_map(|e| match e.weight() {
Wire::Qubit(_) => Some(e.target()),
_ => None,
});
let predecessors: PyResult<Vec<_>> =
filtered.unique().map(|i| self.get_node(py, i)).collect();
Ok(PyTuple::new_bound(py, predecessors?)
.into_any()
.iter()
.unwrap()
.unbind())
}

/// Returns iterator of the successors of a node that are
/// connected by a classical edge as DAGOpNodes and DAGOutNodes.
fn classical_successors(&self, py: Python, node: &DAGNode) -> PyResult<Py<PyIterator>> {
let node = NodeIndex::new(node._node_id as usize);
let edges = self.dag.edges_directed(node, Incoming);
let filtered = edges.filter_map(|e| match e.weight() {
Wire::Clbit(_) => Some(e.target()),
_ => None,
});
let predecessors: PyResult<Vec<_>> =
filtered.unique().map(|i| self.get_node(py, i)).collect();
Ok(PyTuple::new_bound(py, predecessors?)
.into_any()
.iter()
.unwrap()
.unbind())
}
//
// def remove_op_node(self, node):
// """Remove an operation node n.
Expand Down

0 comments on commit 13d18fd

Please sign in to comment.