Tool for merging Stan models

I have been working on a small tool, stanmerge (GitHub - herbps10/stanmerge: Syntax-aware transformation and merging of Stan models), that can merge multiple Stan models in a syntax-aware way. I wanted to share it here in case others may find it useful.

The basic usage is that you can feed multiple Stan files into stanmerge, and it will merge them program block by program block.

There are two main use cases:

  • In my research, I often end up with multiple versions of a model that slightly differ from one another – for example, each version might test a different prior structure for a parameter. To avoid duplicating code, it would be convenient to have the common part of the model separated from the parts that are different. Then, files can be merged to form a final model. See the README for an example of this, where two response distribution setups for a model are stored as separate files and merged individually with a base model specification.
  • Some models may repeat a particular model structure multiple times. For example, we may want to apply the same complex prior to multiple parameters. Rather than duplicate this code multiple times, we can factor out the model component as a separate file, then merge it multiple times into the main model. The Gaussian Process example illustrates this, where the same Gaussian Process prior structure is applied to two model parameters without having to duplicate code.

This all works by manipulating ASTs via the stanc3 compiler – big thanks to all the contributors for making it easy to work with.

More generally, I’m thinking of this as a way to explore modularizing Stan development by thinking of models as combinations of model components. Feedback and suggestions would be welcome!

6 Likes

Hi Herb - cool! Totally agree with the motivation. I actually have a project with some overlapping use cases; here’s a blog post: Drawing Maps of Model Space with Modular Stan | Statistical Modeling, Causal Inference, and Social Science and paper: [2208.06329] Multi-Model Probabilistic Programming. I’d love to hear your thoughts.

4 Likes

This looks really great (as did your work @rybern - any news on development there?). Really keen to see this progress as it feels like it would open up a lot of functionality for model development. No actual useful input I am afraid!

Some waffle about where I would like to use this kind of functionality

We have a use case for this with modelling outbreaks of infectious disease where we have many modular models and we need to combine them in different ways depending on the outbreak setting and data available (i.e cases, deaths, serology etc). Our current implementation (epinowcast/epinowcast.stan at c36d397e67ef17deb8843008b5b2d44298f8f5ae · epinowcast/epinowcast · GitHub) is unwieldy to say the least and will only get more so as we add modules. Either of your solutions would be very useful I think.

In your implementation @hsusmann I really like the use of generic variable names to enable chunks to be reused. We mostly have to go with functions for this kind of use case at the moment but it can be quite limiting.

This is very cool. It’s not directly related to the idea of merging, but I’d love to hear what you thought of using stanc as a library @hsusmann. Allowing for tools like this was one of the original goals of the compiler re-write, but to my knowledge this is one of the first times someone actually built on top of stanc to do something like this

@rybern this looks great, thanks for sharing! I was starting to think about what a module syntax might look like, and your paper fills in all the details.

@seabbs yes I think the idea of rewriting generic variable names is pretty powerful. In the long run it would probably be better to have a module syntax in Stan (like the one @rybern proposes), but stanmerge provides a decent intermediate solution I think.

@WardBrian I don’t have too much feedback at the moment about using stanc3, mostly because it was pretty painless – for the most part I was able to just look in the code base for examples of how things are done (like how to go from a file name to an AST, or how to pretty-print an AST), and then tweak things as needed. If I run into anything I will let you know, though. My next task is a bit more involved, as I’d like to be able to correctly handle comments, which might end up requiring a modified pretty-printer.

1 Like

Comments are really tricky. They were basically the last thing we got implemented in the pretty printer, and as you’ve noticed they still don’t work well under manipulation.

This isn’t a unique problem. I found a lot of insight in the Go language issue on the same thing go/ast: Free-floating comments are single-biggest issue when manipulating the AST · Issue #20744 · golang/go · GitHub

1 Like