From d5a30d3f8e6bcde81d35e6f90708ec29f9fbdc2f Mon Sep 17 00:00:00 2001 From: ax3l Date: Mon, 29 Apr 2024 03:43:51 +0000 Subject: [PATCH] Update Stub Files --- examples/aperture/run_aperture.py | 1 + examples/fodo_tune/run_fodo_tune.py | 1 + examples/initialize_from_array/run_from_array.py | 1 + examples/optimize_triplet/run_triplet.py | 5 +++-- .../pytorch_surrogate_model/run_ml_surrogate_15_stage.py | 3 ++- src/python/impactx/__init__.pyi | 1 + tests/python/conftest.py | 3 ++- tests/python/test_charge_deposition.py | 3 ++- tests/python/test_dataframe.py | 1 + tests/python/test_impactx.py | 1 + tests/python/test_transformation.py | 1 + tests/python/test_xopt.py | 3 ++- 12 files changed, 18 insertions(+), 6 deletions(-) diff --git a/examples/aperture/run_aperture.py b/examples/aperture/run_aperture.py index 7e8cc73fc..490ee3878 100755 --- a/examples/aperture/run_aperture.py +++ b/examples/aperture/run_aperture.py @@ -7,6 +7,7 @@ # -*- coding: utf-8 -*- import amrex.space3d as amr + from impactx import ImpactX, distribution, elements # work-around for https://github.com/ECP-WarpX/impactx/issues/499 diff --git a/examples/fodo_tune/run_fodo_tune.py b/examples/fodo_tune/run_fodo_tune.py index 2315be15b..475b00dd5 100755 --- a/examples/fodo_tune/run_fodo_tune.py +++ b/examples/fodo_tune/run_fodo_tune.py @@ -7,6 +7,7 @@ # -*- coding: utf-8 -*- import amrex.space3d as amr + from impactx import ImpactX, distribution, elements sim = ImpactX() diff --git a/examples/initialize_from_array/run_from_array.py b/examples/initialize_from_array/run_from_array.py index d2400eb37..13a12c054 100644 --- a/examples/initialize_from_array/run_from_array.py +++ b/examples/initialize_from_array/run_from_array.py @@ -9,6 +9,7 @@ import amrex.space3d as amr import numpy as np import transformation_utilities as pycoord + from impactx import Config, ImpactX, elements ################ diff --git a/examples/optimize_triplet/run_triplet.py b/examples/optimize_triplet/run_triplet.py index 4663edbad..14eed93e9 100755 --- a/examples/optimize_triplet/run_triplet.py +++ b/examples/optimize_triplet/run_triplet.py @@ -7,11 +7,12 @@ # -*- coding: utf-8 -*- import amrex.space3d as amr -import impactx import numpy as np -from impactx import ImpactX, distribution, elements from scipy.optimize import minimize +import impactx +from impactx import ImpactX, distribution, elements + # Call MPI_Init and MPI_Finalize only once: if impactx.Config.have_mpi: from mpi4py import MPI # noqa diff --git a/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py b/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py index 0a6101148..2cf8e8f31 100644 --- a/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py +++ b/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py @@ -20,6 +20,8 @@ import numpy as np import scipy.optimize as opt +from surrogate_model_definitions import surrogate_model + from impactx import ( Config, CoordSystem, @@ -29,7 +31,6 @@ distribution, elements, ) -from surrogate_model_definitions import surrogate_model try: import torch diff --git a/src/python/impactx/__init__.pyi b/src/python/impactx/__init__.pyi index b3a8b9bce..0e2d3a9f3 100644 --- a/src/python/impactx/__init__.pyi +++ b/src/python/impactx/__init__.pyi @@ -17,6 +17,7 @@ from __future__ import annotations import os as os from amrex import space3d as amr + from impactx.impactx_pybind import ( Config, CoordSystem, diff --git a/tests/python/conftest.py b/tests/python/conftest.py index b58c32ec7..3dc1be213 100644 --- a/tests/python/conftest.py +++ b/tests/python/conftest.py @@ -3,9 +3,10 @@ import os import amrex.space3d as amr -import impactx import pytest +import impactx + if impactx.Config.have_mpi: from mpi4py import MPI # noqa diff --git a/tests/python/test_charge_deposition.py b/tests/python/test_charge_deposition.py index 216c81919..0719a09a4 100755 --- a/tests/python/test_charge_deposition.py +++ b/tests/python/test_charge_deposition.py @@ -11,11 +11,12 @@ import math import amrex.space3d as amr -import impactx import matplotlib.pyplot as plt import numpy as np from conftest import basepath +import impactx + def test_charge_deposition(save_png=True): """ diff --git a/tests/python/test_dataframe.py b/tests/python/test_dataframe.py index e4ca4734f..42dc0bb4c 100644 --- a/tests/python/test_dataframe.py +++ b/tests/python/test_dataframe.py @@ -11,6 +11,7 @@ import matplotlib.pyplot as plt import pytest + from impactx import ImpactX, amr, distribution, elements diff --git a/tests/python/test_impactx.py b/tests/python/test_impactx.py index b8cd6e4d7..d0b086a59 100755 --- a/tests/python/test_impactx.py +++ b/tests/python/test_impactx.py @@ -10,6 +10,7 @@ import numpy as np import pytest from conftest import basepath + from impactx import ImpactX, distribution, elements # FIXME in AMReX via https://github.com/AMReX-Codes/amrex/pull/3727 diff --git a/tests/python/test_transformation.py b/tests/python/test_transformation.py index 030f5ea72..93b52966c 100644 --- a/tests/python/test_transformation.py +++ b/tests/python/test_transformation.py @@ -8,6 +8,7 @@ import numpy as np import pytest + from impactx import CoordSystem, ImpactX, coordinate_transformation, distribution diff --git a/tests/python/test_xopt.py b/tests/python/test_xopt.py index 83ff23143..6607229eb 100644 --- a/tests/python/test_xopt.py +++ b/tests/python/test_xopt.py @@ -10,9 +10,10 @@ import importlib import amrex.space3d as amr -import impactx import numpy as np import pytest + +import impactx from impactx import ImpactX, distribution, elements # configure the test