Major refactoring for JAX-style classes.#29
Open
EstevaoMGomes wants to merge 132 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
…fty or -1. Adding docstrings to clealry state the scalling. Fixing the unflatten and flatten definitions for the surface pytree. So that custom_losses now see the scaled dofs, and unflatten unscales the dofs. This should make custom_losses take gradients against scaled variables while loss values will be still obtained by applying the loss function to the unscaled dofs-
…asses. Note that s fixed scaled option is added and explained on docstrings. In order to also have the intended behaviour in custom_losses/tree flatten/ unflatten interaction (that is loss(unflatten dofs) but gradient of scaled dofs) the current_scale behaviour had to be slighlty changed. Now it can be given as an input, and if it is None, the current_scale will be calculated at the beggining and kept frozen for the rest of the optimization. This is needed because otherwise the tracing gets confused and gives the wrong gradient.
…and currents dofs consistent with Coils/Curves
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.
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