Improvements in the reader from openPMD file

lasy/laser.py

@@ -398,7 +398,9 @@ def import_from_z(self, field_z, z_axis, z0=0.0, t0=0.0, backend="NP"):
             self.grid.field = prop.z2t(transform_data, t_axis, z0=z0, t0=t0).T
             self.grid.field *= np.exp(1j * (z0 / c + t_axis) * omega0)
 
-    def write_to_file(self, file_prefix="laser", file_format="h5"):
+    def write_to_file(
+        self, file_prefix="laser", file_format="h5", save_as_vector_potential=False
+    ):
         """
         Write the laser profile + metadata to file.
 
@@ -409,6 +411,10 @@ def write_to_file(self, file_prefix="laser", file_format="h5"):
 
         file_format : string
             Format to be used for the output file. Options are ``"h5"`` and ``"bp"``.
+
+        save_as_vector_potential : bool (optional)
+            Whether the envelope is converted to normalized vector potential
+            before writing to file.
         """
         write_to_openpmd_file(
             self.dim,
@@ -417,4 +423,5 @@ def write_to_file(self, file_prefix="laser", file_format="h5"):
             self.grid,
             self.profile.lambda0,
             self.profile.pol,
+            save_as_vector_potential,
         )

lasy/profiles/from_array_profile.py

@@ -33,25 +33,51 @@ class FromArrayProfile(Profile):
         ['x', 'y', 't'] and ['t', 'y', 'x'].
     """
 
-    def __init__(self, wavelength, pol, array, axes, axes_order=["x", "y", "t"]):
+    def __init__(self, wavelength, pol, array, dim, axes, axes_order=["x", "y", "t"]):
         super().__init__(wavelength, pol)
 
+        assert dim in ["xyt", "rt"]
         self.axes = axes
-        assert axes_order in [["x", "y", "t"], ["t", "y", "x"]]
+        self.dim = dim
 
-        if axes_order == ["t", "y", "x"]:
-            self.array = np.swapaxes(array, 0, 2)
-        else:
-            self.array = array
+        if dim == "xyt":
+            assert axes_order in [["x", "y", "t"], ["t", "y", "x"]]
+
+            if axes_order == ["t", "y", "x"]:
+                self.array = np.swapaxes(array, 0, 2)
+            else:
+                self.array = array
+
+            self.field_interp = RegularGridInterpolator(
+                (axes["x"], axes["y"], axes["t"]),
+                array,
+                bounds_error=False,
+                fill_value=0.0,
+            )
 
-        self.field_interp = RegularGridInterpolator(
-            (axes["x"], axes["y"], axes["t"]),
-            array,
-            bounds_error=False,
-            fill_value=0.0,
-        )
+        else:  # dim = "rt"
+            assert axes_order in [["r", "t"], ["t", "r"]]
+
+            if axes_order == ["t", "r"]:
+                self.array = np.swapaxes(array, 0, 2)
+            else:
+                self.array = array
+
+            # The first point is at dr/2.
+            # To avoid problems within the first cell, fill_value in None,
+            # so the value is interpolated. This might cause issues at the upper
+            # boundary.
+            self.field_interp = RegularGridInterpolator(
+                (axes["r"], axes["t"]),
+                array,
+                bounds_error=False,
+                fill_value=None,
+            )
 
     def evaluate(self, x, y, t):
         """Return the envelope field of the scaled profile."""
-        envelope = self.field_interp((x, y, t))
+        if self.dim == "xyt":
+            envelope = self.field_interp((x, y, t))
+        else:
+            envelope = self.field_interp((np.sqrt(x**2 + y**2), t))
         return envelope

lasy/profiles/from_openpmd_profile.py

@@ -4,6 +4,8 @@
 import openpmd_api as io
 from openpmd_viewer import OpenPMDTimeSeries
 from .from_array_profile import FromArrayProfile
+from lasy.utils.laser_utils import get_frequency
+from lasy.utils.grid import Grid
 
 
 class FromOpenPMDProfile(FromArrayProfile):
@@ -43,56 +45,175 @@ class FromOpenPMDProfile(FromArrayProfile):
         Prefix of the openPMD file from which the envelope is read.
         Only used when envelope=True.
         The provided iteration is read from <path>/<prefix>_%T.h5.
+
+    theta : float or None, optional
+        Only used if the openPMD input is in thetaMode geometry.
+        Directly passed to openpmd_viewer.OpenPMDTimeSeries.get_field.
+        The angle of the plane of observation, with respect to the x axis
+        If `theta` is not None, then this function returns a 2D array
+        corresponding to the plane of observation given by `theta`;
+        otherwise it returns a full 3D Cartesian array.
+
+    phase_unwrap_1d : boolean (optional)
+        Whether the phase unwrapping is done in 1D.
+        This is not recommended, as the unwrapping will not be accurate,
+        but it might be the only practical solution when dim is 'xyt'.
+        If None, it is set to False for dim = 'rt' and True for dim = 'xyt'.
+
+    verbose : boolean (optional)
+        Whether to print extended information.
     """
 
     def __init__(
-        self, path, iteration, pol, field, coord=None, envelope=False, prefix=None
+        self,
+        path,
+        iteration,
+        pol,
+        field,
+        coord=None,
+        envelope=False,
+        prefix=None,
+        theta=None,
+        phase_unwrap_1d=None,
+        verbose=False,
     ):
         ts = OpenPMDTimeSeries(path)
-        F, m = ts.get_field(iteration=iteration, field=field, coord=coord, theta=None)
-        assert m.axes in [
-            {0: "x", 1: "y", 2: "z"},
-            {0: "z", 1: "y", 2: "x"},
-            {0: "x", 1: "y", 2: "t"},
-            {0: "t", 1: "y", 2: "x"},
-        ]
-
-        if m.axes in [{0: "z", 1: "y", 2: "x"}, {0: "t", 1: "y", 2: "x"}]:
-            F = F.swapaxes(0, 2)
-
-        if "z" in m.axes.values():
-            t = (m.z - m.z[0]) / c
-        else:
-            t = m.t
-        axes = {"x": m.x, "y": m.y, "t": t}
-
-        # If array does not contain the envelope but the electric field,
-        # extract the envelope with a Hilbert transform
-        if envelope == False:
-            # Assumes z is last dimension!
-            h = hilbert(F)
+        F, m = ts.get_field(iteration=iteration, field=field, coord=coord, theta=theta)
+
+        if theta is None:  # Envelope obtained from the full 3D array
+            assert m.axes in [
+                {0: "x", 1: "y", 2: "z"},
+                {0: "z", 1: "y", 2: "x"},
+                {0: "x", 1: "y", 2: "t"},
+                {0: "t", 1: "y", 2: "x"},
+            ]
+
+            dim = "xyt"
+
+            if m.axes in [{0: "z", 1: "y", 2: "x"}, {0: "t", 1: "y", 2: "x"}]:
+                F = F.swapaxes(0, 2)
 
             if "z" in m.axes.values():
-                # Flip to get complex envelope in t assuming z = -c*t
-                h = np.flip(h, axis=2)
-
-            # Get central wavelength from array
-            phase = np.unwrap(np.angle(h))
-            omg0_h = np.average(np.gradient(-phase, t, axis=-1), weights=np.abs(h))
-            wavelength = 2 * np.pi * c / omg0_h
-            array = h * np.exp(1j * omg0_h * t)
-        else:
-            s = io.Series(path + "/" + prefix + "_%T.h5", io.Access.read_only)
-            it = s.iterations[iteration]
-            omg0 = it.meshes["laserEnvelope"].get_attribute("angularFrequency")
-            wavelength = 2 * np.pi * c / omg0
-            array = F
-
-        axes_order = ["x", "y", "t"]
+                t = (m.z - m.z[0]) / c
+            else:
+                t = m.t
+            axes = {"x": m.x, "y": m.y, "t": t}
+
+            if phase_unwrap_1d is None:
+                phase_unwrap_1d = True
+
+            # If array does not contain the envelope but the electric field,
+            # extract the envelope with a Hilbert transform
+            if not envelope:
+                # Assumes z is last dimension!
+                h = hilbert(F)
+
+                if "z" in m.axes.values():
+                    # Flip to get complex envelope in t assuming z = -c*t
+                    h = np.flip(h, axis=-1)
+
+                # Get central wavelength from array
+                lo = (axes["x"][0], axes["y"][0], axes["t"][0])
+                hi = (axes["x"][-1], axes["y"][-1], axes["t"][-1])
+                npoints = (axes["x"].size, axes["y"].size, axes["t"].size)
+                grid = Grid(dim, lo, hi, npoints)
+                assert np.all(grid.axes[0] == axes["x"])
+                assert np.all(grid.axes[1] == axes["y"])
+                assert np.all(grid.axes[2] == axes["t"])
+                assert grid.field.shape == h.shape
+                grid.field = h
+                omg_h, omg0_h = get_frequency(
+                    grid,
+                    dim,
+                    is_envelope=False,
+                    is_hilbert=True,
+                    phase_unwrap_1d=phase_unwrap_1d,
+                )
+                wavelength = 2 * np.pi * c / omg0_h
+                array = h * np.exp(1j * omg0_h * t)
+            else:
+                s = io.Series(path + "/" + prefix + "_%T.h5", io.Access.read_only)
+                it = s.iterations[iteration]
+                omg0 = it.meshes["laserEnvelope"].get_attribute("angularFrequency")
+                wavelength = 2 * np.pi * c / omg0
+                array = F
+
+            axes_order = ["x", "y", "t"]
+
+        else:  # Envelope assumes axial symmetry processing RZ data
+            assert m.axes in [
+                {0: "r", 1: "z"},
+                {0: "z", 1: "r"},
+                {0: "r", 1: "t"},
+                {0: "t", 1: "r"},
+            ]
+
+            dim = "rt"
+
+            if m.axes in [{0: "z", 1: "r"}, {0: "t", 1: "r"}]:
+                F = F.swapaxes(0, 1)
+
+            if "z" in m.axes.values():
+                t = (m.z - m.z[0]) / c
+            else:
+                t = m.t
+            r = m.r[m.r.size // 2 :]
+            axes = {"r": r, "t": t}
+
+            F = 0.5 * (
+                F[F.shape[0] // 2 :, :] + np.flip(F[: F.shape[0] // 2, :], axis=0)
+            )
+
+            if phase_unwrap_1d is None:
+                phase_unwrap_1d = False
+            # If array does not contain the envelope but the electric field,
+            # extract the envelope with a Hilbert transform
+            if not envelope:
+                # Assumes z is last dimension!
+                h = hilbert(F)
+
+                if "z" in m.axes.values():
+                    # Flip to get complex envelope in t assuming z = -c*t
+                    h = np.flip(h, axis=-1)
+
+                # Get central wavelength from array
+                lo = (axes["r"][0], axes["t"][0])
+                hi = (axes["r"][-1], axes["t"][-1])
+                npoints = (axes["r"].size, axes["t"].size)
+                grid = Grid(dim, lo, hi, npoints, n_azimuthal_modes=1)
+                assert np.all(grid.axes[0] == axes["r"])
+                assert np.allclose(grid.axes[1], axes["t"], rtol=1.0e-14)
+                assert grid.field.shape == h[np.newaxis].shape
+                grid.field = h[np.newaxis]
+                omg_h, omg0_h = get_frequency(
+                    grid,
+                    dim=dim,
+                    is_envelope=False,
+                    is_hilbert=True,
+                    phase_unwrap_1d=phase_unwrap_1d,
+                )
+                wavelength = 2 * np.pi * c / omg0_h
+                array = h * np.exp(1j * omg0_h * t)
+            else:
+                s = io.Series(path + "/" + prefix + "_%T.h5", io.Access.read_only)
+                it = s.iterations[iteration]
+                omg0 = it.meshes["laserEnvelope"].get_attribute("angularFrequency")
+                wavelength = 2 * np.pi * c / omg0
+                array = F
+
+            axes_order = ["r", "t"]
+
+        if verbose:
+            print(
+                "Wavelength used in the definition of the envelope (nm):",
+                wavelength * 1.0e9,
+            )
+
         super().__init__(
             wavelength=wavelength,
             pol=pol,
             array=array,
+            dim=dim,
             axes=axes,
             axes_order=axes_order,
         )

lasy/utils/grid.py

@@ -29,7 +29,7 @@ class Grid:
         used in order to represent the laser field.
     """
 
-    def __init__(self, dim, lo, hi, npoints, n_azimuthal_modes):
+    def __init__(self, dim, lo, hi, npoints, n_azimuthal_modes=None):
         # Metadata
         ndims = 2 if dim == "rt" else 3
         assert dim in ["rt", "xyt"]