Open remilaumont opened 4 months ago
Hi @nabriis, can you please take an initial look at the initial design of incorporating MYULA. This is the PR by @remilaumont and it is the same PR that we eventually want to merge.
Thank you
Thank you @nabriis for your comments. We addressed them and asked again for your review.
Many thanks @nabriis for the review!
@jakobsj, would it be possible that you review the PR, much appreciated!
closes #406
Notes added Aug 12 2024:
Based on the discussion of design we had with the implicit project group, @remilaumont and I (@amal-ghamdi ) did an iteration of updates and we very much appreciate your review @jakobsj
A few points to keep in mind (based on our implicit project meeting a few days ago and based on email notes that we communicated earlier):
We added a
RestorationPrior
class as discussed in the design meeting.restoration_strength
which can be passed at initialization. It is also one of the parameters we made required for the function handle for the restorator (what we called: fun, in the email discussion). fun takes,x
,restoration_strength
and any other parameters the restorator might need.We added a
MoraeuYoshidaPrior
class.The potential that corresponds to this smoothing is "R_alpha" (not R), thus note that we added implementation for
logpdf
inMoraeuYoshidaPrior
to compute R_alpha (only when R is available because it is based on R)At some point in the future, we will also need something similar to MoraeuYoshidaPrior for PnP, as PnP has a more practical appeal in terms of the results it can achieve. One possibility is that we can add a
PnPPrior
class (not in the scope of this PR) to inherit fromMoraeuYoshidaPrior
. It can overwrite thelogpdf
to return Nan but uses the same gradient formula (as both approaches lead to the same gradient formula, although they are coming from two different constructions/derivations)Now in principle,
MoraeuYoshidaPrior
can be tested with MALA/NUTS (suggested by @jeverink at some point), which is really nice, we did not do such tests yet (not necessarily in the scope of this PR but it is nice to try sometime soon)Note: a smoothing like the Moraeu Yoshida can in principle be applied to the likelihood (for non-differentiable likelihoods). Is considering the
MoraeuYoshidaPrior
a prior only, limiting? Should we think of a more generalization? Note that now theMoraeuYoshidaPrior
is part of the implicit prior submodule. We might consider relocating it somewhere if we decide to generalize this smoothing to likelihoods.Compatibility checks for creating a
MoraeuYoshidaPrior
based on aRestorationPrior
We needed to check that the
smoothing_strength
ofMoraeuYoshidaPrior
is the same as therestoration_strength
of theRestorationPrior
, and we implemented that check. One thought here is that we can get rid of storingrestoration_strength
in theRestorationPrior
then we will not need this compatibility check. One side effect is that therestore
method will need to acceptrestoration_strength
in addition to the other restoration parameters.Geometry handling: now we made the
MoraeuYoshidaPrior
take on the geometry of the passedRestorationPrior
. The design of enabling this feature has similarities with what is in theRegularizedGaussian
class.MYULA
MYULA
to work withRestorationPrior
type. We can extend the type of priors it deals with (e.g. any class that has prox) in the futureMYULA
also acceptssmoothing_strength
parameter at initialization and ensures that it is equal to therestoration_strength
of the prior (note that we can get rid of this equality check if we did the suggestion above: not to store therestoration_strength
in theRestorationPrior
).MYULA
now keeps a smoothed target as class attribute (_smoothed_target
) for which the gradient is computed when sampling, besides the original target.