Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Weighted sums reduction #354

Open
DavidLapous opened this issue Jan 30, 2024 · 2 comments
Open

Weighted sums reduction #354

DavidLapous opened this issue Jan 30, 2024 · 2 comments
Assignees

Comments

@DavidLapous
Copy link
Contributor

DavidLapous commented Jan 30, 2024

Hi, and many thanks for your work.
I have the following (reduced) problem, which I didn't manage to encode in pykeops :

  • I have a point measure $\mu$ (i.e. sum of weighted diracs), for each "data", lets say
    • pts with shape (num_data, num_pts, D)
    • weighted by an array weights of shape (num_data,num_pts)

On which I want to evaluate the convolutions against (multiple, but lets say 1 for now) gaussian measure,

  • with covariance matrix K of shape (D,D)
  • of some points x of shape (num_data,num_x,D)

Formally, I want something like
$$\mathcal N(0,K) \star \mu_{\text{data}}(x_i) \quad \text{for } x_i \in x, \text{data}$$
where
$$\mu = \sum_{x\in \mathrm{pts}} \delta_x \mathrm{weights}[x]$$
So my code is as follows:

  • Lazy encoding pts and x with shape (num_data, dummy_num_pts, dummy_num_x, D)
  • Then
z = pts-x
exponent = -(z.weightedsqnorm(K.inverse().flatten()))/2
exponent = exponent.exp()

Now I want to do a weighted sum w.r.t. the dirac's weights weights, but I can't find a way to do the weighted sum

(exponent*weights[:,:,:,None]).sum(dim = pts_axis)

with pykeops, as the last dimension is a 1 (because of the sqnorm), not a D.

My current workaround is to do a (python) for loop over num_data (and the other axis that I removed for this example), but this starts to be slow.
Is there a way to do this with pykeops ? I didn't find anything in the api documentation.

@jeanfeydy jeanfeydy self-assigned this Feb 1, 2024
@jeanfeydy
Copy link
Contributor

Hi @DavidLapous ,

Thanks for your interest in the library!
As far as I can tell, you are trying to compute a batched Gaussian convolution, which is very well supported.

Could you maybe write a full "minimal working example" (with realistic values for num_data, num_pts, num_x and D, since their orders of magnitude matter to select an optimal implementation) of what you are trying to compute?

Then, I will be able to give you specific performance tips.

Best regards,
Jean

@DavidLapous
Copy link
Contributor Author

DavidLapous commented Feb 9, 2024

Hi @jeanfeydy, thanks for your response :)

I think I found a workaround, it seems to work but I'm not super sure as it relies on a "non-valid" operation.

Basically, a simplified version of what I'm doing is something like this:

import numpy as np
num_data = 100 # range : 100 - 5 000
mean_num_pts = 200 # range : 100 - 5 000
D = 2 # less than 5 in general
num_x = 500 # range : num_pts - 1e9
dtype=np.float32

# each data doesn't have the same number of diracs in the measure.
# we unragg the tensor, and the multiplication by the weight in the end will 
# kill the dummy variables. It's fine as the variance is "small" in practice.
num_pts_variance = 100
num_pts_per_data = np.random.normal(mean_num_pts,num_pts_variance, size=(num_data,)).astype(int) 
max_num_pts = num_pts_per_data.max()
pts = np.random.uniform(size=(num_data,max_num_pts, D)).astype(np.float32) ## unragged pts
weights = np.random.choice([-2,-1,1,2], size=(num_data,max_num_pts), replace=True).astype(np.float32)
for i,data_num_pts in enumerate(num_pts_per_data):
    weights[i,data_num_pts:]=0 # kills the dummy variables
# The signed point measure is \sum \delta_x *w for x,w in zip(pts,weights)

# the points on which to evaluate the convolution
x = np.random.uniform(size = (num_data, num_x,D)).astype(np.float32)

from pykeops.numpy import LazyTensor
## (num_data, max_num_pts, num_x, D)
lazy_diracs = LazyTensor(pts[:,:,None,:])
lazy_x      = LazyTensor(x[:,None,:,:])

K = np.random.uniform(size=(D,D))
K = K.T@K + np.diag([.01]*D) # some covariance matrix inverse
K = K.astype(np.float32)

## batch convolutions
z = lazy_diracs - lazy_x
exponent = -(z.weightedsqnorm(K.flatten()))/2
exponent = exponent.exp()
assert exponent.shape == (num_data,max_num_pts, num_x)
assert weights.shape == (num_data,max_num_pts,)

convolutions = (exponent*weights[:,:,None,None]).sum(1)
print(convolutions.shape) # (num_pts, num_x, 1)

The only part that changed from my last part of the code is on the last line

convolutions = (exponent*weights[:,:,None,None]).sum(1)

where I added a None to the weights. (Pure luck on this haha)
This should not work as , in this example

  • exponent is of shape (100, 393, 500) and
  • weights[:,:,None,None] is of shape (100, 393, 1, 1)

are not broadcastable one onto the other.
I'm not really sure what's happening under the hood, but I the "lazy constraint" which makes weights[:,:,None] non-Lazy-ifiable, and the fact that adding a None makes Lazy and the computation valid makes this very fortunate.

I only tested the computation on small examples, but this works. Can you confirm that I'm not doing anything wrong ? I'm also curious on the what's the wizardry that makes this work.

Also, in my setup, there is another axis (usually small), for both the covariance matrices, and the data, but I think this is fine. The "for loop" bottleneck is the "num_data" I think. I'm also using the torch backend, which, fortunately, has the same behaviour.

Kind regards,
-- David

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

2 participants