More check on number of dimensions for rt (#311)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

lasy/profiles/from_openpmd_profile.py

@@ -95,8 +95,16 @@ def __init__(
             dim = "rt"
             axes_order = ["r", "t"]
 
+        # Set r and t axes to the right dimension for LASY.
+        # Note that F is assumed 2D here, no azimuthal data.
         F, axes = reorder_array(F, m, dim)
 
+        if dim == "rt" and F.ndim == 2:
+            F = F[np.newaxis, :]
+        elif F.ndim != 3:
+            raise ValueError(
+                f"F has the wrong number of dimensions: {F.ndim} (should be 3)"
+            )
         # If array does not contain the envelope but the electric field,
         # extract the envelope with a Hilbert transform
         if not is_envelope:

lasy/utils/laser_utils.py

@@ -683,6 +683,7 @@ def create_grid(array, axes, dim, is_envelope=True):
         assert np.all(grid.axes[0] == axes["x"])
         assert np.all(grid.axes[1] == axes["y"])
         assert np.allclose(grid.axes[2], axes["t"], rtol=1.0e-14)
+        assert array.ndim == 3, "Input array should be of dimension 3 [x, y, time]"
         grid.set_temporal_field(array)
     else:  # dim == "rt":
         lo = (axes["r"][0], axes["t"][0])
@@ -691,6 +692,9 @@ def create_grid(array, axes, dim, is_envelope=True):
         grid = Grid(dim, lo, hi, npoints, n_azimuthal_modes=1, is_envelope=is_envelope)
         assert np.all(grid.axes[0] == axes["r"])
         assert np.allclose(grid.axes[1], axes["t"], rtol=1.0e-14)
+        assert (
+            array.ndim == 3
+        ), "Input array should be of dimension 3 [modes, radius, time]"
         grid.set_temporal_field(array)
     return grid