Dark Mode

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Fourier-Chebyshev Fit to B for particle tracing#2012

Draft
YigitElma wants to merge 574 commits intomasterfrom
yge/particle-fit
Draft

Fourier-Chebyshev Fit to B for particle tracing#2012
YigitElma wants to merge 574 commits intomasterfrom
yge/particle-fit

Conversation

Copy link
Collaborator

YigitElma commented Nov 12, 2025 *
edited
Loading

  • Adds a new convenience class FourierChebyshevField which has a couple of methods to be compatible with optimization,
    • init : only resolution info
    • build : creates grid and transforms
    • fit : given params and profiles (these change during optimization), fits a new Fourier-Chebyshev series.
    • evaluate : Evaluate the series at requested point.

Example usage will be,

interpolater = FourierChebyshevField(L=eq.L_grid, M=eq.M_grid, N=eq.N_grid)
interpolater.build(eq)
data_fit = interpolater.fit(eq.params_dict, {"iota": eq.iota})

rtz, _, aux = _trace_particles(
field=interpolator, # normally you pass eq here
y0=x0,
model=model,
model_args=args,
ts=ts,
params=None, # normally you pass eq.params_dict here
max_steps=int((ts[-1] - ts[0]) / 1e-8),
min_step_size=1e-8,
stepsize_controller=PIDController(rtol=1e-4, atol=1e-4, dtmin=1e-8),
saveat=SaveAt(ts=ts),
solver=Tsit5(),
adjoint=RecursiveCheckpointAdjoint(),
event=event,
chunk_size=None,
options={}, # this passes extra stuff to field computation, not needed here
throw=False,
return_aux=True,
)

Todo:

  • trace_particles should initialize FourierChebyshevField and by default use it. Add a flag to use Equilibrium field
  • plot_particle_trajectories etc should also be able to use the interpolated field.
  • Look at the spectral convergence (ideally lower res for basis vectors might be good, but current code is simpler and with batching might not be inefficient)
  • Consider adding spline version too?

f0uriest and others added 30 commits September 10, 2024 11:05
...article based on the dimension of the initial conditions. `particle_tracer.py` and `optimization.py` are adapted to this change.

Optimization doesn't work yet
Automatically chooses single or multiparticle based on the dimension of
the initial conditions. `particle_tracer.py` and `optimization.py` are
adapted to this change.

Optimization doesn't work yet
YigitElma and others added 8 commits December 6, 2025 13:40
...#2027)

Updates the requirements on
[nbsphinx](https://github.com/spatialaudio/nbsphinx) to permit the
latest version.

Release notes

Sourced from href="nbsphinx">https://github.com/spatialaudio/nbsphinx/releases">nbsphinx's
releases.



nbsphinx 0.9.8


href="https://pypi.org/project/nbsphinx/0.9.8/https://pypi.org/project/nbsphinx/0.9.8/">https://pypi.org/project/nbsphinx/0.9.8/>


  • Re-enable Sphinx 8.2+

  • Support text/x-rst MIME type in raw cells

  • Support for mathjax4_config





Changelog

Sourced from href="nbsphinx">https://github.com/spatialaudio/nbsphinx/blob/master/NEWS.rst">nbsphinx's
changelog.



Version 0.9.8 -- 2025-11-28 -- PyPI__ -- diff__



  • Re-enable Sphinx 8.2+

  • Support text/x-rst MIME type in raw cells

  • Support for mathjax4_config


__ href="https://pypi.org/project/nbsphinx/0.9.8/https://pypi.org/project/nbsphinx/0.9.8/">https://pypi.org/project/nbsphinx/0.9.8/>
__ href="https://github.com/spatialaudio/nbsphinx/compare/0.9.7...0.9.8https://github.com/spatialaudio/nbsphinx/compare/0.9.7...0.9.8">https://github.com/spatialaudio/nbsphinx/compare/0.9.7...0.9.8>

Version 0.9.7 -- 2025-03-03 -- PyPI__ -- diff__



  • Disable Sphinx 8.2+ (for now)


__ href="https://pypi.org/project/nbsphinx/0.9.7/https://pypi.org/project/nbsphinx/0.9.7/">https://pypi.org/project/nbsphinx/0.9.7/>
__ href="https://github.com/spatialaudio/nbsphinx/compare/0.9.6...0.9.7https://github.com/spatialaudio/nbsphinx/compare/0.9.6...0.9.7">https://github.com/spatialaudio/nbsphinx/compare/0.9.6...0.9.7>

Version 0.9.6 -- 2024-12-24 -- PyPI__ -- diff__



  • Markdown: allow lists without leading blank line


__ href="https://pypi.org/project/nbsphinx/0.9.6/https://pypi.org/project/nbsphinx/0.9.6/">https://pypi.org/project/nbsphinx/0.9.6/>
__ href="https://github.com/spatialaudio/nbsphinx/compare/0.9.5...0.9.6https://github.com/spatialaudio/nbsphinx/compare/0.9.5...0.9.6">https://github.com/spatialaudio/nbsphinx/compare/0.9.5...0.9.6>

Version 0.9.5 -- 2024-08-13 -- PyPI__ -- diff__



  • Miscellaneous fixes


__ href="https://pypi.org/project/nbsphinx/0.9.5/https://pypi.org/project/nbsphinx/0.9.5/">https://pypi.org/project/nbsphinx/0.9.5/>
__ href="https://github.com/spatialaudio/nbsphinx/compare/0.9.4...0.9.5https://github.com/spatialaudio/nbsphinx/compare/0.9.4...0.9.5">https://github.com/spatialaudio/nbsphinx/compare/0.9.4...0.9.5>

Version 0.9.4 -- 2024-05-06 -- PyPI__ -- diff__



  • Require docutils >= 0.18.1

  • Minor fixes, documentation and CI updates


__ href="https://pypi.org/project/nbsphinx/0.9.4/https://pypi.org/project/nbsphinx/0.9.4/">https://pypi.org/project/nbsphinx/0.9.4/>
__ href="https://github.com/spatialaudio/nbsphinx/compare/0.9.3...0.9.4https://github.com/spatialaudio/nbsphinx/compare/0.9.3...0.9.4">https://github.com/spatialaudio/nbsphinx/compare/0.9.3...0.9.4>

Version 0.9.3 -- 2023-08-27 -- PyPI__ -- diff__



  • Fix gallery regression in Sphinx 7.2


__ href="https://pypi.org/project/nbsphinx/0.9.3/https://pypi.org/project/nbsphinx/0.9.3/">https://pypi.org/project/nbsphinx/0.9.3/>
__ href="https://github.com/spatialaudio/nbsphinx/compare/0.9.2...0.9.3https://github.com/spatialaudio/nbsphinx/compare/0.9.2...0.9.3">https://github.com/spatialaudio/nbsphinx/compare/0.9.2...0.9.3>

Version 0.9.2 -- 2023-05-24 -- PyPI__ -- diff__



  • Improve support for sphinx_immaterial theme

  • Improve support for links starting with #

  • Add support for in-text citations

  • LaTeX: Add support for admonition titles


__ href="https://pypi.org/project/nbsphinx/0.9.2/https://pypi.org/project/nbsphinx/0.9.2/">https://pypi.org/project/nbsphinx/0.9.2/>
__ href="https://github.com/spatialaudio/nbsphinx/compare/0.9.1...0.9.2https://github.com/spatialaudio/nbsphinx/compare/0.9.1...0.9.2">https://github.com/spatialaudio/nbsphinx/compare/0.9.1...0.9.2>

Version 0.9.1 -- 2023-03-14 -- PyPI__ -- diff__



  • pandoc: disable "smart" option only for version 2.0+




... (truncated)




Commits






Dependabot will resolve any conflicts with this PR as long as you don't
alter it yourself. You can also trigger a rebase manually by commenting
`@dependabot rebase`.

[//]: # (dependabot-automerge-start)
[//]: # (dependabot-automerge-end)

---


Dependabot commands and options



You can trigger Dependabot actions by commenting on this PR:
- `@dependabot rebase` will rebase this PR
- `@dependabot recreate` will recreate this PR, overwriting any edits
that have been made to it
- `@dependabot merge` will merge this PR after your CI passes on it
- `@dependabot squash and merge` will squash and merge this PR after
your CI passes on it
- `@dependabot cancel merge` will cancel a previously requested merge
and block automerging
- `@dependabot reopen` will reopen this PR if it is closed
- `@dependabot close` will close this PR and stop Dependabot recreating
it. You can achieve the same result by closing it manually
- `@dependabot show ignore conditions` will show all
of the ignore conditions of the specified dependency
- `@dependabot ignore this major version` will close this PR and stop
Dependabot creating any more for this major version (unless you reopen
the PR or upgrade to it yourself)
- `@dependabot ignore this minor version` will close this PR and stop
Dependabot creating any more for this minor version (unless you reopen
the PR or upgrade to it yourself)
- `@dependabot ignore this dependency` will close this PR and stop
Dependabot creating any more for this dependency (unless you reopen the
PR or upgrade to it yourself)




Signed-off-by: dependabot[bot]
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Rory Conlin
Copy link

review-notebook-app bot commented Dec 8, 2025

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

Base automatically changed from rc/particles to master December 8, 2025 22:43
Copy link
Contributor

github-actions bot commented Dec 9, 2025 *
edited
Loading

Memory benchmark result

| Test Name | %D | Master (MB) | PR (MB) | D (MB) | Time PR (s) | Time Master (s) |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
test_objective_jac_w7x | 3.78 % | 3.813e+03 | 3.957e+03 | 144.13 | 39.78 | 36.26 |
test_proximal_jac_w7x_with_eq_update | -1.27 % | 6.663e+03 | 6.579e+03 | -84.80 | 162.72 | 161.50 |
test_proximal_freeb_jac | 0.41 % | 1.316e+04 | 1.321e+04 | 53.79 | 84.83 | 83.43 |
test_proximal_freeb_jac_blocked | 0.41 % | 7.461e+03 | 7.492e+03 | 30.73 | 72.51 | 73.21 |
test_proximal_freeb_jac_batched | -1.25 % | 7.538e+03 | 7.444e+03 | -94.36 | 72.91 | 73.10 |
test_proximal_jac_ripple | -2.37 % | 3.608e+03 | 3.523e+03 | -85.50 | 65.70 | 65.71 |
test_proximal_jac_ripple_bounce1d | -0.08 % | 3.566e+03 | 3.564e+03 | -2.77 | 76.55 | 76.55 |
test_eq_solve | 0.60 % | 1.992e+03 | 2.004e+03 | 12.00 | 93.79 | 94.69 |

For the memory plots, go to the summary of Memory Benchmarks workflow and download the artifact.

Copy link

codecov bot commented Dec 11, 2025 *
edited
Loading

Codecov Report

Patch coverage is 77.69231% with 29 lines in your changes missing coverage. Please review.
Project coverage is 95.67%. Comparing base (3753ff0) to head (7b21b40).

Files with missing lines Patch % Lines
desc/particles.py 76.98% 29 Missing
Additional details and impacted files
@@ Coverage Diff @@
## master #2012 +/- ##
==========================================
- Coverage 95.75% 95.67% -0.09%
==========================================
Files 102 102
Lines 28344 28471 +127
==========================================
+ Hits 27142 27239 +97
- Misses 1202 1232 +30
Files with missing lines Coverage D
desc/compute/_basis_vectors.py 100.00% <100.00%> (o)
desc/particles.py 88.15% <76.98%> (-4.34%)

... and 2 files with indirect coverage changes

New features to boost your workflow:
  • Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Copy link
Collaborator Author

YigitElma commented Dec 18, 2025

For some reason, this suffers from GPU synchronisation miserably, and for more particles, it is slower than direct approach.

Copy link
Collaborator Author

YigitElma commented Feb 18, 2026

Maybe for these kinds of single node Grid creation, we can use a spline for only the radial part of zernike polynomials. The most expensive part of the transform calculation is the recurrence in the _jacobi which is not very efficient for single-node evaluation. We use vmap over r values but for particle tracing it doesn't perform well. The idea is to compute the zernike radial and its derivatives on a dense 1D grid, pass it to the Transform class, and instead of calling the zernike_radial it will evaluate the spline at the required points. The spline fit can be a callable like interpax.Interpolator1D.

r_dense = jnp.linspace(0.01, 1, 1000)
dense_zernike = zernike_radial(r=r_dense, l=l, m=m)
radial_spline = interpax.Interpolator1D(x=r_dense, f=dense_zernike)
# this should replace zernike call in transform.build
transforms = get_transforms(data_keys, eq, grid, jitable=True, radial_spline=radial_spline)

I haven't tested it yet, but it should potentially solve the problem other options had. FourierChebyshevField is good in practice, but it needs 13 set of higher res spectral coefficients, and each particle needs to access these coefficients, so the GPU memory is bottlenecked (hence 1 particle is faster but as you have more particles, it becomes slower). Interpolator3D should work in practice, but I think it is again bottlenecked by the memory.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

f0uriest f0uriest left review comments

At least 2 approving reviews are required to merge this pull request.

Assignees

YigitElma

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

6 participants