Skip to content

Add type ChordalGMRF.#81

Open
samuelsonric wants to merge 11 commits intotimweiland:mainfrom
samuelsonric:main
Open

Add type ChordalGMRF.#81
samuelsonric wants to merge 11 commits intotimweiland:mainfrom
samuelsonric:main

Conversation

@samuelsonric
Copy link
Copy Markdown
Contributor

@samuelsonric samuelsonric commented Mar 31, 2026 •
edited
Loading

I added a new type ChordalGMRF <: AbstractGMRF, along with files for testing and benchmarking it.

Performance depends on the merging of several PRs:

In the meantime, I added a piracy.jl file that implements the changes locally. Some type piracy also occurs in the Zygote.jl extension.

Make sure to run this with CliqueTrees.

Benchmarks

Here is the output of gaussian_approximation_comparison.jl.

SUMMARY: FORWARD PASS
================================================================================

-----------------------------------------------------------------------------------------------
Matrix                      n        nnz  Correct    GMRF (ms) Chordal (ms)    Speedup
-----------------------------------------------------------------------------------------------
HB/bcsstk15              3948     117816        ✓      116.996      120.558       0.97×
HB/bcsstk16              4884     290378        ✓      132.224      123.099       1.07×
HB/bcsstk17             10974     428650        ✓      252.566      261.991       0.96×
HB/bcsstk18             11948     149090        ✓      273.233      254.486       1.07×
-----------------------------------------------------------------------------------------------

================================================================================
SUMMARY: GRADIENT (Zygote)
================================================================================

-----------------------------------------------------------------------------------------------
Matrix                      n        nnz  Correct    GMRF (ms) Chordal (ms)    Speedup
-----------------------------------------------------------------------------------------------
HB/bcsstk15              3948     117816        ✓      103.143      108.821       0.95×
HB/bcsstk16              4884     290378        ✓      227.489      169.253       1.34×
HB/bcsstk17             10974     428650        ✓      445.003      239.360       1.86×
HB/bcsstk18             11948     149090        ✓      365.012      174.914       2.09×
-----------------------------------------------------------------------------------------------

Overall:
  Forward - All match: ✓ YES, Avg speedup: 1.02×
  Gradient - All match: ✓ YES, Avg speedup: 1.56×

And here is the output of logpdf_comparison.jl:

================================================================================
SUMMARY: FORWARD PASS
================================================================================

-----------------------------------------------------------------------------------------------
Matrix                      n        nnz  Correct    GMRF (ms) Chordal (ms)    Speedup
-----------------------------------------------------------------------------------------------
HB/bcsstk15              3948     117816        ✓        0.093        0.147       0.63×
HB/bcsstk16              4884     290378        ✓        0.187        0.317       0.59×
HB/bcsstk17             10974     428650        ✓        0.760        0.537       1.41×
HB/bcsstk18             11948     149090        ✓        0.222        0.418       0.53×
-----------------------------------------------------------------------------------------------

================================================================================
SUMMARY: GRADIENT (Zygote)
================================================================================

-----------------------------------------------------------------------------------------------
Matrix                      n        nnz  Correct    GMRF (ms) Chordal (ms)    Speedup
-----------------------------------------------------------------------------------------------
HB/bcsstk15              3948     117816        ✓       83.851       27.275       3.07×
HB/bcsstk16              4884     290378        ✓      113.485       32.366       3.51×
HB/bcsstk17             10974     428650        ✓      164.776       45.750       3.60×
HB/bcsstk18             11948     149090        ✓      137.625       24.023       5.73×
-----------------------------------------------------------------------------------------------

Overall:
  Forward - All match: ✓ YES, Avg speedup: 0.79×
  Gradient - All match: ✓ YES, Avg speedup: 3.98×

samuelsonric and others added 10 commits March 31, 2026 13:47
@samuelsonric
Add type ChordalGMRF.
0979b24
@samuelsonric
Remove dependencies.
9c1aabb
@samuelsonric
Bump CliqueTrees compat.
f39711a
@samuelsonric
Clean Project.toml.
8e3c0ed
@samuelsonric
Fix bug.
44f1cec
@samuelsonric
Larger test matrices.
be26f50
@timweiland
Unify gaussian_approximation for GMRF and ChordalGMRF
8f38d11 
Extract solver abstraction helpers (_ga_init_solver, _ga_update_and_solve!,
_ga_make_posterior) so the Newton iteration and IFT-based rrule each exist
once, dispatching on GMRF type. Removes ~130 lines of duplicated code.
@timweiland
Fix ChordalGMRF constructor and rrule dispatch
3ae505c 
Add 4-arg SparseMatrixCSC constructor that wraps Q in Hermitian, matching
the 2-arg constructor's behavior. Fix rrule dispatch from ::typeof(ChordalGMRF)
(resolves to ::UnionAll, matching GMRF too) to ::Type{ChordalGMRF}.
@timweiland
Add sum(::SymTridiagonal) rrule to fix invalidation-exposed bug
71647cc 
The rrules for HermSparse/SymSparse cause method invalidation that exposes
an upstream ChainRulesCore bug: ProjectTo{SymTridiagonal} extracts only one
triangle of the off-diagonal, losing the factor of 2 from symmetry.
@timweiland
Runic formatting
8ace994
@samuelsonric
Copy link
Copy Markdown
Contributor Author

samuelsonric commented Apr 3, 2026

I am going to rewrite this all using Mooncake; probably today.

@timweiland
Copy link
Copy Markdown
Owner

timweiland commented Apr 3, 2026

Ah, alright! I didn't make any major changes to it, just restructured things a bit to avoid code duplication particularly for gaussian_approximation. Let me know when the Mooncake rewrite is ready to review. :)

@samuelsonric
Copy link
Copy Markdown
Contributor Author

samuelsonric commented Apr 7, 2026

@timweiland It is ready. The MooncakeSparse.jl submodule is a intended to be temporary, until MooncakeSparse.jl is finally registered.

@timweiland
Copy link
Copy Markdown
Owner

timweiland commented Apr 14, 2026

I noticed MooncakeSparse.jl is registered now, nice!
Tried to set things up locally with the registered MooncakeSparse.jl, but hit a diamond dep conflict:

  • MooncakeSparse 0.1 needs Mooncake >= 0.5.25
  • DifferentiationInterface 0.7.16 caps at Mooncake <= 0.5.24

Fresh install fails to resolve. DI maintainers are aware though and this should be resolved quite soon from what I've seen on the Julia Slack, so I'm fine with simply waiting for that to happen

@samuelsonric
Copy link
Copy Markdown
Contributor Author

samuelsonric commented Apr 14, 2026

Annoying. What if you relax the compat in your your local MooncakeSparse.jl to 0.5.24?

@timweiland
Copy link
Copy Markdown
Owner

timweiland commented Apr 14, 2026

As far as I can tell MooncakeSparse.jl uses the new API that was only introduced in 0.5.25, in particular the whole friendly tangent business

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

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@samuelsonric @timweiland