Skip to main content

tensorlogic_quantrs_hooks/
lib.rs

1//! TL <-> QuantrS2 hooks (PGM/message passing as reductions).
2//!
3//! **Version**: 0.1.1 | **Status**: Production Ready
4//!
5//! This crate provides integration between TensorLogic and probabilistic graphical models (PGMs).
6//! It maps belief propagation and other message passing algorithms onto einsum reduction patterns.
7//!
8//! # Core Concepts
9//!
10//! - **Factor Graphs**: Convert TLExpr predicates into factors
11//! - **Message Passing**: Sum-product and max-product algorithms as tensor operations
12//! - **Inference**: Marginalization and conditional queries via reductions
13//!
14//! # Architecture
15//!
16//! ```text
17//! TLExpr → FactorGraph → MessagePassing → Marginals
18//!    ↓         ↓              ↓              ↓
19//! Predicates Factors    Einsum Ops    Probabilities
20//! ```
21
22mod cache;
23pub mod convergence;
24pub mod dbn;
25mod elimination_ordering;
26mod error;
27mod expectation_propagation;
28mod factor;
29pub mod factor_graph_viz;
30mod graph;
31mod inference;
32pub mod influence;
33mod junction_tree;
34mod linear_chain_crf;
35pub mod loopy_bp;
36pub mod memory;
37mod message_passing;
38mod models;
39mod parallel_message_passing;
40pub mod parameter_learning;
41pub mod quantrs_hooks;
42pub mod quantum_circuit;
43pub mod quantum_simulation;
44mod sampling;
45pub mod tensor_network_bridge;
46mod variable_elimination;
47mod variational;
48pub mod vmp;
49
50pub use cache::{CacheStats, CachedFactor, FactorCache};
51pub use convergence::{
52    ConvergenceConfig, ConvergenceError, ConvergenceMonitor, ConvergenceState, DampingSchedule,
53    InferenceStats,
54};
55pub use dbn::{CoupledDBN, CouplingFactor, DBNBuilder, DynamicBayesianNetwork, TemporalVar};
56pub use elimination_ordering::{EliminationOrdering, EliminationStrategy};
57pub use error::{PgmError, Result};
58pub use expectation_propagation::{ExpectationPropagation, GaussianEP, GaussianSite, Site};
59pub use factor::{Factor, FactorOp};
60pub use factor_graph_viz::{
61    render_ascii, render_dot, FactorGraphModel, FactorGraphStats, VizFactorNode, VizVariableNode,
62};
63pub use graph::{FactorGraph, FactorNode, VariableNode};
64pub use inference::{ConditionalQuery, InferenceEngine, MarginalizationQuery};
65pub use influence::{
66    InfluenceDiagram, InfluenceDiagramBuilder, InfluenceNode, MultiAttributeUtility, NodeType,
67};
68pub use junction_tree::{Clique, JunctionTree, JunctionTreeEdge, Separator};
69pub use linear_chain_crf::{
70    EmissionFeature, FeatureFunction, IdentityFeature, LinearChainCRF, TransitionFeature,
71};
72pub use loopy_bp::{
73    bethe_free_energy, BetheFreeEnergy, CycleAnalysis, CycleDetector, LbpConvergenceMonitor,
74    LbpDampingPolicy, LbpIterStats, LogMessage, LoopyBeliefPropagation, LoopyBpConfig,
75    LoopyBpResult, UpdateSchedule,
76};
77pub use memory::{
78    BlockSparseFactor, CompressedFactor, FactorPool, LazyFactor, MemoryEstimate, PoolStats,
79    SparseFactor, StreamingFactorGraph,
80};
81pub use message_passing::{
82    ConvergenceStats, MaxProductAlgorithm, MessagePassingAlgorithm, SumProductAlgorithm,
83};
84pub use models::{BayesianNetwork, ConditionalRandomField, HiddenMarkovModel, MarkovRandomField};
85pub use parallel_message_passing::{ParallelMaxProduct, ParallelSumProduct};
86pub use parameter_learning::{
87    BaumWelchLearner, BayesianEstimator, MaximumLikelihoodEstimator, SimpleHMM,
88};
89pub use quantrs_hooks::{
90    AnnealingConfig, DistributionExport, DistributionMetadata, ModelExport, ModelStatistics,
91    QuantRSAssignment, QuantRSDistribution, QuantRSInferenceQuery, QuantRSModelExport,
92    QuantRSParameterLearning, QuantRSSamplingHook, QuantumAnnealing, QuantumInference,
93    QuantumSolution, QuantumSolutionMetadata,
94};
95pub use quantum_circuit::{
96    tlexpr_to_qaoa_circuit, IsingModel, QAOAConfig, QAOAResult, QUBOProblem, QuantumCircuitBuilder,
97};
98pub use quantum_simulation::{
99    run_qaoa, QuantumSimulationBackend, SimulatedState, SimulationConfig,
100};
101pub use sampling::{
102    Assignment, GibbsSampler, ImportanceSampler, LikelihoodWeighting, Particle, ParticleFilter,
103    ProposalDistribution, WeightedSample,
104};
105pub use tensor_network_bridge::{
106    factor_graph_to_tensor_network, linear_chain_to_mps, MatrixProductState, Tensor, TensorNetwork,
107    TensorNetworkStats,
108};
109pub use variable_elimination::VariableElimination;
110pub use variational::{BetheApproximation, MeanFieldInference, TreeReweightedBP};
111pub use vmp::{
112    categorical_kl, dirichlet_kl, gaussian_kl, gaussian_kl_fixed_precision,
113    BetaBernoulliObservation, BetaNP, CategoricalNP, DirichletNP, ExponentialFamily, Family,
114    GammaNP, GammaPoissonObservation, GaussianNP, MessageDirection, VariationalGaussianMixture,
115    VariationalMessagePassing, VariationalState, VgmmConfig, VgmmResult, VmpConfig, VmpFactor,
116    VmpMessage, VmpResult,
117};
118
119use scirs2_core::ndarray::ArrayD;
120use std::collections::HashMap;
121use tensorlogic_ir::TLExpr;
122
123/// Convert a TensorLogic expression to a factor graph.
124///
125/// This function analyzes the logical structure and creates a factor graph
126/// where predicates become factors and quantified variables become nodes.
127pub fn expr_to_factor_graph(expr: &TLExpr) -> Result<FactorGraph> {
128    let mut graph = FactorGraph::new();
129
130    // Recursively extract factors from expression
131    extract_factors(expr, &mut graph)?;
132
133    Ok(graph)
134}
135
136/// Extract factors from a TLExpr and add them to the factor graph.
137fn extract_factors(expr: &TLExpr, graph: &mut FactorGraph) -> Result<()> {
138    match expr {
139        TLExpr::Pred { name, args } => {
140            // Create a factor from predicate
141            let var_names: Vec<String> = args
142                .iter()
143                .filter_map(|term| match term {
144                    tensorlogic_ir::Term::Var(v) => Some(v.clone()),
145                    _ => None,
146                })
147                .collect();
148
149            // Add variables if they don't exist
150            for var_name in &var_names {
151                if graph.get_variable(var_name).is_none() {
152                    graph.add_variable(var_name.clone(), "default".to_string());
153                }
154            }
155
156            if !var_names.is_empty() {
157                graph.add_factor_from_predicate(name, &var_names)?;
158            }
159        }
160        TLExpr::And(left, right) => {
161            // Conjunction creates multiple factors
162            extract_factors(left, graph)?;
163            extract_factors(right, graph)?;
164        }
165        TLExpr::Exists { var, domain, body } | TLExpr::ForAll { var, domain, body } => {
166            // Quantified variables become nodes in the factor graph
167            graph.add_variable(var.clone(), domain.clone());
168            extract_factors(body, graph)?;
169        }
170        TLExpr::Imply(premise, conclusion) => {
171            // Implication can be represented as factors
172            extract_factors(premise, graph)?;
173            extract_factors(conclusion, graph)?;
174        }
175        TLExpr::Not(inner) => {
176            // Negation affects factor values
177            extract_factors(inner, graph)?;
178        }
179        _ => {
180            // Other expressions may not directly map to factors
181        }
182    }
183
184    Ok(())
185}
186
187/// Perform message passing inference on a factor graph.
188///
189/// This function runs belief propagation to compute marginal probabilities.
190pub fn message_passing_reduce(
191    graph: &FactorGraph,
192    algorithm: &dyn MessagePassingAlgorithm,
193) -> Result<HashMap<String, ArrayD<f64>>> {
194    algorithm.run(graph)
195}
196
197/// Compute marginal probability for a variable.
198///
199/// This maps to a reduction operation over all other variables.
200pub fn marginalize(
201    joint_distribution: &ArrayD<f64>,
202    variable_idx: usize,
203    axes_to_sum: &[usize],
204) -> Result<ArrayD<f64>> {
205    use scirs2_core::ndarray::Axis;
206
207    let mut result = joint_distribution.clone();
208
209    // Sum over all axes except the target variable
210    for &axis in axes_to_sum.iter().rev() {
211        if axis != variable_idx {
212            result = result.sum_axis(Axis(axis));
213        }
214    }
215
216    Ok(result)
217}
218
219/// Compute conditional probability P(X | Y = y).
220///
221/// This slices the joint distribution at the evidence values.
222pub fn condition(
223    joint_distribution: &ArrayD<f64>,
224    evidence: &HashMap<usize, usize>,
225) -> Result<ArrayD<f64>> {
226    let mut result = joint_distribution.clone();
227
228    // Slice at evidence values
229    for (&var_idx, &value) in evidence {
230        result = result.index_axis_move(scirs2_core::ndarray::Axis(var_idx), value);
231    }
232
233    // Normalize
234    let sum: f64 = result.iter().sum();
235    if sum > 0.0 {
236        result /= sum;
237    }
238
239    Ok(result)
240}
241
242#[cfg(test)]
243mod tests {
244    use super::*;
245    use approx::assert_abs_diff_eq;
246    use scirs2_core::ndarray::Array;
247    use tensorlogic_ir::Term;
248
249    #[test]
250    fn test_expr_to_factor_graph() {
251        let expr = TLExpr::pred("P", vec![Term::var("x")]);
252        let graph = expr_to_factor_graph(&expr).expect("unwrap");
253        assert!(!graph.is_empty());
254    }
255
256    #[test]
257    fn test_marginalize_simple() {
258        // 2x2 joint distribution: P(X, Y)
259        let joint = Array::from_shape_vec(vec![2, 2], vec![0.25, 0.25, 0.25, 0.25])
260            .expect("unwrap")
261            .into_dyn();
262
263        // Marginalize over Y (axis 1) to get P(X)
264        let marginal = marginalize(&joint, 0, &[0, 1]).expect("unwrap");
265
266        assert_eq!(marginal.ndim(), 1);
267        assert_abs_diff_eq!(marginal.sum(), 1.0, epsilon = 1e-10);
268    }
269
270    #[test]
271    fn test_condition_simple() {
272        // 2x2 joint distribution
273        let joint = Array::from_shape_vec(vec![2, 2], vec![0.1, 0.2, 0.3, 0.4])
274            .expect("unwrap")
275            .into_dyn();
276
277        // Condition on Y=1: P(X | Y=1)
278        let mut evidence = HashMap::new();
279        evidence.insert(1, 1);
280
281        let conditional = condition(&joint, &evidence).expect("unwrap");
282
283        // Should have one dimension less
284        assert_eq!(conditional.ndim(), 1);
285        // Should be normalized
286        assert_abs_diff_eq!(conditional.sum(), 1.0, epsilon = 1e-10);
287    }
288
289    #[test]
290    fn test_factor_graph_construction() {
291        let mut graph = FactorGraph::new();
292        graph.add_variable("x".to_string(), "Domain1".to_string());
293        graph.add_variable("y".to_string(), "Domain2".to_string());
294
295        assert_eq!(graph.num_variables(), 2);
296    }
297}