[WIP] Wrapper for simplified z-t conversions

lasy/profiles/from_openpmd_profile.py

@@ -4,7 +4,7 @@
 from openpmd_viewer import OpenPMDTimeSeries
 from .from_array_profile import FromArrayProfile
 from lasy.utils.laser_utils import field_to_envelope, create_grid
-from lasy.utils.openpmd_input import reorder_array
+from lasy.utils.openpmd_input import refactor_array
 
 
 class FromOpenPMDProfile(FromArrayProfile):
@@ -97,7 +97,7 @@ def __init__(
                 phase_unwrap_1d = False
             axes_order = ["r", "t"]
 
-        F, axes = reorder_array(F, m, dim)
+        F, axes = refactor_array(F, m, dim)
 
         # If array does not contain the envelope but the electric field,
         # extract the envelope with a Hilbert transform

lasy/utils/laser_utils.py

@@ -6,6 +6,14 @@
 from axiprop.lib import PropagatorFFT2, PropagatorResampling
 from axiprop.containers import ScalarFieldEnvelope
 
+try:
+    from axiprop.lib import PropagatorFFT2, PropagatorResampling
+    from axiprop.containers import ScalarFieldEnvelope
+
+    axiprop_installed = True
+except ImportError:
+    axiprop_installed = False
+
 from .grid import Grid
 
 
@@ -842,3 +850,23 @@ def import_from_z(dim, grid, omega0, field_z, z_axis, z0=0.0, t0=0.0, backend="N
         transform_data *= np.exp(-1j * z_axis[0] * (k_z[:, None, None] - omega0 / c))
         grid.field = prop.z2t(transform_data, t_axis, z0=z0, t0=t0).T
         grid.field *= np.exp(1j * (z0 / c + t_axis) * omega0)
+
+
+def convert_z_to_t(array, axes, dim, dummy=False, omega0=None):
+    t = (axes["z"] - axes["z"][0]) / c
+    if dim == "xyt":
+        axes = {"x": axes["x"], "y": axes["y"], "t": t}
+    else:
+        axes = {"r": axes["r"], "t": t}
+
+    if dummy:
+        # Flip to get complex envelope in t assuming z = -c*t
+        array = np.flip(array, axis=-1)
+
+        return array, axes
+
+    else:
+        grid = create_grid(array, axes, dim)
+        assert omega0 is not None
+        import_from_z(dim, grid, omega0, array, axes["z"], z0=0.0, t0=0.0, backend="NP")
+        return grid.field, axes

lasy/utils/openpmd_input.py

@@ -1,9 +1,9 @@
 import numpy as np
-import scipy.constants as ct
+from .laser_utils import dummy_z_to_t
 
 
-def reorder_array(array, md, dim):
-    """Reorder an openPMD array to the lasy representation.
+def refactor_array(array, md, dim):
+    """Refactor an openPMD array to the lasy representation.
 
     Parameters
     ----------
@@ -22,9 +22,16 @@ def reorder_array(array, md, dim):
         A dictionary with the lasy axes information for the array.
     """
     if dim == "xyt":
-        return reorder_array_xyt(array, md)
+        array, axes = reorder_array_xyt(array, md)
     else:
-        return reorder_array_rt(array, md)
+        array, axes = reorder_array_rt(array, md)
+
+    if "z" in axes.keys:
+        # Data uses z representation, need to convert to
+        # t representation used inside lasy
+        array, axes = convert_z_to_t(array, axes, dim, dummy=True)
+
+    return array, axes
 
 
 def reorder_array_xyt(array, md):
@@ -51,16 +58,13 @@ def reorder_array_xyt(array, md):
         {0: "t", 1: "y", 2: "x"},
     ]
 
-    if md.axes in [{0: "z", 1: "y", 2: "x"}, {0: "t", 1: "y", 2: "x"}]:
+    if md.axes == {0: "z", 1: "y", 2: "x"}:
         array = array.swapaxes(0, 2)
+        axes = {"x": md.x, "y": md.y, "z": md.z}
+    elif md.axes == {0: "t", 1: "y", 2: "x"}:
+        array = array.swapaxes(0, 2)
+        axes = {"x": md.x, "y": md.y, "t": md.t}
 
-    if "z" in md.axes.values():
-        t = (md.z - md.z[0]) / ct.c
-        # Flip to get complex envelope in t assuming z = -c*t
-        array = np.flip(array, axis=-1)
-    else:
-        t = md.t
-    axes = {"x": md.x, "y": md.y, "t": t}
     return array, axes
 
 
@@ -88,20 +92,16 @@ def reorder_array_rt(array, md):
         {0: "t", 1: "r"},
     ]
 
-    if md.axes in [{0: "z", 1: "r"}, {0: "t", 1: "r"}]:
+    if md.axes == {0: "z", 1: "r"}:
         array = array.swapaxes(0, 1)
-
-    if "z" in md.axes.values():
-        t = (md.z - md.z[0]) / ct.c
-        # Flip to get complex envelope in t assuming z = -c*t
-        array = np.flip(array, axis=-1)
-    else:
-        t = md.t
-    r = md.r[md.r.size // 2 :]
-    axes = {"r": r, "t": t}
+        axes = {"r": md.r[md.r.size // 2 :], "z": md.z}
+    if md.axes == {0: "t", 1: "r"}:
+        array = array.swapaxes(0, 1)
+        axes = {"r": md.r[md.r.size // 2 :], "t": md.t}
 
     array = 0.5 * (
         array[array.shape[0] // 2 :, :]
         + np.flip(array[: array.shape[0] // 2, :], axis=0)
     )
+
     return array, axes