jaxified monotonicity operator

[pmelchior]

scarlet1 uses proximal operators for constraints, scarlet2 uses transformations. So, things like our good-old positivity constraint goes from np.maximum(0, x) to jnp.exp(x). We could do the former still, but the latter has much better gradients.

Now, is there a way to convert our very important monotonicity operator into a jitable transformation? I have doubts because there are a lot of of conditional replacements in the C layer of the operator. But maybe @fred3m has ideas?

[fred3m]

Sorry, I've been meaning to reply to this for a while but it slipped deep into my pile of TODO's.

When you say jitable, do you mean essentially taking it out of C++ and creating a loop that runs in python (but with JIT will be much faster)?

If so, take a look at the implementation in scarlet lite:

• C++: https://github.com/lsst/scarlet_lite/blob/main/python/lsst/scarlet/lite/operators_pybind11.cc#L14-L34
• python: https://github.com/lsst/scarlet_lite/blob/main/python/lsst/scarlet/lite/operators.py#L42-L252

It should be much easier for you to follow what's going on here, but ping me if you have any questions.