Skip to content

Conversation

@jack-wu05
Copy link

Here is the tree reference I worked on with Trevor for the summer. I also implemented a dense GaussianReference.

A quick summary of what the tree does: Essentially, we fit each variable with a gaussian. We then construct a complete tree with the variables, where each edge is weighted with the mutual information between the two variables. Note that because direction is needed, instead of using one undirected arrow between two variables I pushed the edges so that there are 2 edges: one going each way; mutual information is symmetric so the weights are the same. I then run a max spanning tree algorithm to get the max spanning tree, and store the edge connections. We use these edge connections, in the form of conditional distributions, later on for sampling and density evaluation. The "which_variable" and "which_index" structures I have in the struct def are to keep track of which component of which variable a given node represents, so that component can be appropriately updated. In hindsight, this may not be totally necessary since the only variable is :singleton_variable and the nodes are labelled with integers throughout the code pushed in order, so you may be able to update directly. I'll leave this part for now.

Please note a couple things:

  1. I had to integrate the covariance matrix into a new recorder as this wasn't supported before, hence all the changes to the recorder codes. Please note that I took my new _transformed_online_full recorder and replaced the GaussianReference recorder with it as well, so _transformed_online was totally removed. I wasn't sure how to dispatch this otherwise.

  2. I have not yet tested the gradient nor kernel invariance (which would test the sample_iid! and density evaluation of my new reference), this still needs to be done. I have already tested the correctness of the max spanning tree algorithm.

  3. I made sure to generate randomly from replica.rng, which should preserve SPI.

Here are some promising results from a 601-dimensional HMM test case (which I also added to the codebase) that make this update (and all the effort for full integration) worth pursuing:

GCB-r restarts-r restarts-time tau-r 601-diag 601-tree

Of course, feel free to ask me any questions if any part is confusing. And if my implementation is too poor, please do take inspiration and make an improved version :)

@jack-wu05
Copy link
Author

Oops I think I PRed from the wrong branch (should be main). Will send another later if this PR is seriously considered.

@trevorcampbell
Copy link
Collaborator

trevorcampbell commented Aug 21, 2025

@jack-wu05 PRing from your new-references branch is the right move, so all good. We'll review this shortly :-)

@alexandrebouchard
Copy link
Member

Thanks Jack, really excellent work. I have triggered the CI for now, feel free to ping me during or after the workshop next week.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants