adept icon indicating copy to clipboard operation
adept copied to clipboard

Added test for two-stream instability growth rate optimization

Open physicistphil opened this issue 3 months ago • 0 comments

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.

physicistphil avatar Oct 10 '25 23:10 physicistphil