Gravity simulation with Choclo as engine #1285

santisoler · 2023-09-19T20:04:52Z

Summary

Add the option to use Choclo as the engine for running the gravity simulations, considerably reducing the computation time of the forward simulations and the creation of the sensitivity matrix. Add Choclo as one of the optional dependencies. Update and extend tests for the gravity simulation using pytest. By default, the simulation will use geoana as the engine. Users can alternatively choose choclo instead if it's installed.

PR Checklist

If this is a work in progress PR, set as a Draft PR
Linted my code according to the style guides.
Added tests to verify changes to the code.
Added necessary documentation to any new functions/classes following the
expect style.
Marked as ready for review (if this is was a draft PR), and converted
to a Pull Request
Tagged @simpeg/simpeg-developers when ready for review.

Reference issue

What does this implement/fix?

Additional information

Don't force values to be np.float32, Numba will cast them automatically when trying to write to the preallocated matrix.

When defining the decorated functions as attributes of the class, then every new simulation is created, they need to be recompiled.

Move the test for gravity inversions to pure pytest. Use parametrization to run the same test using geoana and choclo. Add a `random_seed` to `make_synthetic_data` to avoid tests failing randomly.

Add tests for sanity checks of the engine and some other attributes. Add tests for running the forward in parallel and serial. Add test for conversion factor.

codecov · 2023-09-19T20:13:30Z

Codecov Report

Attention: 35 lines in your changes are missing coverage. Please review.

Comparison is base (211f627) 82.36% compared to head (edab185) 82.35%.

Files	Patch %	Lines
...imPEG/potential_fields/gravity/_numba_functions.py	32.55%	29 Missing ⚠️
SimPEG/potential_fields/gravity/simulation.py	94.05%	6 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1285      +/-   ##
==========================================
- Coverage   82.36%   82.35%   -0.02%     
==========================================
  Files         164      165       +1     
  Lines       25099    25239     +140     
==========================================
+ Hits        20674    20785     +111     
- Misses       4425     4454      +29

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Remove the `random_seed` argument added to the `make_synthetic_data` method. The change was moved to a separate PR to check the reason why some tests were failing in CI.

This reverts commit dfec5d8.

thibaut-kobold · 2023-09-28T17:32:20Z

@santisoler : does this mean SimPEG PF will support GPU with this change?

santisoler · 2023-09-28T20:23:40Z

@santisoler : does this mean SimPEG PF will support GPU with this change?

Not with the changes in this PR. I'm using Numba, but running the simulation in CPU.

SimPEG/potential_fields/gravity/simulation.py

jedman

fantastic work @santisoler , this addition dramatically improves the performance of the gravity simulation. looking forward to propagating it to the magnetic simulation too!

santisoler · 2023-10-24T22:25:42Z

Hi @domfournier! Thanks for those suggestions. I've applied some of them, specially the ones related to reducing code duplication. I've also replied the other ones explaining better my design decisions.

I think we can move forward with this PR if you are happy with it, and continue the discussion about design on a future PR. Specially the one I'm going to open implementing the mag simulation.

SimPEG/potential_fields/gravity/_numba_functions.py

domfournier

Thanks for the changes. Looking good on this side.

domfournier · 2023-10-25T12:09:19Z

Hi @domfournier! Thanks for those suggestions. I've applied some of them, specially the ones related to reducing code duplication. I've also replied the other ones explaining better my design decisions.

I think we can move forward with this PR if you are happy with it, and continue the discussion about design on a future PR. Specially the one I'm going to open implementing the mag simulation.

Sounds good yea, we can generalize on the second round for mag. Looking forward to testing the speedup!

jcapriot

I've been playing around a bit with numba myself the last couple days. I like that it's able to jit a lot of the functions in here so we start avoiding the costly operations of initializing new (rather large) numpy arrays and also avoid the stacking operation. Avoiding the costly multiprocessing step here is great too. I've been looking into updating geoana as well so those functions could be used inside numba's functions. At that point the speed is roughly the same between the two (geoana vs choclo).

Eventually we should put in the simplifications when you have gxx, gyy, and gzz, but as those are not currently in the code (for accuracy reasons) no reason to add them in now.

SimPEG/potential_fields/gravity/simulation.py

jcapriot · 2023-10-25T18:30:16Z

SimPEG/potential_fields/gravity/simulation.py

+            raise TypeError(f"Invalid mesh of type {self.mesh.__class__.__name__}.")
+        return cell_nodes
+
+    def _get_tensormesh_cell_nodes(self):


Let's get a new version of discretize out and tagged so we don't need to include this bit.

santisoler · 2023-10-25T22:56:44Z

Thanks @jcapriot for the review! Yes, numba has neat things indeed!

Eventually we should put in the simplifications when you have gxx, gyy, and gzz, but as those are not currently in the code (for accuracy reasons) no reason to add them in now.

Yes. We might want to add some optimizations to this simulation when having multiple components of acceleration or tensor. These are also relevant for the mag simulation that's coming shortly.

Use the new `discretize.TensorMesh.cell_nodes` instead of the private method that was in the simulation. Increase the minimum version of discretize in `setup.py` and `environment_test.yml`.

santisoler · 2023-11-03T20:32:49Z

@jcapriot I've updated the simulation to use the new TensorMesh.cell_nodes method in discretize v0.10.0. I've also increased the minimum version of discretize in setup.py and environment_test.yml.

jcapriot

Throughout the documentation, instead of writing array write either numpy.ndarray (or array_like where appropriate). Also, it could be worth it to run a comparison with the most recent geoana with the jittable prism kernels now available there.

santisoler · 2023-11-20T17:01:11Z

@jcapriot do you think we should change anything else on this PR, or is it ready to be merged?

jcapriot · 2023-11-20T17:33:49Z

@jcapriot do you think we should change anything else on this PR, or is it ready to be merged?

The only thing I would change is the actual argument named after choclo, (choclo_parallel) I would rename it something related to numba (numba_parallel?) as geoana prism kernels should now support being jitted. We should add another PR to update the internals to also support the geoana ones as well if it's not too different.

I feel like the default behavior (in the future) should be: If numba is installed, use numba parallelized geoana (since geoana is already a requirement), else use the vectorized calls, We can add choclo as a further optional thing if that behavior is actually needed.

I also think we should probably advertise it as being numba accelerated as well in that case.

lheagy · 2023-11-20T23:52:22Z

I am really excited to get this in, congrats @santisoler! This accomplishes a major goal of having a concrete connection between the Fatiando and SimPEG communities and the benchmarks you have shown are impressive!

We now have a couple of implementations, and that can be useful for a number of reasons, including testing. I think the discussion about the default behaviour is an important one. choclo is a lightweight dependency and from my perspective, the explicit connection between the SimPEG and Fatiando communities is something that should be valued as we consider this. We can chat more at some of the upcoming meetings.

Again, nice work Santi!

jcapriot · 2023-11-21T00:32:57Z

Oh yeah, I fully support keeping choclo as an option, I just don’t want people to get a bad impression of geoana either… (I.e geoana=slow, choclo=fast). It’s more like they are both good and fast!

santisoler · 2023-11-22T18:33:54Z

Thanks @jcapriot and @lheagy!

Considering the work you've been doing in allowing geoana kernels to be called within Numba jitted functions, it makes sense to rename the argument to numba_parallel. I've just pushed that change.

Oh yeah, I fully support keeping choclo as an option, I just don’t want people to get a bad impression of geoana either… (I.e geoana=slow, choclo=fast). It’s more like they are both good and fast!

I'm sorry for making that impression. The main difference in performance doesn't come from geoana, but from the Python implementation of the simulation in SimPEG. The choclo and geoana kernels have equivalent runtimes.

santisoler · 2023-12-07T21:35:55Z

I just merged this! 🚀

Thanks a lot for the reviews, help and comments. I'm looking forward to open the PR for the mag simulation now!

I'll keep this branch up for a while, so I don't break any environment in case someone was using it. But I'll eventually delete it, so please install the latest version on main or wait for the next release (v0.21.0).

lheagy · 2023-12-07T21:36:55Z

💥 So excited to see this come in!

* main: Fix minor flake8 warning (simpeg#1307) Gravity simulation with Choclo as engine (simpeg#1285)

santisoler added 18 commits September 14, 2023 16:50

Add Choclo as alternative engine in gravity simulation

73f3809

Fix typo

a320069

Improve docstrings

3c9a8c4

Raise error if survey has not source_field attribute

d7ece2e

Allow numba to automatically cast matrix elements

3f3bc0f

Don't force values to be np.float32, Numba will cast them automatically when trying to write to the preallocated matrix.

Define decorated functions as global private ones

8cbcda4

When defining the decorated functions as attributes of the class, then every new simulation is created, they need to be recompiled.

Start adding tests for the new features

c4f4702

Add test for guv

0b064b2

Test gravity inversion with both engines

9cab149

Move the test for gravity inversions to pure pytest. Use parametrization to run the same test using geoana and choclo. Add a `random_seed` to `make_synthetic_data` to avoid tests failing randomly.

Move conversion factor to a private function

627f03a

Specify type of warning

67bea68

Extend tests

8148000

Add tests for sanity checks of the engine and some other attributes. Add tests for running the forward in parallel and serial. Add test for conversion factor.

Use sensitivity_dtype for the fields array in _forward

634d3a1

Parametrize fixture to test with Tensor and Tree meshes

51e1616

Simplify some tests

5a733de

Add tests for sensitivity_dtype

3c0c85f

Simplify some tests with the simple_mesh fixture

6d29ed7

Remove old tests for the simulation

85e678a

santisoler added 5 commits September 20, 2023 09:39

Remove trivial assignment in refactorized test

c7e74ad

Revert make_synthetic_data method

dfec5d8

Remove the `random_seed` argument added to the `make_synthetic_data` method. The change was moved to a separate PR to check the reason why some tests were failing in CI.

Revert "Revert make_synthetic_data method"

72ca755

This reverts commit dfec5d8.

Merge branch 'main' into gravity-choclo

3c67089

Run black after solving conflicts

c0c0788

santisoler marked this pull request as ready for review September 27, 2023 17:25

jedman reviewed Oct 7, 2023

View reviewed changes

SimPEG/potential_fields/gravity/simulation.py Show resolved Hide resolved

SimPEG/potential_fields/gravity/simulation.py Show resolved Hide resolved

Set n_processes to None if passed and using choclo

fc636d4

jedman approved these changes Oct 13, 2023

View reviewed changes

santisoler commented Oct 24, 2023

View reviewed changes

SimPEG/potential_fields/gravity/_numba_functions.py Outdated Show resolved Hide resolved

santisoler commented Oct 24, 2023

View reviewed changes

SimPEG/potential_fields/gravity/_numba_functions.py Outdated Show resolved Hide resolved

Remove uninformative lines in docstrings

7694b09

domfournier self-requested a review October 25, 2023 12:07

domfournier approved these changes Oct 25, 2023

View reviewed changes

jcapriot requested changes Oct 25, 2023

View reviewed changes

Replace private method for TensorMesh.cell_nodes

97d4789

Use the new `discretize.TensorMesh.cell_nodes` instead of the private method that was in the simulation. Increase the minimum version of discretize in `setup.py` and `environment_test.yml`.

jcapriot requested changes Nov 6, 2023

View reviewed changes

Replace array in docstrings for numpy.ndarray

d461204

Rename choclo_parallel to numba_parallel

f9a26a0

santisoler mentioned this pull request Nov 22, 2023

Pin SimPEG version once new gravity simulation PR is merged santisoler/simpeg-benchmarks-gravity#4

Closed

santisoler mentioned this pull request Nov 22, 2023

Jpe/update mag tests santisoler/simpeg#1

Merged

Merge branch 'main' into gravity-choclo

edab185

lheagy added this to the 0.21.0 milestone Nov 29, 2023

santisoler merged commit eb07ef1 into main Dec 7, 2023

thibaut-kobold added a commit to KoBoldMetals/simpeg that referenced this pull request Dec 19, 2023

Merge branch 'main' into update_experimental

e90e0cb

* main: Fix minor flake8 warning (simpeg#1307) Gravity simulation with Choclo as engine (simpeg#1285)

This was referenced Mar 1, 2024

Mention SimPEG collaboration in docs fatiando/choclo#65

Merged

Magnetic simulation with Choclo as engine #1321

Merged

santisoler deleted the gravity-choclo branch September 17, 2025 23:40

Gravity simulation with Choclo as engine #1285

Gravity simulation with Choclo as engine #1285

Uh oh!

Conversation

santisoler commented Sep 19, 2023

Summary

PR Checklist

Reference issue

What does this implement/fix?

Additional information

Uh oh!

codecov bot commented Sep 19, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

thibaut-kobold commented Sep 28, 2023

Uh oh!

santisoler commented Sep 28, 2023

Uh oh!

Uh oh!

Uh oh!

jedman left a comment

Choose a reason for hiding this comment

Uh oh!

santisoler commented Oct 24, 2023

Uh oh!

Uh oh!

Uh oh!

domfournier left a comment

Choose a reason for hiding this comment

Uh oh!

domfournier commented Oct 25, 2023

Uh oh!

jcapriot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jcapriot Oct 25, 2023

Choose a reason for hiding this comment

Uh oh!

santisoler Oct 25, 2023

Choose a reason for hiding this comment

Uh oh!

santisoler commented Oct 25, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

santisoler commented Nov 3, 2023

Uh oh!

jcapriot left a comment

Choose a reason for hiding this comment

Uh oh!

santisoler commented Nov 20, 2023

Uh oh!

jcapriot commented Nov 20, 2023

Uh oh!

lheagy commented Nov 20, 2023

Uh oh!

jcapriot commented Nov 21, 2023

Uh oh!

santisoler commented Nov 22, 2023

Uh oh!

santisoler commented Dec 7, 2023

Uh oh!

lheagy commented Dec 7, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

codecov bot commented Sep 19, 2023 •

edited

Loading

santisoler commented Oct 25, 2023 •

edited

Loading