Separate models into one-per-file

Author: penelopeysm

README.md

@@ -12,14 +12,11 @@ You can modify the list of AD types in `main.jl`.
 
 ## I want to add more models!
 
-You can modify the list of models in `models.jl`.
+You can modify the list of models by adding a new file inside the `models` directory.
+This file should contain the model definition, and call the `@register` macro to register the model with the `ADTests` package.
+See the existing files in that directory for examples.
 
-Note that if you want the model definition to be shown on the website, your model definitions must be of the 'standard' form `@model function name() ... end`.
-This means that:
-- One-liner function definitions like `@model f(x) = ...`will not work.
-- Fancy metaprogramming tricks to generate a family of models at one go (like [this old code](https://github.com/TuringLang/ADTests/blob/266d7ab85fea2e01e7e05af6cee179d7f6200b0f/models.jl#L108-L129)) will not work.
-
-To understand why, see the `get_model_definition` function in `ad.py`.
+To make sure that the definition is included in the final website, you will have to make sure that the filename is the same as the model name (i.e. a model `@model function f()` is in `models/f.jl`).
 
 ## I want to edit the website!
 

ad.py

@@ -102,26 +102,8 @@ def run_ad(args):
 
 
 def get_model_definition(model_key):
-    """Get the model definition from the Julia script."""
-    lines = []
-    submodels = []
-    record = False
-    with open("models.jl", "r") as file:
-        for line in file:
-            line = line.rstrip()
-            if line.startswith(f"@model function {model_key}"):
-                record = True
-            if record:
-                lines.append(line)
-
-                if "to_submodel" in line:
-                    submodel_name = line.split("to_submodel(")[1].split("(")[0]
-                    submodels.append(submodel_name)
-                if line == "end":
-                    break
-    for submodel in submodels:
-        lines = [get_model_definition(submodel), *lines]
-    return "\n".join(lines)
+    """Get the model definition from the file that contains it."""
+    return Path(f"models/{model_key}.jl").read_text().strip()
 
 
 def html(_args):

main.jl

@@ -23,9 +23,44 @@ ADTYPES = Dict(
     "Zygote" => AutoZygote(),
 )
 
-# Models to test.
-include("models.jl")
-using .Models: MODELS
+# Models to test. The convention is that:
+#    a, b, c, ... are assumed variables
+#    x, y, z, ... are observed variables
+# although it's hardly a big deal.
+MODELS = Dict{String,DynamicPPL.Model}()
+macro register(model)
+    :(MODELS[string($(esc(model)).f)] = $(esc(model)))
+end
+
+# These imports tend to get used a lot in models
+using DynamicPPL: @model
+using Distributions
+using LinearAlgebra: I
+
+include("models/assume_dirichlet.jl")
+include("models/assume_lkjcholu.jl")
+include("models/assume_mvnormal.jl")
+include("models/assume_normal.jl")
+include("models/assume_submodel.jl")
+include("models/assume_wishart.jl")
+include("models/control_flow.jl")
+include("models/dot_assume_observe_index.jl")
+include("models/dot_assume.jl")
+include("models/dot_observe.jl")
+include("models/dynamic_constraint.jl")
+include("models/models.jl")
+include("models/multiple_constraints_same_var.jl")
+include("models/multithreaded.jl")
+include("models/n010.jl")
+include("models/n050.jl")
+include("models/n100.jl")
+include("models/n500.jl")
+include("models/observe_index.jl")
+include("models/observe_literal.jl")
+include("models/observe_multivariate.jl")
+include("models/observe_submodel.jl")
+include("models/pdb_eight_schools_centered.jl")
+include("models/pdb_eight_schools_noncentered.jl")
 
 # The entry point to this script itself begins here
 if ARGS == ["--list-model-keys"]

models.jl

@@ -0,0 +1,5 @@
+@model function assume_beta()
+    a ~ Beta(2, 2)
+end
+
+@register assume_beta()

models/assume_beta.jl

@@ -0,0 +1,5 @@
+@model function assume_dirichlet()
+    a ~ Dirichlet([1.0, 5.0])
+end
+
+@register assume_dirichlet()

models/assume_dirichlet.jl

@@ -0,0 +1,5 @@
+@model function assume_lkjcholu()
+    a ~ LKJCholesky(5, 1.0, 'U')
+end
+
+@register assume_lkjcholu()

models/assume_lkjcholu.jl

@@ -0,0 +1,5 @@
+@model function assume_mvnormal()
+    a ~ MvNormal([0.0, 0.0], [1.0 0.5; 0.5 1.0])
+end
+
+@register assume_mvnormal()

models/assume_mvnormal.jl

@@ -0,0 +1,5 @@
+@model function assume_normal()
+    a ~ Normal()
+end
+
+@register assume_normal()

models/assume_normal.jl

@@ -0,0 +1,9 @@
+@model function inner1()
+    return a ~ Normal()
+end
+@model function assume_submodel()
+    a ~ to_submodel(inner1())
+    x ~ Normal(a)
+end
+
+@register assume_submodel()