Add Genesis4 reader and writer

docs/examples/wavefront/wavefront.ipynb

@@ -32,6 +32,8 @@
     "\n",
     "import numpy as np\n",
     "\n",
+    "import os\n",
+    "\n",
     "import matplotlib.pyplot as plt"
    ]
   },
@@ -326,6 +328,75 @@
    "source": [
     "1 / W.sigma_x / 2"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1dd88a11-eeec-4d7d-8349-eb5d18c73a82",
+   "metadata": {},
+   "source": [
+    "# Genesis4 conversion"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "471e3478-f4df-45f5-bd49-df2dc971df01",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "W.write_genesis4(\"genesis4_field.h5\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3090d7d2-8f9d-4f14-8210-6946b4600b8e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "W2 = Wavefront.from_genesis4(\"genesis4_field.h5\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b95ae703-1826-4ce6-ab0e-982793155e48",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "np.allclose(W.Ex, W2.Ex)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3af5d343-667a-4f81-8e2f-f591e7765a16",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "W.plot()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "41b5335d-2ecf-47db-b86e-b27efbcc8b6a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "W2.plot()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a20d81dd-cedd-4969-83af-8d0f823786a7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Cleanup\n",
+    "os.remove(\"genesis4_field.h5\")"
+   ]
   }
  ],
  "metadata": {

pmd_beamphysics/interfaces/genesis.py

@@ -4,7 +4,7 @@
 from h5py import File, Group
 
 from pmd_beamphysics.statistics import twiss_calc
-from pmd_beamphysics.units import c_light, mec2, unit, write_unit_h5
+from pmd_beamphysics.units import c_light, mec2, unit, write_unit_h5, Z0
 
 # Genesis 1.3
 # -------------
@@ -582,3 +582,157 @@ def genesis4_par_to_data(h5, species="electron", smear=True):
     }
 
     return data
+
+
+def load_genesis4_fields(h5):
+    """
+    Copied from: https://github.com/slaclab/lume-genesis/blob/52dc2815bb0bb42c1e057d1393c6ada07b8580b3/genesis/version4/readers.py#L8
+    TODO: Point LUME-Geneis' fuction to here instead.
+
+    Loads the field data into memory from an open h5 handle.
+
+    Example usage:
+
+    import h5py
+    with h5py.File('rad_field.fld.h5', 'r') as h5:
+        dfl, param = load_genesis4_fields(h5)
+
+    Returns tuple (dfl, param) where
+
+        dfl is a 3d complex dfl grid with shape (nx, ny, nz)
+
+        param is a dict with:
+            gridpoints:    number of gridpoints in one transverse dimension, equal to nx and ny above
+            gridsize:      gridpoint spacing (meter)
+            refposition:   starting position (meter)
+            wavelength:    radiation wavelength (meter)
+            slicecount:    number of slices
+            slicespacing   slice spacing (meter)
+
+        These params correspond to v2 params:
+            gridpoints:   ncar
+            gridsize:     dgrid*2 / (ncar-1)
+            wavelength:   xlamds
+            slicespacing: xlamds * zsep
+
+
+    """
+
+    # Get params
+    param = {
+        key: h5[key][0]
+        for key in [
+            "gridpoints",
+            "gridsize",
+            "refposition",
+            "wavelength",
+            "slicecount",
+            "slicespacing",
+        ]
+    }
+
+    # transverse grid points in each dimension
+    nx = param["gridpoints"]
+
+    # slice list
+    slist = sorted(
+        [
+            g
+            for g in h5
+            if g.startswith("slice") and g not in ["slicecount", "slicespacing"]
+        ]
+    )
+
+    # Note from Sven:
+    #   The order of the 1D array of the wavefront is with the x coordinates as the inner loop.
+    #   So the order is (x1,y1),(x2,y1), ... (xn,y1),(x1,y2),(x2,y2),.....
+    #   This is done int he routine getLLGridpoint in the field class.
+    # Therefore the transpose is needed below
+
+    dfl = np.stack(
+        [
+            h5[g]["field-real"][:].reshape(nx, nx).T
+            + 1j * h5[g]["field-imag"][:].reshape(nx, nx).T
+            for g in slist
+        ],
+        axis=-1,
+    )
+
+    return dfl, param
+
+
+def wavefront_write_genesis4(
+    w,
+    h5: File,
+    polarization: str = None,
+    refposition: float = 0,
+) -> None:
+    """
+    Write the field data in the Genesis 4 format.
+
+    The full field data is stored as a 3D array of complex numbers `DFL` in units of `sqrt(W)`.
+
+    The relation of this and the electric field `E` in V/m is:
+    ```
+    E = DFL * sqrt(2*Z0) / Δ
+    ```
+    Where `Z0 = π * 119.9169832 V^2/W` exactly and `Δ` is the grid spacing.
+
+
+    Parameters
+    ----------
+    w: Wavefront
+
+    h5: h5py.File
+
+    polarization: str
+
+    refposition: float
+
+
+    """
+    nx, ny, nz = w.shape
+    dx, dy, dz = w.dx, w.dy, w.dz
+    wavelength = w.wavelength
+
+    # Auto-select
+    if polarization is None:
+        if w.Ey is None:
+            E = w.Ex
+        elif w.Ex is None:
+            E = w.Ey
+        else:
+            raise ValueError("Can only write one component: 'x' or 'y'")
+    else:
+        assert polarization in ("x", "y")
+        if polarization == "x":
+            E = w.Ex
+        else:
+            E = w.Ey
+
+    dfl = E * dx / np.sqrt(2 * Z0)
+
+    if nx != ny:
+        raise ValueError(f"Genesi4 requires nx = ny. This data has {nx=}, {ny=}")
+
+    if dx != dy:
+        raise ValueError(f"Genesi4 requires dx = dy. This data has {dx=}, {dy=}")
+
+    h5["gridpoints"] = np.asarray([nx])
+    h5["gridsize"] = np.asarray([dx])
+    h5["refposition"] = np.asarray([refposition])
+    h5["wavelength"] = np.asarray([wavelength])
+    h5["slicecount"] = np.asarray([nz])
+    h5["slicespacing"] = np.asarray([dz])
+
+    # Note from Sven:
+    #   The order of the 1D array of the wavefront is with the x
+    #   coordinates as the inner loop.
+    #   So the order is (x1,y1),(x2,y1), ... (xn,y1),(x1,y2),(x2,y2),.....
+    #   This is done in the routine getLLGridpoint in the field class.
+    # Therefore the transpose is needed below
+    for z in range(nz):
+        slice_index = z + 1
+        slice_group = h5.create_group(f"slice{slice_index:06}")
+        slice_group["field-real"] = dfl[:, :, z].real.T.flatten()
+        slice_group["field-imag"] = dfl[:, :, z].imag.T.flatten()

pmd_beamphysics/wavefront/wavefront.py

@@ -2,18 +2,27 @@
 from enum import Enum
 from dataclasses import dataclass, replace
 from copy import deepcopy
+from typing import Union
 
 from math import pi
 import numpy as np
 from numpy.fft import fftfreq, fftshift, ifftshift, ifftn
 
 from scipy.constants import epsilon_0, c
 
+import h5py
+import pathlib
+
 import matplotlib.pyplot as plt
 from matplotlib.colors import LogNorm
 
 from pmd_beamphysics.statistics import mean_calc, mean_variance_calc
 from pmd_beamphysics.plot import plot_2d_density_with_marginals
+from pmd_beamphysics.units import Z0
+from pmd_beamphysics.interfaces.genesis import (
+    wavefront_write_genesis4,
+    load_genesis4_fields,
+)
 
 
 def fftfreq_max(n, d=1.0):
@@ -830,3 +839,46 @@ def sigma_z(self):
         sqrt(<z^2> - <z>^2) in meters
         """
         return self._std("z")
+
+    @classmethod
+    def from_genesis4(
+        cls,
+        file: Union[pathlib.Path, str, h5py.Group],
+    ):
+        """
+        Wavefront from a Genesis4 field file
+        """
+
+        if isinstance(file, (str, pathlib.Path)):
+            with h5py.File(file, "r") as h5:
+                dfl, param = load_genesis4_fields(h5)
+
+        elif isinstance(file, h5py.Group):
+            dfl, param = load_genesis4_fields(h5)
+        else:
+            raise ValueError(f"{file=} must be a str, pathlib.Path, or h5py.Group")
+
+        dx = param["gridsize"]
+        dz = param["slicespacing"]
+        wavelength = param["wavelength"]
+
+        Ex = dfl * np.sqrt(2 * Z0) / dx
+
+        return cls(Ex=Ex, dx=dx, dy=dx, dz=dz, wavelength=wavelength)
+
+    def write_genesis4(
+        self,
+        file: Union[pathlib.Path, str, h5py.Group],
+    ):
+        """
+        Write a Genesis4-style field field
+        """
+        if isinstance(file, (str, pathlib.Path)):
+            with h5py.File(file, "w") as h5:
+                wavefront_write_genesis4(self, h5)
+            return
+
+        if isinstance(file, h5py.Group):
+            wavefront_write_genesis4(self, h5)
+
+        raise ValueError(type(file))  # type: ignore[unreachable]