adept
adept copied to clipboard
Added test for two-stream instability growth rate optimization
This PR did the following:
- Added a script for optimizing the growth of the two-stream instability as measured by the average
e2from the 10th to 200th timestep. Measuring the growth rate by adiff-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
numpyoperations tojax.numpyto the vlasov1d helpers file for differentiability - Changed the
__call__function in theBaseVlasov1Dmodule 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.