Allow users to specify prior process for row and column partitions

book/lace_preprocess_mdbook_yaml/src/lib.rs

@@ -1,5 +1,3 @@
-use std::collections::HashMap;
-
 use anyhow::anyhow;
 use log::debug;
 use mdbook::{
@@ -12,7 +10,15 @@ use regex::Regex;
 
 use serde::Deserialize;
 
-type GammaMap = HashMap<String, lace_stats::rv::dist::Gamma>;
+#[derive(Deserialize)]
+struct ViewPriorProcess {
+    pub view_prior_process: lace_codebook::PriorProcess,
+}
+
+#[derive(Deserialize)]
+struct StatePriorProcess {
+    pub state_prior_process: lace_codebook::PriorProcess,
+}
 
 fn check_deserialize_yaml<T>(input: &str) -> anyhow::Result<()>
 where
@@ -54,17 +60,14 @@ macro_rules! check_deserialize_arm {
     }
 }
 
-fn check_deserialize_dyn(
-    input: &str,
-    type_name: &str,
-    format: &str,
-) -> anyhow::Result<()> {
+fn check_deserialize_dyn(input: &str, type_name: &str, format: &str) -> anyhow::Result<()> {
     check_deserialize_arm!(
         input,
         type_name,
         format,
         [
-            GammaMap,
+            ViewPriorProcess,
+            StatePriorProcess,
             lace_codebook::ColMetadata,
             lace_codebook::ColMetadataList
         ]
@@ -80,35 +83,21 @@ impl YamlTester {
         YamlTester
     }
 
-    fn examine_chapter_content(
-        &self,
-        content: &str,
-        re: &Regex,
-    ) -> anyhow::Result<()> {
+    fn examine_chapter_content(&self, content: &str, re: &Regex) -> anyhow::Result<()> {
         let parser = Parser::new(content);
         let mut code_block = Some(String::new());
 
         for event in parser {
             match event {
-                Event::Start(Tag::CodeBlock(CodeBlockKind::Fenced(
-                    ref code_block_string,
-                ))) => {
+                Event::Start(Tag::CodeBlock(CodeBlockKind::Fenced(ref code_block_string))) => {
                     if re.is_match(code_block_string) {
-                        debug!(
-                            "YAML Block Start, identifier string={}",
-                            code_block_string
-                        );
+                        debug!("YAML Block Start, identifier string={}", code_block_string);
                         code_block = Some(String::new());
                     }
                 }
-                Event::End(Tag::CodeBlock(CodeBlockKind::Fenced(
-                    ref code_block_string,
-                ))) => {
+                Event::End(Tag::CodeBlock(CodeBlockKind::Fenced(ref code_block_string))) => {
                     if let Some(captures) = re.captures(code_block_string) {
-                        debug!(
-                            "Code Block End, identifier string={}",
-                            code_block_string
-                        );
+                        debug!("Code Block End, identifier string={}", code_block_string);
 
                         let serialization_format = captures
                             .get(1)
@@ -119,21 +108,13 @@ impl YamlTester {
                             .get(2)
                             .ok_or(anyhow!("No deserialize type found"))?
                             .as_str();
-                        debug!(
-                            "Target deserialization type is {}",
-                            target_type
-                        );
+                        debug!("Target deserialization type is {}", target_type);
 
                         let final_block = code_block.take();
-                        let final_block =
-                            final_block.ok_or(anyhow!("No YAML text found"))?;
+                        let final_block = final_block.ok_or(anyhow!("No YAML text found"))?;
                         debug!("Code block ended up as\n{}", final_block);
 
-                        check_deserialize_dyn(
-                            &final_block,
-                            target_type,
-                            serialization_format,
-                        )?;
+                        check_deserialize_dyn(&final_block, target_type, serialization_format)?;
                     }
                 }
                 Event::Text(ref text) => {
@@ -154,11 +135,7 @@ impl Preprocessor for YamlTester {
         "lace-yaml-tester"
     }
 
-    fn run(
-        &self,
-        _ctx: &PreprocessorContext,
-        book: Book,
-    ) -> anyhow::Result<Book> {
+    fn run(&self, _ctx: &PreprocessorContext, book: Book) -> anyhow::Result<Book> {
         debug!("Starting the run");
         let re = Regex::new(r"^(yaml|json).*,deserializeTo=([^,]+)").unwrap();
         for book_item in book.iter() {

book/src/SUMMARY.md

@@ -16,6 +16,7 @@
     - [Data Simulation](./pcc/simulate.md)
     - [In- and out-of-table operations](./pcc/inouttable.md)
     - [Adding data to a model](./pcc/add-data.md)
+    - [Prior processes](./pcc/prior-processes.md)
 - [Preparing your data](./data/basics.md)
 - [Codebook reference](./codebook-ref.md)
 - [Appendix](./appendix/appendix.md)

book/src/codebook-ref.md

@@ -16,30 +16,50 @@ information about
 
 String name of the table. For your reference.
 
-### `state_alpha_prior`
+### `state_prior_process`
 
-A gamma prior on the Chinese Restaurant Process (CRP) alpha parameter assigning
-columns to views.
+The prior process used for assigning columns to views. Can either be a Dirichlet process with a Gamma prior on alpha:
 
-Example with a gamma prior
+```yaml,deserializeTo=StatePriorProcess
+state_prior_process: !dirichlet
+  alpha_prior:
+    shape: 1.0
+    rate: 1.0
+```
+
+or a Pitman-Yor process with a Gamma prior on alpha and a Beta prior on d.
 
-```yaml,deserializeTo=GammaMap
-state_alpha_prior:
-  shape: 1.0
-  rate: 1.0
+```yaml,deserializeTo=StatePriorProcess
+state_prior_process: !pitman_yor
+  alpha_prior:
+    shape: 1.0
+    rate: 1.0
+  d_prior:
+    alpha: 0.5
+    beta: 0.5
 ```
 
-### `view_alpha_prior`
+### `view_prior_process`
 
-A gamma prior on the Chinese Restaurant Process (CRP) alpha parameter assigning
-rows within views to categories.
+The prior process used for assigning rows to categories. Can either be a Dirichlet process with a Gamma prior on alpha:
+
+```yaml,deserializeTo=ViewPriorProcess
+view_prior_process: !dirichlet
+  alpha_prior:
+    shape: 1.0
+    rate: 1.0
+```
 
-Example with a gamma prior
+or a Pitman-Yor process with a Gamma prior on alpha and a Beta prior on d.
 
-```yaml,deserializeTo=GammaMap
-view_alpha_prior:
-  shape: 1.0
-  rate: 1.0
+```yaml,deserializeTo=StatePriorProcess
+view_prior_process: !pitman_yor
+  alpha_prior:
+    shape: 1.0
+    rate: 1.0
+  d_prior:
+    alpha: 0.5
+    beta: 0.5
 ```
 
 ### `col_metadata`

book/src/pcc/prior-processes.md

@@ -0,0 +1,12 @@
+# Prior Processes
+
+In Lace (and in Bayesian nonparametrics) we put a prior on the number of parameters. This *prior process* formalizes how instances are distributed to an unknown number of categories. Lace gives you two options
+
+- The one-parameter Dirichlet process, `DP(α)`
+- The two-parameter Pitman-Yor process, `PYP(α, d)`
+
+The Dirichlet process more heavily penalizes new categories with an exponential fall off while the Pitman-Yor process has a power law fall off in the number for categories. When d = 0, Pitman-Yor is equivalent to the Dirichlet process.
+
+While Pitman-Yor may fit the data better it will create more parameters, which will cause model training to take longer.
+
+For those looking for a good introduction to prior process, [this slide deck](https://www.gatsby.ucl.ac.uk/~ywteh/teaching/probmodels/lecture5bnp.pdf) from Yee Whye Teh is a good resource.

book/src/workflow/model.md

@@ -219,7 +219,7 @@ engine.update(
     save_path="mydata.lace",
     transitions=[
         StateTransition.row_assignment(RowKernel.slice()),
-        StateTransition.view_alphas(),
+        StateTransition.view_prior_process_params(),
     ],
 )
 ```

cli/README.md

@@ -1 +1 @@
-# Lace CLI
+# Lace CLI

cli/resources/datasets/animals/codebook.yaml

@@ -1,10 +1,12 @@
-table_name: my_table
-state_alpha_prior:
-  shape: 1.0
-  rate: 1.0
-view_alpha_prior:
-  shape: 1.0
-  rate: 1.0
+table_name: animals
+state_prior_process: !dirichlet
+  alpha_prior:
+    shape: 1.0
+    rate: 1.0
+view_prior_process: !dirichlet
+  alpha_prior:
+    shape: 1.0
+    rate: 1.0
 col_metadata:
 - name: black
   coltype: !Categorical

cli/resources/datasets/satellites/codebook.yaml

@@ -1,10 +1,12 @@
-table_name: my_table
-state_alpha_prior:
-  shape: 1.0
-  rate: 1.0
-view_alpha_prior:
-  shape: 1.0
-  rate: 1.0
+table_name: satellites
+state_prior_process: !dirichlet
+  alpha_prior:
+    shape: 1.0
+    rate: 1.0
+view_prior_process: !dirichlet
+  alpha_prior:
+    shape: 1.0
+    rate: 1.0
 col_metadata:
 - name: Country_of_Operator
   coltype: !Categorical

cli/src/opt.rs

@@ -21,8 +21,8 @@ pub enum Transition {
     ColumnAssignment,
     ComponentParams,
     RowAssignment,
-    StateAlpha,
-    ViewAlphas,
+    StatePriorProcessParams,
+    ViewPriorProcessParams,
     FeaturePriors,
 }
 
@@ -33,8 +33,8 @@ impl std::str::FromStr for Transition {
         match s {
             "column_assignment" => Ok(Self::ColumnAssignment),
             "row_assignment" => Ok(Self::RowAssignment),
-            "state_alpha" => Ok(Self::StateAlpha),
-            "view_alphas" => Ok(Self::ViewAlphas),
+            "state_prior_process_params" => Ok(Self::StatePriorProcessParams),
+            "view_prior_process_params" => Ok(Self::ViewPriorProcessParams),
             "feature_priors" => Ok(Self::FeaturePriors),
             "component_params" => Ok(Self::ComponentParams),
             _ => Err(format!("cannot parse '{s}'")),
@@ -142,17 +142,17 @@ impl RunArgs {
         let transitions = match self.transitions {
             None => vec![
                 StateTransition::ColumnAssignment(col_alg),
-                StateTransition::StateAlpha,
+                StateTransition::StatePriorProcessParams,
                 StateTransition::RowAssignment(row_alg),
-                StateTransition::ViewAlphas,
+                StateTransition::ViewPriorProcessParams,
                 StateTransition::FeaturePriors,
             ],
             Some(ref ts) => ts
                 .iter()
                 .map(|t| match t {
                     Transition::FeaturePriors => StateTransition::FeaturePriors,
-                    Transition::StateAlpha => StateTransition::StateAlpha,
-                    Transition::ViewAlphas => StateTransition::ViewAlphas,
+                    Transition::StatePriorProcessParams => StateTransition::StatePriorProcessParams,
+                    Transition::ViewPriorProcessParams => StateTransition::ViewPriorProcessParams,
                     Transition::ComponentParams => StateTransition::ComponentParams,
                     Transition::RowAssignment => StateTransition::RowAssignment(row_alg),
                     Transition::ColumnAssignment => StateTransition::ColumnAssignment(col_alg),
@@ -264,9 +264,6 @@ pub struct CodebookArgs {
     /// Parquet input filename
     #[clap(long = "parquet", group = "src")]
     pub parquet_src: Option<PathBuf>,
-    /// CRP alpha prior on columns and rows
-    #[clap(long, default_value = "1.0, 1.0")]
-    pub alpha_prior: GammaParameters,
     /// Codebook out. May be either json or yaml
     #[clap(name = "CODEBOOK_OUT")]
     pub output: PathBuf,
@@ -343,16 +340,16 @@ mod tests {
     #[test]
     fn view_alphas_from_str() {
         assert_eq!(
-            Transition::from_str("view_alphas").unwrap(),
-            Transition::ViewAlphas
+            Transition::from_str("view_prior_process_params").unwrap(),
+            Transition::ViewPriorProcessParams
         );
     }
 
     #[test]
     fn state_alpha_from_str() {
         assert_eq!(
-            Transition::from_str("state_alpha").unwrap(),
-            Transition::StateAlpha
+            Transition::from_str("state_prior_process_params").unwrap(),
+            Transition::StatePriorProcessParams
         );
     }