PyLops / pylops

PyLops – A Linear-Operator Library for Python
https://pylops.readthedocs.io
GNU Lesser General Public License v3.0
422 stars 101 forks source link

Complex LinearOperator #202

Closed Sketos closed 2 years ago

Sketos commented 3 years ago

Hi,

I have a LinearOperator object that is complex and a data vector that is also complex. I want to feed these into the NormalEquationsInversion but I want my output solution to be real.

I was wondering if there is a way to split the "Op" LinearOperator to one that has the real and one on the imag components so that when I want to perform an Inversion I can get the desired outcome by means that I will explain with an example. Assume I have the following

Aop = pylops.MatrixMult(
    sparse.bsr_matrix(matrix), dtype="complex128"
)

FFTop = pylops.signalprocessing.FFT2D(dims)

Op = Aop * FFTop

I also have complex data and complex weights that I can feed to a modified version of the "NormalEquationsInversion",

x = NormalEquationsInversion(Op, data_real, data_imag, Weight_real, Weight_imag)

where

NormalEquationsInversion(Op, data_real, data_imag, Weight_real, Weight_imag, Regs):

    #NOTE: Is there a way to get the real and imag components of this operator as separate operators?
    Op_real = ...
    Op_imag = ...

    OpH_real = Op_real.H
    OpH_imag = Op_imag.H

    Op_normal = (OpH_real * Weight_real * Op_real  + OpH_imag * Weight_imag * Op_imag) + RegH * Reg

    y_normal = OpH_real * Weight_real * data_real + OpH_imag * Weight_imag * data_imag

    xinv, istop = cg(Op_normal, y_normal, **kwargs_cg)

    return xinv

Is there a way to fill in the blanks above to make it do what I want?

cako commented 3 years ago

I don't see a problem with creating a real and imaginary operators. You just need to wrap the forward and adjoint functions with np.real and np.imag, and make sure they pass the dot test. The issue I see with this approach is that you will likely be doubling computation, unless you memoize the operation. Could you be a little bit more specific about your application?

To be honest I would rethink your problem a little bit. Conventionally, linear operators are defined between vector spaces over the same field (e.g. only over the reals or only over the complex numbers). If you have complex data and complex operators, it makes sense to use the complex numbers as the field for your problem. In this case, the fact that you want a real solution can be enforced in different ways. You can for example project the solution at each iteration onto the reals. Similarly, if it makes sense for the problem, you can project the adjoint operation onto the reals. Basically take np.real of whatever comes out of Op.H. This sometimes makes sense for FFTs when you have real data (see this issue for an example). If none of these solutions make sense for your problem, you could regularize your problem so as to damp the imaginary part of the solution. Alternatively, you could use scipy.optimize.minimize and write your constraint as scipy.optimize.LinearConstraint. This option would mean abandoning NormalEquationsInversion, however.

mrava87 commented 3 years ago

Hello, I echo @cako here. I think what he suggests makes sense. Also I think that the issue https://github.com/PyLops/pylops/issues/162 he provides it is worth a read as the problem is very similar (and actually looking at your profile I guess you work with @Jammy2211?)

Nevertheless, this problem/discussion happened a few too many times to ignore and I took some time to make a very simple example with the following aim:

  1. show the equivalence of inverting the complex problem or the augmented (real/imag) problem
  2. investigate if we could add some methods to pylops to allow the real/imag problem without losing efficiency (re to the memoize link of Carlos.)

I think my answer to both questions is yes, although I must admit I am still puzzled why the real/imag problem would add anything to the table compared to solving the full complex problem directly (especially with linear operators where we can easily add reality enforcement like discussed by Carlos and in the other issue). But being puzzled doesn't mean I know by fact why you want to do it, so perhaps you have good reasons (will be great if you can elaborate :) )

Here is the link to the notebook https://github.com/PyLops/pylops_notebooks/blob/master/developement/ComplexMatmul.ipynb

@cako would be great to also get your opinion on a few drafts for how I see making this available through the real and imag methods in LinearOperator and some sort of Memoize operator that wraps any operator - so far fully coupled with MatMul in MatMult_memoize (it was an half an hour attempt so probably we can chew on it for a bit longer... but if you are happy we can try to work on it together until we are both satisfied ;) )

cako commented 3 years ago

I had a look at the notebook and frankly there's not much to add. I just thought I'd take the memoization a bit further by storing a limited history of inputs and crawling through it. It's a pretty simple but general implementation for any LinearOperator. I added it to the end of this notebook:

https://github.com/PyLops/pylops_notebooks/blob/master/developement/ComplexMatmul_v2.ipynb

Let me know what you guys think!

mrava87 commented 3 years ago

Hi @Sketos, I have now added what @cako and I sketched in the notebook to PyLops.

You can see now that there is a new operator called MemoizeOperator and each LinearOperator can use the methods toreal and toimag which do take the real/imag parts of the output of the operator and adjoint. This way, we can now solve complex-valued problems in two different ways: the one you used in the matrix formation with stacked system of eqs for real part and imag part or just one complex system of equations. See this tutorial: https://pylops.readthedocs.io/en/latest/tutorials/realcomplex.html#sphx-glr-tutorials-realcomplex-py

We can also more easily force the output of the forward/adjoint of an operator to be its real or imag part without having to hard-code this behaviour in each operator directly. So if you have a complex-valued operator that you want to force to have real numbers in the model you can do:

A = np.ones((n,n)) + 1j*np.ones((n,n)
Aop = MatrixMult(A)
Aop = Aop.toreal(forw=False, adj=True)

Test it out and let us know how it goes :D

Sketos commented 3 years ago

Hello @cako and @mrava87,

First of all, thank you both very much not only for helping me understand better how we can make our matrix approach and the one we are trying to develop with Pylops agree, but also for the overall contribution you had to our project.

I think I understand a bit batter why the solution we get with Pylops does not agree with the one we get with our matrix approach, but I would be interested to hear your thoughts on that. The way the problem is set up in our matrix implementation at the moment does not satisfy the fundamental assumption of the pylops implementation that a single "x" vector will be the solution to either the real or the imag parts of our matrix "A" and data-vector "y". This could be one of the reason the two approaches disagree.

I tried to give a short description with an example here (https://github.com/Sketos/github_issues/blob/master/pylops/%23202/main.ipynb) in case that will be useful to you. This is the reason why I am exploring the alternative approach of creating separate "real" and "imag" operator so that I can add them together before trying to find the solution with the linalg.cg method and make the two approaches equivalent.

Another worry I have, which I am currently trying to understand, is whether the forward and adjoint methods of our NUFFT are set up properly to deal with what I am trying to do (I am reading into the cg method details). If we think the problem in terms of our matrix approach, for the linear system Ax = y,

A = A_real + A_imag = f_real^T f_real + f_imag^T f_imag == f_real^H f_real + f_imag^H f_imag != f^H f

where f = f_real + f_imag * 1j (this is because f^T != f^H for a complex matrix). Therefore, as is currently set up, the adjoint is not doing what I think we want and unless we split the "Op" operator into a real and imag so that the transpose is equivalent to the hermitian then we won't be getting the same answer. I still trying to wrap my head around this so my explanation might not be as clear as I would like it to be (maybe the example at the end of the above link gives a better explanation).

I am now going to have a more thorough look at the new operator you developed (i.e. MemoizeOperator) to see if this is the way forward. If you have any more feedback to what I have described I would be interested to hear your thoughts.

Cheers

cako commented 3 years ago

Thanks for the notebook, it's helped me understand a bit more where you're coming from.

If I understood correctly, you a problem which you have chosen to pose it the following way:

(F_real + F_imag) x = D_real + D_imag

You have traditionally solve this problem using a matrix formulation, I'm supposing something like x = (F_real + F_imag) \ (D_real + D_imag). Now you're trying to figure out whats the best way to frame it using PyLops and linear operators. If this is the case, then x as per your matrix solution is a purely real vector.

Now, you want to implement this in PyLops, but you only have access to the complex operator F (not the matrix) defined by F = F_real + i F_image. One suggestion was to memoize the operator so that you had access to F_real and F_imag as operators and not matrices without any additional computation. This approach gives the exact result as x in your example:

# --------------------------------------
# NOTE: Method 3 (Memoize)
# --------------------------------------
# Build complex operator, in practice you only have access to F_clpx
F_cplx = pylops.MatrixMult(F_real, dtype='float64') + 1j*pylops.MatrixMult(F_imag, dtype='float64')
# Memoize the operator, giving access to F_real and F_image as operators
F_memo = pylops.MemoizeOperator(F_cplx, max_neval=1)
F_real_op = F_memo.toreal()
F_imag_op = F_memo.toimag()
# Ugly hack to get the corresponding real type. This should probably be default
#    complex64 -> float32, complex128 -> float64, complex256 -> float128
dtype_real = np.array(0, dtype=F_memo.dtype).real.dtype
F_real_op.dtype = dtype_real
F_imag_op.dtype = dtype_real

x_lops_method_3 = NormalEquationsInversion_method_2(
    Op=F_real_op+F_imag_op,
    data=D_real+D_imag
)

Side-note: @mrava87 the toreal and toimag methods do not change the dtype of the base operator, which it should. I'll submit a commit to fix this.

I think if (F_real + F_imag) x = D_real + D_imag is the problem you want to solve, that's probably the simplest way. But before you go ahead and do that, consider what this system is actually solving: (f_real^T f_real + f_imag^T f_imag) x = f_real^T d_real + f_imag^T d_imag

This system does not correspond to the normal equations of any system I can think of. The (real) system (f_real + f_imag) x = d_real + d_imag would give rise to (f_real^T + f_imag^T) (f_real + f_imag) x = (f_real^T + f_imag^T)(d_real + d_imag) which if you expand you will have the cross-terms which your system lacks. On the other hand (f_real + i f_imag) (x_r + i x_imag) = d_real + i d_imag would give rise to (f_real - i f_imag)^T (f_real + i f_imag) (x_r + i x_imag) = (f_real - i f_imag)^T(d_real + i d_imag) which would also have a bunch of cross terms which you are negleting in the equation in your post.

Again I'm not sure what's the best way to frame your problem is, but (f_real^T f_real + f_imag^T f_imag) x = f_real^T d_real + f_imag^T d_imag seems like a very strange system to me since it mixes up real and imaginary components. Personally I would go back and check if this is the system you are actually interested in, but I may be completely off base as well!

Anyways, let me know if this makes sense and hopefully we can sort it out soon!

mrava87 commented 3 years ago

Side-note: @mrava87 the toreal and toimag methods do not change the dtype of the base operator, which it should. I'll submit a commit to fix this.

Good point, go for it :)

Sketos commented 3 years ago

Hello @cako and @mrava87,

Thank you for clarifying the use of the MemoizeOperator.

I have now added a new example (Example 2) at this link. With this example I would like to bring together the use of the separate real and imag Operators with the discussion in the #162.

As @Jammy2211 described in that ticket we would like our Pylops Formalism to never have to explicitly create the $f_tilde$ matrix, which for complex data is the NUFFT of the $f$ matrix ($f_tilde = Df$).

In Example 2 you will find our current pylops Formalism (method 1) of the problem where the solution we get is complex (we dont want that). There is also a modified version of the pylops Formalism (method 2) where I create separate real and imag operators (for which the MemoizeOperator could come in handy).

The problem I currently have is that method 2 does not work either (i.e. does not give the same solution as the Matrix Formalism) and I dont fully understand why. I believe there is something wrong with the way the NUFFT is implemented because of the following (which would be the first step in the solver method)

# NOTE: x != x_inv_real where,
 b_real = Op_real * x
 x_inv_real = Op_real.H * b_real

If we understand how to make the Matrix Formalism and the Pylops Formalism agree with this example we will then be able to implement pylops in to our project.

Please let me know if something is not clear.

Cheers

mrava87 commented 3 years ago

Thanks a lot, I will take a look into your code and try to understand what the problem may be.

One thing I am still confused though is what f actually is? The reason I ask is because a linear operator formalism is based on the concept that you have a vector x (your model), a vector y (your data) and a chain of linear operations that you can represent on a paper as many matrices which multiply each other and the resulting matrix is then inverted, or alternatively you can just apply the first operator to x, the second to what comes out, and so on and so forth.

Now in your case I would expect that you have x, can make f (say by a code that mimics what it does or if it is sparse as a sparse matrix) and apply it to x, and then apply D to the resulting vector. Since D is NUFFT this is easy and you do it already I guess. But you say f_tilde matrix, which for complex data is the NUFFT of the f matrix? Could you not all along avoid this and apply f to x followed by the NUFFT of the resulting vector? Note that I always talk about vector, but eventually they can be reshaped and physically be a 2/Nd array within the linearoperator code? I kind of remember this is what me and @Jammy2211 discussed some time ago?

Regarding the MemoizeOperator, this is a smart way of dealing with complex operators you want to evaluate both their real and imag outputs, but to start you don't need it. All you need is to use the new .real() and .imag() methods that are now available for any linear operator, so you can do F=NUFFTOp(...); Fr = F.real(); Fi = F.imag() and they will extract the real and imaginary part. In https://github.com/PyLops/pylops_notebooks/blob/master/developement/ComplexMatmul.ipynb we show how solving the complex or real+imag problem should be the same :)

But as I said, I will go through your notebook step by step and get back to you with more on this :)

mrava87 commented 3 years ago

And if you look at the equations I wrote there, this is how i explain myself the two problems are the same. You start from y=Ax

You then say that this is y_r+jy_i=(A_r+jA_i)x and that since x is real the problem is decoupled... when x multiplies A_r it must give y_r, and when x multiplies A_i it must give y_i. Now this is the same as writing two set of equations:

y_r=A_r x 
y_i=A_i x

which if you write the normal equation of this system gives what you solve with the matrix formalism.

A_r^Hy_r + A_i^Hy_i=(A_r^HA_r + A_i^HA_i) x

What I would suggest is to work on making an operator A that takes a real vector x and produces a complex vector y which real and imaginary parts are the same as those produced in your matrix formalism by f^tilde_R and f^tilde_I. If you can do that and then apply the .real() and .imag() methods to this operator you will have the same effect as np.dot of f^tilde_R and f^tilde_I on a vector. Do then the same for their adjoint. If this is the case the inverse problem will have to give the same result as effectively every step is mimicked 1-to-1. I just worry that going straight for the inverse problems makes your life harder as its much harder to tell where the two pairs of forward/adjoint do not match each other

Sketos commented 3 years ago

Hello @mrava87,

I am just posting a quick answer to your question, but I ll think about your suggestion and get back to you.

Your understanding of the f is correct. It essentially maps each of the n_s pixels in the source-plane to pixels in the image-plane, n_p (f itself is a sparse matrix of 0's and 1's of shape (n_p, n_s)).

Let's assume we know the true model vector x which has a shape of (n_s, 1) then using the mapping matrix, f, we can construct our model image-plane vector, y_model, as

y_model = f @ x

For complex data we can then apply the NUFFT operator, D, which has a shape (n_v, n_p) to y_model and the resulting complex model vector which has a shape (n_v, 1) is compared to the data vector, y, (e.g. if we want to evaluate the likelihood of that model).

However, for strong lens modelling as @Jammy2211 has already described (in way more details that I will try to do here; see posts from July) we have a model for the galaxy acting as the lens and that model essentially creates the f matrix and is that model that we want to evaluate the likelihood of. Given a model for the lens, which gives us f, we want to find "x" above.

Cheers

Sketos commented 3 years ago

Hello @mrava87 and @cako,

It seems by adding the line of code

Op = Op.toreal(
    forw=False, adj=True
 )

before passing Op to the "NormalEquationsInversion" makes the Matrix and Pylops formalisms equivalent, hurray!!! I also implemented a weight operator, WOp, and we still get consistent results between the two formalisms. I still have a couple of questions but I will bring then up later as for now things seems to be working.

Now the last piece of the puzzle is to implement the regularization properly. In the matrix formalism with regularization, the system we are solving is given by the following formula,

Screenshot 2021-01-02 at 17 05 48

For the Pylops formalism we compute the H = R^T R matrix which we wrap as a LinearOperator and we then pass it to the "NormalEquationsInversion" (this has already been discussed in #162 where we opted for the option to pass the already computed H matrix to the "NormalEquationsInversion" instead of the R martix; this way only the '_matvec' of the 'H_Op' class is being called in the solver).

You will find the updated example here with the option to turn regularization on or off (when turned on the Matrix and Pylops formalisms are equivalent). Is there a way to implement regularization in the Pylops Formalism so that the normal equations are as the equation above (I am not sure exactly how the normal equations look like the way regularization is currently implemented)?

Cheers

cako commented 3 years ago

Good to hear you got the equivalent results. I'll admit I'm not entirely sure why those options were needed but well done on finding them 😄

Now that you're moving on the to regularization, I'd highly recommend trying out some of the sparse solvers (good tutorial here) coupled with sparsifying transforms (e.g. wavelet, seislet or curvelet)