add model architecture using mermaid

Author: florence-bockting

.pre-commit-config.yaml

@@ -11,13 +11,14 @@ repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: "v6.0.0"
     hooks:
-      #- id: check-added-large-files
+      - id: check-added-large-files
       - id: check-ast
       - id: check-case-conflict
       - id: check-json
       - id: check-merge-conflict
       - id: check-symlinks
       - id: check-yaml
+        exclude: ^mkdocs\.yml$ # otherwise !!python/name:mermaid2.fence_mermaid in .mkdocs.yaml make it fail
       - id: debug-statements
       - id: detect-private-key
       - id: end-of-file-fixer

docs/NAVIGATION.md

@@ -9,6 +9,7 @@ See https://oprypin.github.io/mkdocs-literate-nav/
 - [Tutorials](tutorials/index.md)
 - [Further background](further-background/index.md)
     - [Dependency pinning and testing](further-background/dependency-pinning-and-testing.md)
+    - [Model architecture](further-background/model-architecture.md)
 - [Development](development.md)
 - [API reference](api/pymc_elicito/)
 - [Changelog](changelog.md)

docs/further-background/index.md

@@ -11,5 +11,4 @@ Points we will aim to cover:
 - Why it was created
 - Help the reader make connections
 
-We will aim to avoid writing instructions or technical descriptions here,
-they belong elsewhere.
+We will aim to avoid writing instructions or technical descriptions here, they belong elsewhere.

docs/further-background/model-architecture.md

@@ -0,0 +1,66 @@
+```mermaid
+flowchart TB
+    subgraph eliobj["Elicit class (user input)"]
+            input["Model, Parameters, Targets, \n Network, Initializer, Optimizer, \n Meta-Setting"]
+            expert_dat["expert elicited summaries"]
+    end
+    subgraph initialization["initialization"]
+            initializer["initializer"]
+            hyper_parametric("prior parameters")
+            hyper_deep("weights/biases of DNNs")
+    end
+
+    subgraph subGraph3["priors"]
+            transformation["transform to constrained space"]
+            prior["sample from prior in \n unconstrained space"]
+    end
+    subgraph model["model"]
+            generative_model["run generative model \n in forward mode"]
+    end
+    subgraph summaries["summaries"]
+            elicits["elicited summaries"]
+            summary["target quantities"]
+    end
+    subgraph loss["loss"]
+            total_loss["total loss"]
+            indiv_loss["individual losses"]
+    end
+    subgraph optimization["optimization"]
+            gradients["gradients"]
+    end
+    subgraph training["training"]
+        subGraph3
+        prior_samples[/"prior samples"/]
+        model
+        model_output[/"model simulations"/]
+        summaries
+        simulated_summaries[/"simulated summaries"/]
+        loss
+        optimization
+        check{"convergence \n criterion"}
+        train_vars[/"trainable variables \n (=hyperparameters)"/]
+    end
+    input --> initializer
+    initializer -- if parametric prior --> hyper_parametric
+    initializer -- if deep prior --> hyper_deep
+    hyper_parametric --> train_vars
+    hyper_deep --> train_vars
+    train_vars --> prior
+    expert_dat --> indiv_loss & input
+    prior --> transformation
+    transformation --> prior_samples
+    prior_samples --> generative_model
+    generative_model --> model_output
+    model_output --> summary
+    summary --> elicits
+    elicits --> simulated_summaries
+    simulated_summaries --> indiv_loss
+    indiv_loss --> total_loss
+    total_loss --> gradients
+    gradients --> check
+    check -- not reached:<br> update --> train_vars
+    check -- reached --> stop((("stop")))
+
+    classDef data shape: lean-r
+    class prior_samples data
+```

mkdocs.yml

@@ -28,6 +28,7 @@ theme:
         name: Switch to light mode
 
 plugins:
+  - mermaid2
   # https://mkdocstrings.github.io/autorefs/
   - autorefs
   # Required for auto-generating our documentation stubs
@@ -121,6 +122,11 @@ plugins:
   - section-index
 
 markdown_extensions:
+  - pymdownx.superfences:
+      custom_fences:
+        - name: mermaid
+          class: mermaid
+          format: !!python/name:mermaid2.fence_mermaid
   # https://squidfunk.github.io/mkdocs-material/setup/extensions/python-markdown/#attribute-lists
   - attr_list
   - footnotes

pyproject.toml

@@ -96,6 +96,7 @@ dev = [
 docs = [
     # Key dependencies
     # ----------------
+    "mkdocs-mermaid2-plugin>=1.2.3",
     "attrs==25.3.0",
     "mkdocs-autorefs==1.4.2",
     "mkdocs-gen-files==0.5.0",

requirements-docs-locked.txt

@@ -25,6 +25,7 @@ cycler==0.12.1
 debugpy==1.8.11
 decorator==5.1.1
 defusedxml==0.7.1
+editorconfig==0.17.1
 exceptiongroup==1.3.0 ; python_full_version < '3.11'
 executing==2.1.0
 fastjsonschema==2.21.1
@@ -46,6 +47,7 @@ ipython-pygments-lexers==1.1.1 ; python_full_version >= '3.11'
 isoduration==20.11.0
 jedi==0.19.2
 jinja2==3.1.5
+jsbeautifier==1.15.4
 json5==0.10.0
 jsonpointer==3.0.0
 jsonschema==4.23.0
@@ -80,6 +82,7 @@ mkdocs-jupyter==0.25.1
 mkdocs-literate-nav==0.6.2
 mkdocs-material==9.6.16
 mkdocs-material-extensions==1.3.1
+mkdocs-mermaid2-plugin==1.2.3
 mkdocs-section-index==0.3.10
 mkdocstrings==0.30.0
 mkdocstrings-python==1.16.12