Skip to content

Conversation

@physicistphil
Copy link

@physicistphil physicistphil commented Oct 10, 2025

This PR did the following:

  • Added a script for optimizing the growth of the two-stream instability as measured by the average e2 from the 10th to 200th timestep. Measuring the growth rate by a diff-style metric may be too sensitive to numerical instabilities and fp64 precision
  • Added test conditions for this optimization target and for testing the symmetry of the optimized variables, which are the temperature and drift velocity of each beam
  • Changed numpy operations to jax.numpy to the vlasov1d helpers file for differentiability
  • Changed the __call__ function in the BaseVlasov1D module so that the inputs to be differentiated (i.e., state and cfg) are explicit (ope -- reverted this so it doesn't break other scripts)

When ran, the script creates an mlflow experiment "twostream-optimize" with child runs of "opt-iter-x".

Efforts were made to make the optimization appear functional-like despite the OOP paradigm. To Jax, it shouldn't make a difference whether self.state is exploited or not because the function to be differentiated is still a pure function: the __call__ takes in params and returns the simulation results. Nevertheless the optimization target and mutable parameters should be clear.

from diffrax import Solution
from jax import numpy as jnp
import matplotlib
matplotlib.use("Agg")
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This probably shouldn't be called here in the library module since it'll apply to all scripts that import the module. Can you move it to the test or the script that needs the headless plotting?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ope, fixed

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.

2 participants