Flatten/Reshape pattern: Ideas for the future

[cako]

Description

Many linear operators act along a specific dimension(s), not on a flattened array. For example a 2D array x shaped (nx, nt) can be differentiated along the x-direction by appying FirstDerivative((nx, nt), axis=0) * x.ravel(). The .ravel() call is necessary to flatten (without copy). Like FirstDerivative, these operators take dims and sometimes axis (or dir before #310) as keyword arguments. Since #331 all operators which receive dims have class attributes dims (model shape) and dimsd (data shape).

This pattern is very common throughout the PyLops and indeed for many operators in the wild. I hope to promote in this discussion the best practices—and possibly new features—for operators which follow this pattern.

Decorators for developers

Flattening by decorator

One suggestion is the use of decorators to simplify the implementation of several operators which follow the reshape/flatten pattern. As an example the Flip operator defines _matvec as follows:

class Flip(LinearOperator):
    def _matvec(self, x):
        x = np.reshape(x, self.dims)
        y = np.flip(x, axis=self.axis)
        y = y.ravel()
        return y

One possibility for simplifying this pattern is defining a decorator along the lines of:

def reshape_flatten(forward=True):
    inp_dims = "dims" if forward else "dimsd"
    def decorator(func):
        def wrapper(self, x):
            x = x.reshape(getattr(self, inp_dims))
            y = func(self, x)
            return y.ravel()
        return wrapper
    return decorator

class Flip(LinearOperator):
    @reshape_flatten(forward=True)
    def _matvec(self, x):
        y = np.flip(x, axis=self.axis)
        return y

By providing this decorator, as long as users implement dims and dimsd, they can implement their functions normally without worrying about reshaping and flattening.

Flexible operators

Arguably doing reshape/flatten does not require a large boilerplate. So the above code might not be super useful. However, they become useful if we allow operators to receive not only flattened arrays, but reshaped ones as well. Currently, running FirstDerivative((nx, nt), axis=0) * x will produce a ValueError unless x is shaped (nx * nt,) or (nx * nt, 1). However, it makes sense that it should work on an array shaped (nx, nt). Implementing this requires some changes to LinearOperator, and the following decorator could be used to automatically flatten/reshape an array depending on the input size.

def autoravel(accept_reshaped=True, forward=True):
    dims_attr = "dims" if forward else "dimsd"
    dimsd_attr = "dimsd" if forward else "dims"

    if accept_reshaped:
        def decorator(func):
            def wrapper(self, x):
                flatten_out = False
                if x.shape != getattr(self, dims_attr):
                    flatten_out = True
                    x = x.ravel()
                y = func(self, x)
                if flatten_out:
                    y = y.ravel()
                return y
            return wrapper
    else:
        funcname = "_matvec" if forward else "_rmatvec"
        def decorator(func):
            def wrapper(self, x):
                K = np.prod(getattr(self, dims_attr))
                L = np.prod(getattr(self, dimsd_attr))
                if x.shape != (K,) and x.shape != (K, 1):
                    raise ValueError("dimension mismatch")
                y = func(self, x.reshape(dims_attr))
                if x.ndim == 1:
                    y = y.reshape(L)
                elif x.ndim == 2:
                    y = y.reshape(L, 1)
                else:
                    raise ValueError(
                        f"invalid shape returned by user-defined {funcname}()"
                    )
                return y
            return wrapper
    return decorator

Closing remarks

The purpose of this discussion is twofold. First, I would like to know if accepting non-flattened array is something that makes sense and would be interesting to users. Second, I would like to hear any thoughts, comments and criticisms on the above approaches.

[mrava87]

Thanks for starting a discussion on this. Here are some thoughts:

• Flattening by decorator: I really like this idea and as you say helps keeping the matvec and rmatvec codes fully focused on the actual implementation of the operator at hand.

• Flexible operators: I am a bit unsure here, and the main reason comes from this:

import numpy as np
import pylops

nx, nt = 100, 50
Fop = pylops.FirstDerivative((nx, nt), axis=0)
Fop * np.ones((nx*nt, 20))

which works just fine because any PyLops operator is also setup to implement _matmat

pylops/pylops/LinearOperator.py

Line 75 in 3529da6

def _matmat(self, X):

and this can be explicitely called here

pylops/pylops/LinearOperator.py

Line 183 in 3529da6

def matmat(self, X):

or implicitly as part of the @ and * operator overloads here:

pylops/pylops/LinearOperator.py

Line 258 in 3529da6

elif x.ndim == 2:

Now, we can of course make changes to LinearOperator and check if the first dimension equals Op.shape[1] then this would mean calling matzot if there is a second dimensions, whilst if this is not the case we will call matvec under the assumption that x.size is Op.shape[1] (or maybe checking first). Is this the kind of change you were envisioning?

[cako]

Currently, we check that the vector is "flat" before we call matvec:

pylops/pylops/LinearOperator.py

Lines 256 to 261 in 3529da6

	if x.ndim == 1: # or x.ndim == 2 and x.shape[1] == 1:
	return self.matvec(x)
	elif x.ndim == 2:
	return self.matmat(x)
	else:
	raise ValueError("expected 1-d or 2-d array or matrix, got %r" % x)

Instead, we could substitute this by:

if np.prod(x.shape) == self.shape[1]:
    return self.matvec(x)
elif x.ndim == 2: # Optionally, we can add "and x.shape[0] == self.shape[1]"
    return self.matmat(x)
else:
    raise ValueError("expected 1-d or 2-d array or matrix, got %r" % x)

Using your example

Fop * np.ones((nx * nt, 20)) # This will call matmat
Fop * np.ones((nx * nt * 20)) # This will error
Fop * np.ones((nx, nt, 20)) # This will error (but maybe we want it to call matmat...?)
Fop * np.ones((nx, nt)) # This will call matvec
Fop * np.ones((nx, nt, 1)) # This will call matvec
Fop * np.ones((nx * nt,)) # This will call matvec

How about I submit a draft PR and we can play around with this?

[mrava87]

Sounds good. I just wanted to make sure you were aware of the presence of matmat and the fact we need to make sure we don't end up having possibly contrasting cases when we allow passing non-flatten vectors to an operator :)

Let me finalise the describe PR, I should be done by the end of today and then you can make a draft PR and we take it from there