Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ML surrogate example #482

Merged
merged 17 commits into from
Dec 13, 2023
Merged

Conversation

RTSandberg
Copy link
Member

@RTSandberg RTSandberg commented Dec 7, 2023

An example using PyTorch neural network surrogates in ImpactX. The neural networks are trained on data from WarpX simulations.

This is an example of 9 stages of Laser-plasma accelerator elements coupled with simple 3 cm beam transport between stages using idealized plasma lenses as illustrated in the below schematic.
schema_9_stages

Some initial phase space projections are as shown:
initial_phase_spaces

At the end of the 9 stages, the beam evolves to this:
stage_8_phase_spaces

The beam moments evolve as shown:
lpa_ml_surrogate_moments

Actions

  • move data to Zenodo archive
  • add pytorch to CI @ax3l
  • (optional) consider having data on HuggingFace

@RTSandberg RTSandberg changed the title Set up ml surrogate example ML surrogate example Dec 7, 2023
@RTSandberg RTSandberg added this to the Advanced Methods (SciDAC) milestone Dec 7, 2023
@ax3l ax3l self-assigned this Dec 9, 2023
Copy link
Member

@ax3l ax3l left a comment

Choose a reason for hiding this comment

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

Thank you, excellent!
I added a first round of review! 🙏

examples/pytorch_surrogate_model/README.rst Outdated Show resolved Hide resolved
examples/pytorch_surrogate_model/README.rst Outdated Show resolved Hide resolved
examples/pytorch_surrogate_model/README.rst Outdated Show resolved Hide resolved
examples/pytorch_surrogate_model/README.rst Outdated Show resolved Hide resolved
examples/pytorch_surrogate_model/run_ml_surrogate.py Outdated Show resolved Hide resolved
Sigmoid = 4


def get_enum_type(type_to_test, EnumClass):

Check notice

Code scanning / CodeQL

Explicit returns mixed with implicit (fall through) returns Note

Mixing implicit and explicit returns may indicate an error as implicit returns always return None.
def download_and_unzip(url, data_dir):
request.urlretrieve(url, data_dir)
with tarfile.open(data_dir) as tar_dataset:
tar_dataset.extractall()

Check failure

Code scanning / CodeQL

Arbitrary file write during tarfile extraction High

This file extraction depends on a
potentially untrusted source
.
@ax3l ax3l merged commit 3759e52 into ECP-WarpX:development Dec 13, 2023
14 checks passed
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