Major refactoring for JAX-style classes.#29
Open
EstevaoMGomes wants to merge 135 commits into
Open
Conversation
- Implemented `fo_integrators.py` for full orbit tracing with various methods and parameters. - Implemented `gc_integrators.py` for guiding center dynamics with adaptative and constant step sizes. - Enhanced `Tracing` class in `dynamics.py` to support multiple methods and step sizes.
…d adjust num_steps based on dt
… class for improved step size handling
…for adaptive step size
…ters for performance
…ance plots, and improve layout for better visualization
En/review eg/augmented lagrangian : Modifying constraints and augmented_lagrange formalismo to fit the custom_losses machinery added in eg/analysis PR #29
En/review eg/coil perturbations fix : This PR changes coils.py and coil_pertubation.py in order to allow for the stochastic optimization and coil perturbation machinery to work with the custom_losses machinery addeed in PR #29. Coils and surfaces scaling was also adjusted so that the scaling is properly treated by the custom_losses, by correctly evaluating losses(unscalled_dofs(scaled_dofs)).
En/review eg/dynamics fix : This PR fixes dynamics.py in order to work properly with PR #29 changes. A function to initialize particles around a given axis is also added.
En/review eg/losse val grad : This adds some small features to the custom_losses implementation of PR #29, with the most relevant being the ability to call val_and_grad method for this custom loss class. This allows for more efficieny in optimizers that require both quantities at the same time like LBFGSB.
…ective_functions_mod En/review eg/fields surfaces objective functions mod : This PR checks objective_function.py so that the loss functions there now follow the custom_losses dependencies added in PR #29.
…Tree traversal, fixing 'None is not a valid value for jnp.array' crash in tracing (from_simsopt and from_json paths)
Guard Coils.dofs_currents against None/bool sentinels (fixes tracing crash on #42)
The requested changes have been addressed so I am going to merge into eg/analysis.
…ive and correct simsopt tracing comparison (adaptive models, Dopri5, matched output points). Verified: 12 examples FAIL->PASS, zero regressions.
…mocks, near-axis shapes matched to pyqsc_jax), fix grad_pytree tracer leak in losses.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Refactor of coils & surfaces to be proper PyTrees;
Added a loss wrapper to differentiate with respect to the dogs (PyTree leaves);
Added analysis & validation of the code