Optimise rge matrix multiplication

[LucaMantani]

This PR aims at optimising the theory predictions when RGE matrices are included.
The main idea is to exploit the following equation

redefining the linear correction matrix and the quadratic correction matrix by absorbing the RGE matrix Gamma into them. Since this operation needs to be done only once at the beginning, the theory prediction evaluation later on becomes much faster.

NOTE: while the linear correction matrix was already in that form, the quadratic was not produced in that way before. I changed the way it is defined now in order to conform to that expression, which I find also more intuitive and makes the code more streamlined. In particular, the quadratic correction matrix is defined as a triangular matrix. One could have decided instead to make it a symmetric matrix filling both Op1xOp2 and Op2xOp1 with half the number we have in the theory tables but I preferred the option of the triangular matrix.

The fits with RGE are now much faster, e.g. the fits in the following reports are identical:
report_example.pdf

the one in main took 32 minutes, while the one in this new branch took 4 minutes.

Another considerable advantage of this approach is that it's going to make the adaptation of other parts of the code to account for RGE easier. For example, in the analytic fit routine, one will not need to modify too much the code since the linear correction matrix will have the rge embedded automatically (if loaded). The same is true for the Fisher information matrix.

[giacomomagni]

Thank you @LucaMantani, indeed this implementation is more clean!!

I have two small suggestions:

• maybe you can considering adding a check/test that the Quad corrections are always upper triangular after the loader.
• did you check that the fits with constrained coefficients (in particular non linear ones) are still working? I.e that all the cases where QuadraticCorrections was used are now in good shape ? Or you were planning to do it in [WIP] Fisher with RG #93, which now seems superseded by this PR ?

[giacomomagni]

Also the function flatten now can be dropped.

[LucaMantani]

True, I will remove it.

[LucaMantani]

Thank you @LucaMantani, indeed this implementation is more clean!!

I have two small suggestions:

• maybe you can considering adding a check/test that the Quad corrections are always upper triangular after the loader.

• did you check that the fits with constrained coefficients (in particular non linear ones) are still working? I.e that all the cases where QuadraticCorrections was used are now in good shape ? Or you were planning to do it in [WIP] Fisher with RG #93, which now seems superseded by this PR ?

Thanks for these comments.

1. I will add a test checking that.

2. Do you have an example I could test to check this? The Fisher PR is indeed superseded and it will need to be readapted but it should be relatively easy with the new loader.

[giacomomagni]

Mmmm there should be some test around about non linear constrain.
So for instance you can add a constrain like: Op3 = 1.23 * Op1 ^ 2 + 4.56 * Op2 with:

  Op3:
    constrain:
    - Op1: 
      - 1.23
      - 2 
    - Op2: -4.56

to the runcard and check on both branches.

For the FIsher, I think new_QuadraticCorrections has already the good shape (n_c, n_c, n_dat), as I recall the older format was not handy to do manipulations.

[LucaMantani]

if there are tests, they pass, but I will check that explicitly now.

Regarding the Fisher, it has now the quadratic in a good shape but it's not implementing the RGE. One will need to read the matrices when doing reports and then they will indeed be included automatically by the Loader.

[giacomomagni]

Regarding the Fisher, it has now the quadratic in a good shape but it's not implementing the RGE. One will need to read the matrices when doing reports and then they will indeed be included automatically by the Loader.

I'm confused... the report is already using the Loader, so all the RGE effect will be included with tis PR no ?
Or you mean other multiplication have to be done ?

[LucaMantani]

Regarding the Fisher, it has now the quadratic in a good shape but it's not implementing the RGE. One will need to read the matrices when doing reports and then they will indeed be included automatically by the Loader.

I'm confused... the report is already using the Loader, so all the RGE effect will be included with tis PR no ? Or you mean other multiplication have to be done ?

Yes, the report already uses the Loader, so things will be easier to implement, but the Loader needs to receive the rgemat argument and at the moment that happens only in the ultranest routine. Same problem for the analytic fit. I think they are easy fix, or maybe there is even a more unified way of doing it, but I don't think it should be taken care of in this PR

[jacoterh]

Hi @LucaMantani , great we get such a nice speed up, it makes much more sense to do it this way. I added some extra comments in the code just to make sure we understand what we're doing in a year from now. I have one question regarding the triangular property after accounting the RG, see inline.

[jacoterh]

Without RG, this is upper triangular indeed. But this property is lost when the RG is included due to the jnp.einsum contraction above. Are we double counting now? Probably not, but I don't see why atm.

[LucaMantani]

Mmm, good spot...

[LucaMantani]

I tested for a fake UV model with non-linear constraints (card used in attachment) and it works fine:
main:

┏━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━┓
┃ Parameter ┃ Best value ┃ Error ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━┩
│ g         │ 0.000      │ 0.332 │
└───────────┴────────────┴───────┘

branch:

┏━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━┓
┃ Parameter ┃ Best value ┃ Error ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━┩
│ g         │ -0.001     │ 0.331 │
└───────────┴────────────┴───────┘

card.txt

[LucaMantani]

@jacoterh when you have the time, please do an independent check yourself, with a different card just for the sake of it, so that we feel confident enough to merge.

[giacomomagni]

I tested for a fake UV model with non-linear constraints (card used in attachment) and it works fine: main:

card.txt

okay thanks!

[jacoterh]

I also found agreement between main and #96

card.txt

[LucaMantani]

@jacoterh what about the basis_rotation?

[jacoterh]

Not sure why we are worried about the basis rotation actually, this branch gives the same result as main for the EWPO:

[LucaMantani]

Not sure why we are worried about the basis rotation actually, this branch gives the same result as main for the EWPO:

I think the linear part is identical, but the quadratic one is not. The function "flatten" is still used there, while it should be completely dropped out of the code at this point.