Merge pull request #68 from AngelFP/add_new_tests

Add new tests

.github/workflows/test-package.yml

@@ -15,7 +15,10 @@ jobs:
         python-version: [3.7, 3.8, 3.9, '3.10', 3.11]
 
     steps:
-    - uses: actions/checkout@v4
+    - name: Checkout
+      uses: actions/checkout@v4
+    - name: Setup headless display
+      uses: pyvista/setup-headless-display-action@v2
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
       with:

setup.cfg

@@ -46,6 +46,7 @@ visualpic =
 test =
     pytest
     flake8
+    pyvista
 
 [options.entry_points]
 console_scripts =

tests/test_data_container.py

@@ -1,14 +1,51 @@
+import scipy.constants as ct
 from visualpic import DataContainer
 
 
+# Intensity
+def calculate_intensity(data_list, sim_geometry, sim_params):
+    """Recipe to calculate intensity as a derived field."""
+    if sim_geometry == "1d":
+        Ez = data_list[0]
+        E2 = Ez**2
+    elif sim_geometry == "2dcartesian":
+        Ez, Ex = data_list
+        E2 = Ez**2 + Ex**2
+    elif sim_geometry == "3dcartesian":
+        Ez, Ex, Ey = data_list
+        E2 = Ez**2 + Ex**2 + Ey**2
+    elif sim_geometry == "cylindrical":
+        raise NotImplementedError
+    elif sim_geometry == "thetaMode":
+        Ez, Er, Et = data_list
+        E2 = Ez**2 + Er**2 + Et**2
+    return ct.c * ct.epsilon_0 / 2 * E2
+
+
+intensity = {
+    "name": "I",
+    "units": "W/m^2",
+    "requirements": {
+        "1d": ["Ez"],
+        "2dcartesian": ["Ez", "Ex"],
+        "3dcartesian": ["Ez", "Ex", "Ey"],
+        "cylindrical": ["Ez", "Er"],
+        "thetaMode": ["Ez", "Er", "Et"],
+    },
+    "recipe": calculate_intensity,
+}
+
+
 def test_data_container():
     """Test basic functionality of the DataContainer using sample data."""
     data_path = "./test_data/example-3d/hdf5"
     diags = DataContainer("openpmd", data_path)
     diags.load_data()
+    diags.add_derived_field(intensity)
+
     available_fields = diags.get_list_of_fields()
     available_species = diags.get_list_of_species()
-    assert available_fields == ["Ex", "Ey", "Ez", "rho", "I", "A", "a"]
+    assert set(["Ex", "Ey", "Ez", "rho", "I"]).issubset(available_fields)
     assert available_species == ["electrons"]
     assert diags.get_list_of_species(required_data=["x"]) == ["electrons"]
     assert diags.get_list_of_species(required_data=["magic"]) == []
@@ -22,7 +59,10 @@ def test_data_container():
         species = diags.get_species(sp)
         for it in species.timesteps:
             comps = species.get_list_of_available_components()
-            assert comps == ["q", "m", "x", "y", "z", "px", "py", "pz", "w", "x_prime"]
+            assert set(["q", "m", "x", "y", "z", "px", "py", "pz", "w"]).issubset(comps)
+            sp_data = species.get_data(
+                time_step=it, components_list=comps, data_units="SI", time_units="SI"
+            )
             sp_data = species.get_data(it)
 
 

tests/test_vtk_visualizer.py

@@ -0,0 +1,23 @@
+from visualpic import DataContainer, VTKVisualizer
+
+
+def test_vtk_visualizer():
+    """Make a field and particle render, and save it to file."""
+    data_path = "./test_data/example-3d/hdf5"
+    diags = DataContainer("openpmd", data_path)
+    diags.load_data()
+    vis = VTKVisualizer(use_qt=True)
+    rho = diags.get_field("rho")
+    elec = diags.get_species("electrons")
+    vis.add_field(rho)
+    vis.add_species(elec)
+    vis.render_to_file(
+        vis.available_time_steps[0],
+        "./test_render_3d.png",
+        resolution=[400, 400],
+        ts_is_index=False,
+    )
+
+
+if __name__ == "__main__":
+    test_vtk_visualizer()