smc (particle filter) stitch

Cargo.toml

@@ -22,6 +22,7 @@ rustc-hash = "1.1.0"
 # lambdas = "0.2.0"
 lambdas = { git = "https://github.com/mlb2251/lambdas", rev = "2c9bfd0"}
 test-case = "3.3.1"
+modppl = "0.3.0"
 rand_chacha = "0.9.0"
 shlex = "1.3.0"
 

data/regression/furniture-small.json

@@ -0,0 +1,10 @@
+[
+    "(T (repeat (repeat (T (C (C (T (T c (M 2 0 0 0)) (M 1.125 0 0 0)) (T (T c (M 2 0 0 0)) (M 1.625 0 0 0))) (T r (M 0.84375 0 0 0))) (M 1 0 0 0)) 1 (M 1 0 0 0)) 1 (M 1 0 0 0)) (M 1 0 -6.9375 0))",
+    "(T (repeat (repeat (T (C (T r (M 1.125 0 0 0)) (T c (M 0.84375 0 0 0))) (M 1 0 0 0)) 1 (M 1 0 0 0)) 1 (M 1 0 0 0)) (M 1 0 -7.1875 0))",
+    "(T (repeat (repeat (T (C (T (T (repeat (T l (M 1 0 -0.5 1.20711)) 8 (M 1 0.785398 0 0)) (M 0.75 0 0 0)) (M 1.125 0 0 0)) (T c (M 0.84375 0 0 0))) (M 1 0 0 0)) 1 (M 1 0 0 0)) 1 (M 1 0 0 0)) (M 1 0 -7.1875 0))",
+    "(T (repeat (repeat (T (C (T (T (repeat (T l (M 1 0 -0.5 1.20711)) 8 (M 1 0.785398 0 0)) (M 0.75 0 0 0)) (M 1.125 0 0 0)) (T c (M 0.84375 0 0 0))) (M 1 0 0 0)) 2 (M 1 0 14.375 0)) 1 (M 1 0 0 0)) (M 1 0 -7.1875 0))",
+    "(C (C (T (T (repeat (repeat (T (C (T (T (r_s 13.5 3) (M 1 0 0 0)) (M 1 0 0 0)) (T (repeat (repeat (T (T c (M 0.84375 0 0 0)) (M 1 0 0 0)) 2 (M 1 0 5.625 0)) 1 (M 1 0 0 0)) (M 1 0 -2.8125 0))) (M 1 0 0 1.5)) 1 (M 1 0 0 0)) 2 (M 1 0 0 3.75)) (M 1 0 0 0)) (M 1 0 0 0.75)) (T (T (r_s 15 8.25) (M 1 0 0 4.125)) (M 1 0 0 0))) (T (repeat (repeat (T (T (T (T l (M 1 0 -0.5 0)) (M 1 1.5708 0 0)) (M 3 0 0 0)) (M 1 0 0 -1.5)) 3 (M 1 0 7.5 0)) 1 (M 1 0 0 0)) (M 1 0 -7.5 0)))",
+    "(C (C (T (T (repeat (repeat (T (C (T (T (r_s 13.5 3) (M 1 0 0 0)) (M 1 0 0 0)) (T (repeat (repeat (T (T c (M 0.84375 0 0 0)) (M 1 0 0 0)) 2 (M 1 0 5.625 0)) 1 (M 1 0 0 0)) (M 1 0 -2.8125 0))) (M 1 0 0 1.5)) 1 (M 1 0 0 0)) 2 (M 1 0 0 3.75)) (M 1 0 0 0)) (M 1 0 0 0.75)) (T (T (r_s 15 8.25) (M 1 0 0 4.125)) (M 1 0 0 0))) (T (repeat (repeat (T (T (T (T l (M 1 0 -0.5 0)) (M 1 1.5708 0 0)) (M 6 0 0 0)) (M 1 0 0 -3)) 2 (M 1 0 15 0)) 1 (M 1 0 0 0)) (M 1 0 -7.5 0)))",
+    "(C (T (repeat (repeat (T (C (T (T (r_s 9 4) (M 1 0 0 0)) (M 1 0 0 0)) (T (repeat (repeat (T (C (T (T (repeat (T l (M 1 0 -0.5 1.20711)) 8 (M 1 0.785398 0 0)) (M 0.75 0 0 0)) (M 1.125 0 0 0)) (T c (M 0.84375 0 0 0))) (M 1 0 0 0)) 2 (M 1 0 2.875 0)) 1 (M 1 0 0 0)) (M 1 0 -1.4375 0))) (M 1 0 0 0)) 1 (M 1 0 0 0)) 4 (M 1 0 0 5)) (M 1 0 0 -7.5)) (T (T (r_s 11 21) (M 1 0 0 0)) (M 1 0 0 0)))",
+    "(C (C (T (T (repeat (repeat (T (C (T (T (r_s 13.5 3) (M 1 0 0 0)) (M 1 0 0 0)) (T (repeat (repeat (T (T r (M 0.84375 0 0 0)) (M 1 0 0 0)) 2 (M 1 0 5.625 0)) 1 (M 1 0 0 0)) (M 1 0 -2.8125 0))) (M 1 0 0 1.5)) 1 (M 1 0 0 0)) 2 (M 1 0 0 3.75)) (M 1 0 0 0)) (M 1 0 0 0.75)) (T (T (r_s 15 8.25) (M 1 0 0 4.125)) (M 1 0 0 0))) (T (repeat (repeat (T (T (T (T l (M 1 0 -0.5 0)) (M 1 1.5708 0 0)) (M 6 0 0 0)) (M 1 0 0 -3)) 4 (M 1 0 5 0)) 1 (M 1 0 0 0)) (M 1 0 -7.5 0)))"
+]

src/bin/benchmark_cogsci.rs

@@ -12,13 +12,15 @@ struct Args {
     /// Number of repetitions
     #[clap(short, long, default_value = "1")]
     count: usize,
+    #[clap(short, long, default_value = "false")]
+    smc: bool,
 }
 
 fn main() {
     let args = Args::parse();
     let mut geomeans = Vec::with_capacity(args.count);
     for _ in 0..args.count {
-        let geo = benchmark_cogsci_geomean();
+        let geo = benchmark_cogsci_geomean(args.smc);
         println!("{geo:.2}");
         geomeans.push(geo);
     }
@@ -28,7 +30,7 @@ fn main() {
     }
 }
 
-fn benchmark_cogsci_geomean() -> f64 {
+fn benchmark_cogsci_geomean(smc: bool) -> f64 {
     let cogsci_dir = PathBuf::from("data/cogsci");
     // Check if directory exists
     if !cogsci_dir.exists() {
@@ -54,8 +56,15 @@ fn benchmark_cogsci_geomean() -> f64 {
     }
     // Configuration for compression: -a3 -i10
     let mut cfg = MultistepCompressionConfig::default();
-    cfg.step.max_arity = 3;  // -a3
-    cfg.iterations = 10;     // -i10
+    if smc {
+        cfg.step.smc = true;
+        cfg.step.smc_particles = 1_000;
+        cfg.step.smc_extra_steps = 25;
+        cfg.iterations = 3;     // -i3
+    } else {
+        cfg.step.max_arity = 3;  // -a3
+        cfg.iterations = 10;     // -i10
+    }
     cfg.silent = true;       // Reduce output noise
     let mut compression_times = Vec::new();
     for file_path in json_files.iter() {

src/compression.rs

@@ -232,10 +232,43 @@ pub struct CompressionStepConfig {
     #[clap(long)]
     pub quiet: bool,
 
-    // Fused lambda tags
+    /// Fused lambda tags
     #[clap(long, value_parser = clap::value_parser!(FusedLambdaTags), default_value="")]
     pub fused_lambda_tags: FusedLambdaTags,
 
+    /// If true, we will use Sequential Monte Carlo (SMC) to sample patterns
+    #[clap(long)]
+    pub smc: bool,
+
+    /// Seed for the random number generator used in SMC
+    #[clap(long, default_value = "0")]
+    pub seed: u64,
+
+    /// Number of particles to use in SMC
+    #[clap(long, default_value = "1000")]
+    pub smc_particles: usize,
+
+    /// If true, we will use "fast utility" in SMC, which is a heuristic that
+    /// estimates the utility of a pattern without fully rewriting it.
+    #[clap(long)]
+    pub smc_fast_utility: bool,
+
+    /// Number of smc steps we will run after finding a best pattern
+    /// before stopping. e.g., if --smc-extra-step=50, we will stop at
+    /// step 69 if a best pattern is found at step 19.
+    #[clap(long, default_value = "50")]
+    pub smc_extra_steps: usize,
+
+    /// Number of expansions to make before recalculating the utility of a pattern
+    /// in SMC. This is used to control the tradeoff between exploration and exploitation.
+    #[clap(long, default_value = "1")]
+    pub smc_expand_per_step: usize,
+
+    /// SMC temperature, used to control the exploration-exploitation tradeoff.
+    /// A higher temperature means more exploration, while a lower temperature means more exploitation.
+    #[clap(long, default_value = "1.0")]
+    pub smc_temperature: f32,
+
     /// TDFA settings
     #[clap(flatten)]
     pub tdfa: TDFAConfig,
@@ -510,6 +543,28 @@ impl Pattern {
             tracked: follow.is_some(),
         }
     }
+
+    #[allow(clippy::too_many_arguments)]
+    fn single_var(corpus_span: &Span, cost_fn: &ExprCost, cost_of_node_all: &[Cost], num_paths_to_node: &[Cost], tdfa_global_annotations: &Option<TDFAGlobalAnnotations>, set: &ExprSet, cfg: &CompressionStepConfig, follow: &Option<Invention>) -> Self {
+        let mut pattern = Self::single_hole(corpus_span, cost_fn, cost_of_node_all, num_paths_to_node, tdfa_global_annotations, set, cfg, follow);
+        let hole_zid = pattern.holes.pop().unwrap();
+        pattern.pattern_args.add_variable_at(hole_zid, 0);
+        pattern
+    }
+
+    pub fn single_var_from_shared(shared: &SharedData) -> Self {
+        Self::single_var(
+            &shared.corpus_span,
+            &shared.cost_fn,
+            &shared.cost_of_node_all,
+            &shared.num_paths_to_node,
+            &shared.tdfa_global_annotations,
+            &shared.set,
+            &shared.cfg,
+            &None, // no follow in single var, since this is only used by SMC
+        )
+    }
+
     /// convert pattern to an Expr
     fn to_expr(&self, shared: &SharedData) -> ExprOwned {
         let mut set = ExprSet::empty(Order::ChildFirst, false, false);
@@ -1172,7 +1227,7 @@ pub struct FinishedPattern {
 
 impl FinishedPattern {
     //#[inline(never)]
-    fn new(pattern: Pattern, shared: &SharedData) -> Self {
+    pub fn new(pattern: Pattern, shared: &SharedData) -> Self {
         let arity = pattern.pattern_args.arity();
         let usages = pattern.match_locations.iter().map(|loc| shared.num_paths_to_node[*loc]).sum();
         let compressive_utility = compressive_utility(&pattern,shared);
@@ -1380,7 +1435,7 @@ pub struct CompressionStepResult {
 }
 
 impl CompressionStepResult {
-    fn new(done: FinishedPattern, inv_name: &str, shared: &mut SharedData, very_first_cost: Cost, name_mapping: &[(String,String)], dc_comparison_millis: Option<usize>) -> Self {
+    pub fn new(done: FinishedPattern, inv_name: &str, shared: &mut SharedData, very_first_cost: Cost, name_mapping: &[(String,String)], dc_comparison_millis: Option<usize>) -> Self {
 
         let inv = done.to_invention(inv_name, shared);
         let rewritten = rewrite_fast(&done, shared, &Node::Prim(inv.name.clone().into()), &shared.cost_fn);
@@ -1504,7 +1559,7 @@ fn utility_upper_bound(
 /// to changes in size that come from rewriting with an invention. Currently this is just the
 /// size of the abstraction itself
 //#[inline(never)]
-fn noncompressive_utility(
+pub fn noncompressive_utility(
     body_utility: Cost,
     cfg: &CompressionStepConfig,
 ) -> Cost {
@@ -1580,7 +1635,7 @@ pub fn compressive_utility_from_marginals(
 }
 
 //#[inline(never)]
-fn compressive_utility(pattern: &Pattern, shared: &SharedData) -> UtilityCalculation {
+pub fn compressive_utility(pattern: &Pattern, shared: &SharedData) -> UtilityCalculation {
 
     // * BASIC CALCULATION
     // Roughly speaking compressive utility is num_usages(invention) * size(invention), however there are a few extra
@@ -1830,19 +1885,29 @@ pub fn multistep_compression_internal(
         } else {
             (format!("{}{}", cfg.abstraction_prefix, cfg.previous_abstractions + step_results.len()), cfg.step.follow.as_ref().map(|x| Invention::from_string("inv", x, cfg.step.follow_types.clone())))
         };
-
-        // call actual compression
-        let res: Vec<CompressionStepResult> = compression_step(
-            &rewritten,
-            &inv_name,
-            &cfg,
-            &tasks,
-            &weights,
-            very_first_cost,
-            &name_mapping,
-            &step_results,
-            &follow_iter,
-        );
+        let res: Vec<CompressionStepResult> = if cfg.step.smc {
+            assert!(follow_iter.is_none());
+            smc::compression_step_smc(
+                &rewritten,
+                &inv_name,
+                &cfg,
+                &tasks,
+                &weights,
+                &step_results,
+            )
+        } else {
+            compression_step(
+                &rewritten,
+                &inv_name,
+                &cfg,
+                &tasks,
+                &weights,
+                very_first_cost,
+                &name_mapping,
+                &step_results,
+                &follow_iter,
+            )
+        };
 
         if !res.is_empty() {
             // rewrite with the invention

src/expansion.rs

@@ -4,7 +4,7 @@ use itertools::Itertools;
 use lambdas::{Idx, Node, Symbol, Tag, ZId, ZNode};
 use rustc_hash::{FxHashMap, FxHashSet};
 
-use crate::{invalid_metavar_location, Arg, Cost, LocationsForReusableArgs, Pattern, PatternArgs, SharedData, SymvarInfo, VariableType, ZIdExtension};
+use crate::{compatible_locations, invalid_metavar_location, Arg, Cost, LocationsForReusableArgs, Pattern, PatternArgs, SharedData, SymvarInfo, VariableType, ZIdExtension};
 
 /// Tells us what a hole will expand into at this node.
 #[derive(Debug, Clone, PartialEq, Eq, Hash, Ord, PartialOrd)]
@@ -71,7 +71,7 @@ impl ExpandsTo {
     #[inline]
     pub fn local_expansion_utility(&self, shared: &SharedData) -> Cost {
         let ExpandsTo(s) = self;
-        let res = match &s {
+        let res = match s {
             ExpandsToInner::Lam(_) => shared.cost_fn.cost_lam,
             ExpandsToInner::App => shared.cost_fn.cost_app,
             ExpandsToInner::Var(_, _) => shared.cost_fn.cost_var,
@@ -162,11 +162,15 @@ pub fn expands_to_of_node(node: &Node) -> ExpandsTo {
 pub fn get_syntactic_expansions(arg_of_loc: &FxHashMap<usize, Arg>, match_locations: Vec<usize>, sym_var_info: &Option<SymvarInfo>) -> Vec<(ExpandsTo, Vec<Idx>)> {
     match_locations.into_iter()
         .group_by(|loc| &arg_of_loc[loc].expands_to).into_iter()
-        .filter(|(expands_to, _)| sym_var_info.as_ref().is_none_or(|s| !expands_to.is_prim_symbol(s)))
+        .filter(|(expands_to, _)| valid_syntactic_expansion_loc(sym_var_info, expands_to))
         .map(|(expands_to, locs)| (expands_to.clone(), locs.collect::<Vec<Idx>>()))
         .collect::<Vec<_>>()
 }
 
+pub fn valid_syntactic_expansion_loc(sym_var_info: &Option<SymvarInfo>, expands_to: &ExpandsTo) -> bool {
+    sym_var_info.as_ref().is_none_or(|s| !expands_to.is_prim_symbol(s))
+}
+
 
 //#[inline(never)]
 /// Return options for what abstraction arguments (aka ivars, #i) can expand into. When expanding to an ivar that
@@ -211,6 +215,92 @@ pub fn get_ivars_expansions(original_pattern: &Pattern, arg_of_loc: &FxHashMap<I
 }
 
 
+/* Perform expansions on variables -- largely for SMC */
+
+pub fn perform_expansion_variable(
+    pattern: Pattern,
+    shared: &SharedData,
+    variable_ivar: i32,
+    expands_to: ExpandsTo,
+) -> Option<Pattern> {
+    let mut pattern = pattern;
+    let mut expands_to = expands_to;
+
+    let variable_zids: Vec<usize> = pattern.pattern_args.remove_variable_at(variable_ivar, match &mut expands_to {
+        ExpandsTo(ExpandsToInner::IVar(i, _)) => {
+            Some(i)
+        }
+        _ => None,
+
+    });
+
+    let body_utility = pattern.body_utility +  expands_to.local_expansion_utility(shared) * variable_zids.len() as Cost;
+    pattern.body_utility = body_utility;
+
+    let num_vars = pattern.pattern_args.arity() as i32;
+
+    for variable_zid in &variable_zids {
+        expands_to.add_variables(*variable_zid, &mut pattern.pattern_args);
+        expands_to.syntactic_expansion(&shared.extensions_of_zid[*variable_zid], |zid, i| {
+            pattern.pattern_args.add_variable_at(zid, num_vars + (i as i32));
+        });
+    }
+    Some (pattern)
+}
+
+
+pub fn get_num_variables(pattern: &Pattern) -> usize {
+    pattern.pattern_args.arity()
+}
+
+
+fn sample_in_range_without_value(max_range_excl: usize, avoid: i32, rng: &mut impl rand::Rng) -> Option<i32> {
+    if max_range_excl <= 1 {
+        return None; // no other variable to expand to
+    }
+    let mut new_ivar = rng.gen_range(0..max_range_excl - 1) as i32;
+    if new_ivar >= avoid {
+        new_ivar += 1; // skip the variable we are expanding
+    }
+    Some(new_ivar)
+}
+
+pub fn sample_variable_reuse_expansion(
+    pattern: &Pattern,
+    shared: &SharedData,
+    variable_ivar: i32,
+    match_location: usize,
+    rng: &mut impl rand::Rng,
+) -> Option<(Pattern, ExpandsTo)> {
+    let new_ivar = sample_in_range_without_value(get_num_variables(pattern), variable_ivar, rng)?;
+    let zid_original = pattern.pattern_args.zid_for_ivar(variable_ivar);
+    let zid_new = pattern.pattern_args.zid_for_ivar(new_ivar);
+    let arg_of_loc_original = &shared.arg_of_zid_node[zid_original];
+    let arg_of_loc_new = &shared.arg_of_zid_node[zid_new];
+    if compatible_locations(
+        shared,
+        &[match_location],
+        arg_of_loc_original,
+        arg_of_loc_new,
+        pattern.pattern_args.type_for_ivar(new_ivar)
+    ).is_empty() {
+        return None; // no compatible locations for the new ivar
+    }
+
+    let locs = compatible_locations(
+        shared,
+        &pattern.match_locations,
+        arg_of_loc_original,
+        arg_of_loc_new,
+        pattern.pattern_args.type_for_ivar(new_ivar)
+    );
+    assert!(!locs.is_empty());
+    let mut pattern = pattern.clone();
+    pattern.match_locations = locs;
+    let expands_to = ExpandsTo(ExpandsToInner::IVar(new_ivar,  pattern.pattern_args.type_for_ivar(new_ivar)));
+    Some((pattern, expands_to))
+}
+
 pub fn svar_locations(original_pattern: &Pattern, arg_of_loc: &FxHashMap<Idx,Arg>, reusable_locs: FxHashSet<Idx>, sym_var_info: &SymvarInfo) -> Vec<Idx> {
 
     let mut locations = vec![];

src/lib.rs

@@ -4,6 +4,7 @@ pub mod rewriting;
 pub mod egraphs;
 pub mod util;
 pub mod formats;
+pub mod smc;
 pub mod tdfa;
 pub mod expansion;
 pub mod pattern_args;
@@ -18,6 +19,7 @@ pub use {
     util::*,
     formats::*,
     lambdas::*,
+    smc::*,
     tdfa::*,
     expansion::*,
     pattern_args::*,