#[allow(unused_imports)]
use crate::mwpf::shim::prelude::*;
use crate::mwpf::decoding_hypergraph::*;
use crate::mwpf::invalid_subgraph::*;
use crate::mwpf::model_hypergraph::*;
use crate::mwpf::num_traits::{FromPrimitive, One, Signed, ToPrimitive, Zero};
use crate::mwpf::pointers::*;
use crate::mwpf::primal_module::Affinity;
use crate::mwpf::primal_module_serial::PrimalClusterPtr;
use crate::mwpf::relaxer_optimizer::OptimizerResult;
use crate::mwpf::util::*;
use crate::mwpf::visualize::*;
use hashbrown::{HashMap, HashSet};
#[cfg(feature = "python")]
use pyo3::prelude::*;
use std::sync::Arc;
#[derive(Default, Debug, PartialEq, Eq, Clone)]
#[cfg_attr(feature = "python", pyclass(module = "stabilizer_qec.mwpf", eq, eq_int))]
pub enum DualModuleMode {
#[default]
Search,
Tune, }
impl DualModuleMode {
pub fn new() -> Self {
Self::default()
}
pub fn advance(&mut self) {
match self {
Self::Search => *self = Self::Tune,
Self::Tune => panic!("dual module mode is already in tune mode"),
}
}
pub fn reset(&mut self) {
*self = Self::Search;
}
}
macro_rules! add_shared_methods {
() => {
fn mode(&self) -> &DualModuleMode {
&self.mode
}
fn mode_mut(&mut self) -> &mut DualModuleMode {
&mut self.mode
}
};
}
pub(crate) use add_shared_methods;
pub struct DualNode {
pub index: NodeIndex,
pub invalid_subgraph: Arc<InvalidSubgraph>,
pub grow_rate: Rational,
global_time: Option<ArcRwLock<Rational>>,
pub last_updated_time: Rational,
pub dual_variable_at_last_updated_time: Rational,
}
impl DualNode {
pub fn get_dual_variable(&self) -> Rational {
match self.global_time.clone() {
Some(global_time) => {
let global_time = global_time.read_recursive();
if global_time.ge(&self.last_updated_time) {
(global_time.clone() - self.last_updated_time.clone()) * self.grow_rate.clone()
+ self.dual_variable_at_last_updated_time.clone()
} else {
self.dual_variable_at_last_updated_time.clone()
}
}
None => self.dual_variable_at_last_updated_time.clone(),
}
}
pub fn set_dual_variable(&mut self, new_dual_variable: Rational) {
self.dual_variable_at_last_updated_time = new_dual_variable;
}
pub fn init_time(&mut self, global_time_ptr: ArcRwLock<Rational>) {
self.last_updated_time = global_time_ptr.read_recursive().clone();
self.global_time = Some(global_time_ptr);
}
}
pub type DualNodePtr = ArcRwLock<DualNode>;
pub type DualNodeWeak = WeakRwLock<DualNode>;
impl std::fmt::Debug for DualNodePtr {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
let dual_node = self.read_recursive(); f.debug_struct("DualNode")
.field("index", &dual_node.index)
.field("dual_variable", &dual_node.get_dual_variable())
.field("grow_rate", &dual_node.grow_rate)
.field("hair", &dual_node.invalid_subgraph.hair)
.finish()
}
}
impl std::fmt::Debug for DualNodeWeak {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
self.upgrade_force().fmt(f)
}
}
#[derive(Debug)]
#[cfg_attr(feature = "python", pyclass(module = "stabilizer_qec.mwpf"))]
pub struct DualModuleInterface {
pub nodes: Vec<DualNodePtr>,
pub hashmap: HashMap<Arc<InvalidSubgraph>, NodeIndex>,
pub decoding_graph: DecodingHyperGraph,
}
pub type DualModuleInterfacePtr = ArcRwLock<DualModuleInterface>;
pub type DualModuleInterfaceWeak = WeakRwLock<DualModuleInterface>;
impl std::fmt::Debug for DualModuleInterfacePtr {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
let interface = self.read_recursive();
write!(f, "{}", interface.nodes.len())
}
}
impl std::fmt::Debug for DualModuleInterfaceWeak {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
self.upgrade_force().fmt(f)
}
}
#[derive(PartialEq, Eq, Clone, Debug)]
pub struct OrderedDualNodePtr {
pub index: NodeIndex,
pub ptr: DualNodePtr,
}
impl std::hash::Hash for OrderedDualNodePtr {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
self.index.hash(state);
}
}
impl OrderedDualNodePtr {
pub fn new(index: NodeIndex, ptr: DualNodePtr) -> Self {
Self { index, ptr }
}
}
impl From<DualNodePtr> for OrderedDualNodePtr {
fn from(ptr: DualNodePtr) -> Self {
let index = ptr.read_recursive().index;
Self { index, ptr }
}
}
impl PartialOrd for OrderedDualNodePtr {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.index.cmp(&other.index))
}
}
impl Ord for OrderedDualNodePtr {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
self.index.cmp(&other.index)
}
}
#[derive(PartialEq, Eq, Clone, Debug)]
pub struct OrderedDualNodeWeak {
pub index: NodeIndex,
pub weak_ptr: DualNodeWeak,
}
impl OrderedDualNodeWeak {
pub fn new(index: NodeIndex, weak_ptr: DualNodeWeak) -> Self {
Self { index, weak_ptr }
}
pub fn upgrade_force(&self) -> OrderedDualNodePtr {
OrderedDualNodePtr::new(self.index, self.weak_ptr.upgrade_force())
}
}
impl PartialOrd for OrderedDualNodeWeak {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.index.cmp(&other.index))
}
}
impl Ord for OrderedDualNodeWeak {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
self.index.cmp(&other.index)
}
}
#[derive(Clone, Debug, Default)]
pub enum DualReport {
#[default]
Unbounded,
ValidGrow(Rational),
Obstacles(Vec<Obstacle>),
}
pub trait DualModuleImpl {
fn new_empty(initializer: &Arc<SolverInitializer>) -> Self;
fn clear(&mut self);
fn add_defect_node(&mut self, dual_node_ptr: &DualNodePtr);
fn add_dual_node(&mut self, dual_node_ptr: &DualNodePtr);
fn set_grow_rate(&mut self, dual_node_ptr: &DualNodePtr, grow_rate: Rational);
fn report(&mut self) -> DualReport;
fn grow_dual_node(&mut self, _dual_node_ptr: &DualNodePtr, _length: Rational) {
panic!("the dual module implementation doesn't support this function, please use another dual module")
}
fn grow(&mut self, length: Rational);
fn get_edge_nodes(&self, edge_index: EdgeIndex) -> Vec<DualNodePtr>;
fn get_edge_slack(&self, edge_index: EdgeIndex) -> Rational;
fn is_edge_tight(&self, edge_index: EdgeIndex) -> bool;
fn mode(&self) -> &DualModuleMode;
fn mode_mut(&mut self) -> &mut DualModuleMode;
fn advance_mode(&mut self) {
eprintln!("this dual_module does not implement different modes");
}
fn reset_mode(&mut self) {
*self.mode_mut() = DualModuleMode::default();
}
fn end_tuning(&mut self) {
panic!("this module doesn't work with tuning")
}
fn get_total_tuning_time(&self) -> Option<f64> {
panic!("this module doesn't work with tuning")
}
fn clear_tuning_time(&mut self) {
panic!("this module doesn't work with tuning")
}
fn add_dual_node_tune(&mut self, dual_node_ptr: &DualNodePtr) {
eprintln!("this dual_module does not implement tuning");
self.add_dual_node(dual_node_ptr);
}
fn sync(&mut self) {
panic!("this dual_module does not have global time and does not need to sync");
}
fn grow_edge(&self, _edge_index: EdgeIndex, _amount: &Rational) {
panic!("this dual_module doesn't support edge growth");
}
fn is_edge_tight_tune(&self, edge_index: EdgeIndex) -> bool {
eprintln!("this dual_module does not implement tuning");
self.is_edge_tight(edge_index)
}
fn get_edge_slack_tune(&self, edge_index: EdgeIndex) -> Rational {
eprintln!("this dual_module does not implement tuning");
self.get_edge_slack(edge_index)
}
fn debug_print(&self) {
println!("this dual_module doesn't support debug print");
}
fn calculate_cluster_affinity(&mut self, _cluster: PrimalClusterPtr) -> Option<Affinity> {
eprintln!("not implemented, skipping");
Some(Affinity::from(100.0))
}
fn get_obstacles_tune(
&self,
optimizer_result: OptimizerResult,
dual_node_deltas: FastIterMap<OrderedDualNodePtr, (Rational, NodeIndex)>,
) -> FastIterSet<Obstacle> {
let mut obstacles = FastIterSet::new();
match optimizer_result {
OptimizerResult::EarlyReturned => {
for (dual_node_ptr, (grow_rate, _)) in dual_node_deltas.into_iter() {
let node_ptr_read = dual_node_ptr.ptr.read_recursive();
if grow_rate.is_negative() && node_ptr_read.dual_variable_at_last_updated_time.is_zero() {
obstacles.insert(Obstacle::ShrinkToZero {
dual_node_ptr: dual_node_ptr.clone(),
});
}
for edge_index in node_ptr_read.invalid_subgraph.hair.iter() {
if grow_rate.is_positive() && self.is_edge_tight_tune(*edge_index) {
obstacles.insert(Obstacle::Conflict { edge_index: *edge_index });
}
}
}
}
OptimizerResult::Skipped => {
for (dual_node_ptr, (grow_rate, _cluster_index)) in dual_node_deltas.into_iter() {
let mut actual_grow_rate = Rational::from_usize(usize::MAX).unwrap();
let node_ptr_read = dual_node_ptr.ptr.read_recursive();
for edge_index in node_ptr_read.invalid_subgraph.hair.iter() {
actual_grow_rate = std::cmp::min(actual_grow_rate, self.get_edge_slack_tune(*edge_index));
}
if actual_grow_rate.is_zero() {
for edge_index in node_ptr_read.invalid_subgraph.hair.iter() {
if grow_rate.is_positive() && self.is_edge_tight_tune(*edge_index) {
obstacles.insert(Obstacle::Conflict { edge_index: *edge_index });
}
}
if grow_rate.is_negative() && node_ptr_read.dual_variable_at_last_updated_time.is_zero() {
obstacles.insert(Obstacle::ShrinkToZero {
dual_node_ptr: dual_node_ptr.clone(),
});
}
} else {
drop(node_ptr_read);
let mut node_ptr_write = dual_node_ptr.ptr.write();
for edge_index in node_ptr_write.invalid_subgraph.hair.iter() {
self.grow_edge(*edge_index, &actual_grow_rate);
if actual_grow_rate.is_positive() && self.is_edge_tight_tune(*edge_index) {
obstacles.insert(Obstacle::Conflict { edge_index: *edge_index });
}
}
node_ptr_write.dual_variable_at_last_updated_time += actual_grow_rate.clone();
if actual_grow_rate.is_negative() && node_ptr_write.dual_variable_at_last_updated_time.is_zero() {
obstacles.insert(Obstacle::ShrinkToZero {
dual_node_ptr: dual_node_ptr.clone(),
});
}
}
}
}
_ => {
let mut edge_deltas = FastIterMap::new();
for (dual_node_ptr, (grow_rate, _cluster_index)) in dual_node_deltas.into_iter() {
let mut node_ptr_write = dual_node_ptr.ptr.write();
node_ptr_write.dual_variable_at_last_updated_time += grow_rate.clone();
if grow_rate.is_negative() && node_ptr_write.dual_variable_at_last_updated_time.is_zero() {
obstacles.insert(Obstacle::ShrinkToZero {
dual_node_ptr: dual_node_ptr.clone(),
});
}
for edge_index in node_ptr_write.invalid_subgraph.hair.iter() {
match edge_deltas.entry(*edge_index) {
FastIterEntry::Vacant(v) => {
v.insert(grow_rate.clone());
}
FastIterEntry::Occupied(mut o) => {
o.insert(o.get() + grow_rate.clone());
}
}
}
}
for (edge_index, grow_rate) in edge_deltas.into_iter() {
if grow_rate.is_zero() {
continue;
}
self.grow_edge(edge_index, &grow_rate);
if grow_rate.is_positive() && self.is_edge_tight_tune(edge_index) {
obstacles.insert(Obstacle::Conflict { edge_index });
}
}
}
}
obstacles
}
fn get_edge_free_weight(
&self,
edge_index: EdgeIndex,
participating_dual_variables: &hashbrown::HashSet<usize>,
) -> Weight;
fn get_edge_weight(&self, edge_index: EdgeIndex) -> Weight;
fn get_subgraph_weight(&self, subgraph: &Subgraph) -> Weight {
let mut weight = Weight::zero();
for &edge_index in subgraph {
weight += self.get_edge_weight(edge_index);
}
weight
}
fn adjust_weights_for_negative_edges(&mut self) {
unimplemented!()
}
fn update_weights(&mut self, _new_weights: Vec<Weight>, _mix_ratio: Weight) {
unimplemented!()
}
fn set_weights(&mut self, _new_weights: FastIterMap<EdgeIndex, Weight>) {
unimplemented!()
}
fn get_negative_weight_sum(&self) -> Rational {
unimplemented!()
}
fn get_negative_edges(&self) -> HashSet<EdgeIndex> {
unimplemented!()
}
fn get_flip_vertices(&self) -> HashSet<VertexIndex> {
unimplemented!()
}
}
#[derive(PartialEq, Eq, Debug, Clone, PartialOrd, Ord)]
pub enum Obstacle {
Conflict { edge_index: EdgeIndex },
ShrinkToZero { dual_node_ptr: OrderedDualNodePtr },
}
impl std::hash::Hash for Obstacle {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
match self {
Obstacle::Conflict { edge_index } => {
(0, *edge_index as u64).hash(state);
}
Obstacle::ShrinkToZero { dual_node_ptr } => {
(1, dual_node_ptr.index).hash(state);
}
}
}
}
impl DualReport {
pub fn new() -> Self {
Self::default()
}
pub fn add_obstacle(&mut self, obstacle: Obstacle) {
match self {
Self::Unbounded | Self::ValidGrow(_) => {
*self = Self::Obstacles(vec![obstacle]);
}
Self::Obstacles(obstacles) => {
obstacles.push(obstacle);
}
}
}
pub fn is_unbounded(&self) -> bool {
matches!(self, Self::Unbounded)
}
pub fn get_valid_growth(&self) -> Option<Rational> {
match self {
Self::Unbounded => {
panic!("please call DualReport::is_unbounded to check if it's unbounded");
}
Self::ValidGrow(length) => Some(length.clone()),
_ => None,
}
}
pub fn pop(&mut self) -> Option<Obstacle> {
match self {
Self::Unbounded | Self::ValidGrow(_) => {
panic!("please call DualReport::get_valid_growth to check if this group is none_zero_growth");
}
Self::Obstacles(obstacles) => obstacles.pop(),
}
}
pub fn peek(&self) -> Option<&Obstacle> {
match self {
Self::Unbounded | Self::ValidGrow(_) => {
panic!("please call DualReport::get_valid_growth to check if this group is none_zero_growth");
}
Self::Obstacles(obstacles) => obstacles.last(),
}
}
pub fn iter(&self) -> Option<std::slice::Iter<Obstacle>> {
match self {
Self::Unbounded | Self::ValidGrow(_) => None,
Self::Obstacles(obstacles) => Some(obstacles.iter()),
}
}
}
impl DualModuleInterfacePtr {
pub fn new(model_graph: Arc<ModelHyperGraph>) -> Self {
Self::new_value(DualModuleInterface {
nodes: Vec::new(),
hashmap: HashMap::new(),
decoding_graph: DecodingHyperGraph::new(model_graph, Arc::new(SyndromePattern::new_empty())),
})
}
pub fn new_load(decoding_graph: DecodingHyperGraph, dual_module_impl: &mut impl DualModuleImpl) -> Self {
let interface_ptr = Self::new(decoding_graph.model_graph.clone());
interface_ptr.load(decoding_graph.syndrome_pattern, dual_module_impl);
interface_ptr
}
pub fn load(&self, syndrome_pattern: Arc<SyndromePattern>, dual_module_impl: &mut impl DualModuleImpl) {
self.write().decoding_graph.set_syndrome(syndrome_pattern.clone());
for vertex_idx in syndrome_pattern.defect_vertices.iter() {
self.create_defect_node(*vertex_idx, dual_module_impl);
}
}
pub fn sum_dual_variables(&self) -> Rational {
let interface = self.read_recursive();
interface
.nodes
.iter()
.map(|node_ptr| node_ptr.read_recursive().get_dual_variable())
.sum()
}
pub fn clear(&self) {
let mut interface = self.write();
interface.nodes.clear();
interface.hashmap.clear();
}
#[allow(clippy::unnecessary_cast)]
pub fn get_node(&self, node_index: NodeIndex) -> Option<DualNodePtr> {
let interface = self.read_recursive();
interface.nodes.get(node_index as usize).cloned()
}
fn create_defect_node(&self, vertex_idx: VertexIndex, dual_module: &mut impl DualModuleImpl) -> DualNodePtr {
let mut interface = self.write();
let invalid_subgraph = Arc::new(InvalidSubgraph::new_complete(
vec![vertex_idx].into_iter().collect(),
FastIterSet::new(),
&interface.decoding_graph,
));
let node_index = interface.nodes.len() as NodeIndex;
let node_ptr = DualNodePtr::new_value(DualNode {
index: node_index,
invalid_subgraph: invalid_subgraph.clone(),
grow_rate: Rational::one(),
dual_variable_at_last_updated_time: Rational::zero(),
global_time: None,
last_updated_time: Rational::zero(),
});
interface.nodes.push(node_ptr.clone());
interface.hashmap.insert(invalid_subgraph, node_index);
drop(interface);
dual_module.add_defect_node(&node_ptr);
node_ptr
}
#[allow(clippy::unnecessary_cast)]
pub fn find_node(&self, invalid_subgraph: &Arc<InvalidSubgraph>) -> Option<DualNodePtr> {
let interface = self.read_recursive();
interface
.hashmap
.get(invalid_subgraph)
.map(|index| interface.nodes[*index as usize].clone())
}
pub fn create_node(&self, invalid_subgraph: Arc<InvalidSubgraph>, dual_module: &mut impl DualModuleImpl) -> DualNodePtr {
self.create_node_internal(invalid_subgraph, dual_module, Rational::one(), DualModuleImpl::add_dual_node)
}
pub fn create_node_tune(
&self,
invalid_subgraph: Arc<InvalidSubgraph>,
dual_module: &mut impl DualModuleImpl,
) -> DualNodePtr {
self.create_node_internal(
invalid_subgraph,
dual_module,
Rational::zero(),
DualModuleImpl::add_dual_node_tune,
)
}
pub fn find_or_create_node(
&self,
invalid_subgraph: &Arc<InvalidSubgraph>,
dual_module: &mut impl DualModuleImpl,
) -> (bool, DualNodePtr) {
match self.find_node(invalid_subgraph) {
Some(node_ptr) => (true, node_ptr),
None => (false, self.create_node(invalid_subgraph.clone(), dual_module)),
}
}
pub fn find_or_create_node_tune(
&self,
invalid_subgraph: &Arc<InvalidSubgraph>,
dual_module: &mut impl DualModuleImpl,
) -> Option<(bool, DualNodePtr)> {
match self.find_node(invalid_subgraph) {
Some(node_ptr) => Some((true, node_ptr)),
None => Some((false, self.create_node_tune(invalid_subgraph.clone(), dual_module))),
}
}
fn create_node_internal<D: DualModuleImpl>(
&self,
invalid_subgraph: Arc<InvalidSubgraph>,
dual_module: &mut D,
grow_rate: Rational,
add_dual_node_fn: fn(&mut D, &DualNodePtr),
) -> DualNodePtr {
debug_assert!(
self.find_node(&invalid_subgraph).is_none(),
"do not create the same node twice"
);
let mut interface = self.write();
let node_index = interface.nodes.len() as NodeIndex;
interface.hashmap.insert(invalid_subgraph.clone(), node_index);
let node_ptr = DualNodePtr::new_value(DualNode {
index: node_index,
invalid_subgraph,
grow_rate,
dual_variable_at_last_updated_time: Rational::zero(),
global_time: None,
last_updated_time: Rational::zero(),
});
interface.nodes.push(node_ptr.clone());
drop(interface);
add_dual_node_fn(dual_module, &node_ptr);
node_ptr
}
}
impl DualModuleInterfacePtr {
pub fn create_node_vec(&self, edges: &[EdgeIndex], dual_module: &mut impl DualModuleImpl) -> DualNodePtr {
let invalid_subgraph = Arc::new(InvalidSubgraph::new(
edges.iter().cloned().collect(),
&self.read_recursive().decoding_graph,
));
self.create_node(invalid_subgraph, dual_module)
}
pub fn create_node_complete_vec(
&self,
vertices: &[VertexIndex],
edges: &[EdgeIndex],
dual_module: &mut impl DualModuleImpl,
) -> DualNodePtr {
let invalid_subgraph = Arc::new(InvalidSubgraph::new_complete(
vertices.iter().cloned().collect(),
edges.iter().cloned().collect(),
&self.read_recursive().decoding_graph,
));
self.create_node(invalid_subgraph, dual_module)
}
}
impl MWPSVisualizer for DualModuleInterfacePtr {
fn snapshot(&self, abbrev: bool) -> serde_json::Value {
let interface = self.read_recursive();
let mut dual_nodes = Vec::<serde_json::Value>::new();
for dual_node_ptr in interface.nodes.iter() {
let dual_node = dual_node_ptr.read_recursive();
dual_nodes.push(json!({
if abbrev { "e" } else { "edges" }: dual_node.invalid_subgraph.edges,
if abbrev { "v" } else { "vertices" }: dual_node.invalid_subgraph.vertices,
if abbrev { "h" } else { "hair" }: dual_node.invalid_subgraph.hair,
if abbrev { "d" } else { "dual_variable" }: dual_node.get_dual_variable().to_f64(),
if abbrev { "dn" } else { "dual_variable_numerator" }: numer_of(&dual_node.get_dual_variable()),
if abbrev { "dd" } else { "dual_variable_denominator" }: denom_of(&dual_node.get_dual_variable()),
if abbrev { "r" } else { "grow_rate" }: dual_node.grow_rate.to_f64(),
if abbrev { "rn" } else { "grow_rate_numerator" }: numer_of(&dual_node.grow_rate),
if abbrev { "rd" } else { "grow_rate_denominator" }: denom_of(&dual_node.grow_rate),
}));
}
let sum_dual = self.sum_dual_variables();
json!({
"interface": {
"sum_dual": sum_dual.to_f64(),
"sdn": numer_of(&sum_dual),
"sdd": denom_of(&sum_dual),
},
"dual_nodes": dual_nodes,
})
}
}