Skip to content

feat(python): add differentiable PyTorch gravity interface#54

Merged
darioizzo merged 4 commits into
esa:mainfrom
sasso-effe:feat/torch-submodule
Jun 17, 2026
Merged

feat(python): add differentiable PyTorch gravity interface#54
darioizzo merged 4 commits into
esa:mainfrom
sasso-effe:feat/torch-submodule

Conversation

@sasso-effe

Copy link
Copy Markdown
Contributor

Adds polyhedral_gravity.torch. Includes tests, notebooks, and CI coverage.

Closes #53

@sasso-effe

Copy link
Copy Markdown
Contributor Author

The CI is failing on two pipelines:

  • There are 2 errors in the "C/C++ Library on windows-latest" pipeline.
    • One of them should have been fixed by my latest commit.
    • The other one is pre-existing and due to Windows runner having updated CMake version to 4.3, However, I am not sure whether it is blocking or not, so I would try to run the pipeline again. @darioizzo could you do it?
  • 3 tests out of 50 are failing in the newly added PyTorch Interface (CPU) pipeline. This is due to the bug in the C++ implementation that I proposed to fix in fix: correct vertex singularity detection in computeSingularityTerms #52. I can either:
    • Wait for that PR to be merged, and then rebase this branch. @schuhmaj do you have an ETA for that?
    • Modify the test suite to exclude those singularity cases.
    • Mark them to be skipped for now, and enable them again when the fix is merged.

@darioizzo darioizzo requested review from Copilot and removed request for Copilot June 15, 2026 16:07
@sasso-effe

Copy link
Copy Markdown
Contributor Author

My last commit definitely fixes the build on Windows. This time I tested it on my fork (See here). Sorry if I did not think about doing it before.

PyTorch tests are still going to fail, though. For those, we should do one of the things I proposed before:

I can either:

  • Wait for that PR to be merged, and then rebase this branch. @schuhmaj do you have an ETA for that?
  • Modify the test suite to exclude those singularity cases.
  • Mark them to be skipped for now, and enable them again when the fix is merged.

@sasso-effe sasso-effe force-pushed the feat/torch-submodule branch from c431d2d to 9e6201e Compare June 17, 2026 08:18
@darioizzo darioizzo merged commit ffd7561 into esa:main Jun 17, 2026
7 checks passed
@schuhmaj

schuhmaj commented Jun 17, 2026

Copy link
Copy Markdown
Collaborator

@sasso-effe Thanks for the contribution. The idea to add some implementation that allows us to determine gradients w.r.t. position is great! Though I do have some comments (see below) for future reference. I would love to hear if this is interesting to you. Also, thanks a lot (!) for reporting the runtimes in the notebook AND specifying the hardware (CPU <=> GPU comparisons are in other places way too often reduced to unfair speedup without mentioning the hardware in other places 😢)

@darioizzo Currently, this breaks existing functionality/ is not 100% compliant with code standard/ docs:

  • JOSS requires proper documentation; the new modules are not mentioned in the docs https://esa.github.io/polyhedral-gravity-model/ - you could link the jupyter notebooks here
  • The variable/parameter naming is not compliant with PEP 8, something like parameter G should be renamed to gravitational_constant (for private variables inside functions, it's imo okay to use mathematical notation, but for functions it can become quiet messy when one sees calls like fn(F=F, G=G) because constants and parameter names are less distinguishable)
  • Existing functionality breaks! Namely, the magic attributes which are now in _core, e.g., polyhedral_gravity.__version__ --> This needs to be refactored!
    • A run of the script/measure_runtime.py should always be obligatory to check if the runtime is not affected. We do not do this in the CI, since it would be a flaky test

I know that I am not always super-quick here, but I was on vacation, and I am mostly maintaining this in my free time, so things can proceed a bit slower. If a feature is urgent, one can always leave it on a branch for a while or inside a fork, but making main unstable is not best practice.

@sasso-effe If you could fix the above-mentioned things, we can make it a release to simplify installation via conda and pip with version 3.4 - maybe with some ultimate fix for the epsilon handling as well.


@sasso-effe We do have a proper GPU implementation lying around (actually multiple 😄- that is my current research work) and are currently undergoing testing on Nvidia, AMD, and Intel GPU hardware to have the polyhedral model portable and fast on any hardware (https://github.com/schuhmaj/performance-portability-benchmark/tree/main/src/polyhedralGravity).

So the way to go would be to update the current build process and use nanobind, which allows for easier exchange of nd-arrays between Python and C++. And not just numpy, but also cupy, torch or jax arrays (https://nanobind.readthedocs.io/en/latest/api_extra.html#n-dimensional-array-type).

Plug in the above-mentioned GPU implementations, and we have a model that can run anywhere within any Python framework and on any hardware, fast and optimized to our needs.
Just a backward pass would need to be implemented manually. Torch cannot automatically deduce what a C++ function is doing, unfortunately.

Now my question: @sasso-effe, I think the speed gain on GPUs is probably good to take, but I presume you also require the autodiff for your current work? Or is it just the GPU acceleration?
Then we leave this torch submodule in for the moment, until we have a fully fleshed out C++ solution that supports everything.

@sasso-effe

Copy link
Copy Markdown
Contributor Author

@schuhmaj, first of all, sorry for having left main unstable. I have almost fixed all the points you raised, and I will push them on a new branch and open a PR tomorrow. I'd be glad to help with the version 3.4 release.

Regarding your question, I did this implementation because I needed the autodiff for my research project; the GPU acceleration was just a nice side effect. I've paused my efforts on the mentioned project because I realized it was not leading anywhere, but I had already done all the implementation, and I thought it was nice to contribute it to this repo to not lose it. Maybe it can be helpful to someone else.

I will have a look at the repo you linked!

@schuhmaj

schuhmaj commented Jun 17, 2026

Copy link
Copy Markdown
Collaborator

@schuhmaj, first of all, sorry for having left main unstable. I have almost fixed all the points you raised, and I will push them on a new branch and open a PR tomorrow. I'd be glad to help with the version 3.4 release.

Cool! No worries, as long as the release is fine, we're good ;)

Regarding your question, I did this implementation because I needed the autodiff for my research project; the GPU acceleration was just a nice side effect. I've paused my efforts on the mentioned project because I realized it was not leading anywhere, but I had already done all the implementation, and I thought it was nice to contribute it to this repo to not lose it. Maybe it can be helpful to someone else.

But good to have it in. For the PyPI, we can enable the torch dependency somewhere in the build process in the setup.py to allow the easy installation in one go à la pip install 'polyhedral-gravity[torch]' . Though we should definitely write somewhere that this is only useful if one has a GPU/ or wants the autodiff capabilities. Maybe also migrate this just into the documentation to not confuse the user with two APIs (torch and classic) in the README.md which should be somewhat lightweight.

I will have a look at the repo you linked!

Mostly experimental, at least from the SE side of things; it's more a performance study. Most likely the Kokkos implementation will come to this repository at some point (and replace the current thrust dependency - or we properly enable Cuda-Thrust here as well)

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

Successfully merging this pull request may close these issues.

feat: add differentiable PyTorch interface (polyhedral_gravity.torch)

3 participants