Custom step function

General question/request:

There are very few examples for creating a custom step function for sampling, the last one in the documentation being for pymc3 (and some of the functions it references are gone/deprecated). The topics about it on this forum also contain dead links. Is it possible to write up a new guide that covers the topic for pymc4 or pymc5?

Specific problem I’m trying to solve:

I am trying to implement the Stochastic Block Model, used to model the clustering of nodes in a network and the probabilities of edges existing that connect nodes within and between said clusters. In short, there are multiple variables whose shapes depend on the number of clusters K. This makes sampling complicated as said shapes need to change if K is changed.

The MC sampling steps in that paper are as follows:

  • MK: Metropolis move to increase or decrease K, adding or removing an empty cluster.
  • GS: Gibbs sampling on a randomly-selected node. Fixing all but one node in z, select a new cluster assignment for that node.
  • M3: Metropolis–Hastings on the labels in two clusters. This is the M3 move proposed in Nobile and Fearnside (2007). Two clusters are selected at random and the nodes are reassigned to the two clusters using a novel scheme fully described in that paper. K is not affected by this move.
  • AE: The absorb–eject move is a Metropolis–Hasting merge/split cluster move, as described in Nobile and Fearnside (2007). This move does affect K along with z.

Are there any resources (that maybe I missed) that would help me with implementing these steps? Or could you point me in the direction of which functions/what part of the documentation I need to look at so I can implement these sampling steps? So far I’ve read and tried to replicate this old documentation article in order to get the hang of custom step functions but I’m afraid I do not fully understand the example in the documentation (e.g. the objects/functions stick_breaking and BlockedStep and why they are used vs their alternatives).

CC @junpenglao

The link is: Compound Steps in Sampling — PyMC3 3.11.4 documentation - should be straightforward to upgrade to v4

One thing I had to change upon upgrade from v4 to v5, is to return point, [] instead of just return point.