CI: PyTorch Surrogate Example

Cover our PyTorch surrogate example in CI.
ECP-WarpX · Jan 13, 2025 · 436709a · 436709a
1 parent 337f25b
commit 436709a
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt
@@ -960,6 +960,7 @@ add_impactx_test(spectrometer.py
     OFF  # no plot script yet
 )
 
+
 # Chicane with CSR ###########################################################
 #
 if(ImpactX_FFT)
@@ -1056,3 +1057,16 @@ add_impactx_test(linac-segment.py
     examples/linac_segment/analysis_linac_segment.py
     OFF  # no plot script yet
 )
+
+
+# PyTorch Surrogate: Staged LPA ##############################################
+#
+add_impactx_test(pytorch_surrogate_model
+    examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py
+      OFF  # ImpactX MPI-parallel
+      ON   # ImpactX Python interface
+    examples/pytorch_surrogate_model/analyze_ml_surrogate_15_stage.py
+    examples/pytorch_surrogate_model/visualize_ml_surrogate_15_stage.py
+)
+label_impactx_test(pytorch_surrogate_model slow)
+label_impactx_test(pytorch_surrogate_model.MPI slow)