Add PDHG solver with support for non-linear operators #322
Conversation
Codecov Report
@@ Coverage Diff @@
## main #322 +/- ##
==========================================
+ Coverage 93.82% 93.88% +0.05%
==========================================
Files 54 55 +1
Lines 3418 3451 +33
==========================================
+ Hits 3207 3240 +33
Misses 211 211
Flags with carried forward coverage won't be shown. Click here to find out more.
Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here. |
| x_mag /= x_mag.max() | ||
| # Create reference image with structured magnitude and random phase | ||
| x_gt = x_mag * snp.exp(-1j * scico.random.randn(x_mag.shape, seed=0)[0]) | ||
| x_gt = jax.device_put(x_gt) # convert to jax type, push to device |
There was a problem hiding this comment.
@Michael-T-McCann: Is this really necessary? Since x_gt is created by an snp function, it's already a DeviceArray.
There was a problem hiding this comment.
Agreed, remove it.
| @@ -273,12 +273,33 @@ def jvp(self, primals, tangents): | |||
|
|
|||
| return jax.jvp(self, primals, tangents) | |||
|
|
|||
| def jhvp(self, *primals): | |||
There was a problem hiding this comment.
@Michael-T-McCann: Thoughts on the name of the method? I'm open to suggestions if you think this one is not so great.
There was a problem hiding this comment.
I don't mind it given that jax has settled on jvp and vjp.
Add PDHG solver with support for non-linear operators.