Does elastic support batched velocity models?

[CharlelieLrt]

Hi,

To my understanding, elastic does not support generation over a batch of multiple input subsurface models (the subsurface properties are documented as 2D arrays).
Is there a way to perform generation over batched-subsurface models?
Would using torch.func.vmap work with elastic?

[ar4]

Hi @CharlelieLrt,

I'm glad to hear that this is something you are interested in, as I have added this feature to the next release. There are a few other features that I had planned to add before making the release, so it will probably be a few weeks before releasing it. If you would like, I can provide you with a pre-release wheel that includes this feature.

[CharlelieLrt]

@ar4 thank you for the quick reply, and yes that would be awesome if you could provide a pre-release version!

Just for my information, what kind of batching mechanism are you planning to support:

1. Native batching for elastic, i.e. accept velocity models of shape (B, depth, width)

OR

2. Non-native batching, but compatibility with torch.func.vmap

... or both maybe? :)

[ar4]

If the action runs successfully, the wheels should be available as artifacts from this action in about an hour and a half, or you can download the sdist artifact, which is already available, and build it yourself now.

It is the first type. Each property model can now be provided with shape (ny, nx) or (n_shots, ny, nx). Does the second type have advantages that would be useful to you?

[ar4]

Ah, I see that some of the new input validation I added to the code is causing test failures. I'll fix that tomorrow. If you would like to try the new feature before then, I have attached the sdist. You should, I hope, be able to install it using pip install deepwave-0.0.21.tar.gz. You will need compilers available and it will take several minutes as it will compile the C and CUDA code.

deepwave-0.0.21.tar.gz

[CharlelieLrt]

Each property model can now be provided with shape (ny, nx) or (n_shots, ny, nx).

By (n_shots, ny, nx), you mean N independent subsurface models, correct? AFAIK the elastic and other operators already support multiple shots (sources).

Does the second type have advantages that would be useful to you?

Yes, it does. For example if we want to compute per-example gradients, rather than a single gradient summed over multiple subsurface models.

[ar4]

I decided to delay the other features until the next release, so I have now made a new release (v0.0.21) that includes this feature.

By (n_shots, ny, nx), you mean N independent subsurface models, correct?

Correct.

AFAIK the elastic and other operators already support multiple shots (sources).

Correct. Previously you could propagate multiple shots simultaneously, but they would all share the same property models. You now have the option of using different property models for each.

For example if we want to compute per-example gradients, rather than a single gradient summed over multiple subsurface models.

I think you can do what you want using the "native batching approach". Would it enable you to do something else useful if I also implemented the "non-native batching" approach?

I have added an example to the documentation demonstrating the new batched feature.

I hope that you find it useful. Please let me know if you encounter any problems.

[CharlelieLrt]

Awesome, thank you so much for releasing this feature so quickly!

I think you can do what you want using the "native batching approach"

Yes, the per-example gradients can be done with native batching in some standard cases. It's probably sufficient for most of my use-cases for now. There are however some more advanced cases where combining the elastic operator with vmap could be beneficial (e.g. using it with forward mode AD from torch.func).

[ar4]

e.g. using it with forward mode AD from torch.func

These are the kinds of interesting applications that Deepwave is intended to facilitate, so adding support for vmap sounds like a good idea. There are a few other high priority features that I am keen to add, but I will look into it once those are done. Thank you for the suggestion.