feat[FieldData]: export to ZBF

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Methods `EMEExplicitGrid.from_structures` and `EMECompositeGrid.from_structure_groups` to place EME cell boundaries at structure bounds.
 - 'ModeSimulation' now supports 'PermittivityMonitor'.
 - Classmethod `from_frequency_range` in `GaussianPulse` for generating a pulse whose amplitude in the frequency_range [fmin, fmax] is maximized, which is particularly useful for running broadband simulations.
+- `FieldData` supports exporting E fields to a Zemax Beam File (ZBF) with `FieldData.to_zbf()`.
 
 ### Changed
 - Performance enhancement for adjoint gradient calculations by optimizing field interpolation.

tests/test_data/test_monitor_data.py

 def run_emulated(simulation: td.Simulation, path=None, **kwargs) -> td.SimulationData: 
 def run_emulated(simulation: td.Simulation, path=None, **kwargs) -> td.SimulationData: 
@@ -1,5 +1,7 @@
 """Tests tidy3d/components/data/monitor_data.py"""
 
+from struct import unpack
+
 import matplotlib.pyplot as plt
 import numpy as np
 import pydantic.v1 as pydantic
@@ -23,7 +25,7 @@
 )
 from tidy3d.exceptions import DataError
 
-from ..utils import AssertLogLevel
+from ..utils import AssertLogLevel, get_spatial_coords_dict
 from .test_data_arrays import (
     AUX_FIELD_TIME_MONITOR,
     DIFFRACTION_MONITOR,
@@ -865,3 +867,225 @@ def test_no_nans():
     )
     with pytest.raises(pydantic.ValidationError):
         td.CustomMedium(eps_dataset=eps_dataset_nan)
+
+
+class TestZBF:
+    """Tests exporting field data to a zbf file"""
+
+    simulation = SIM
+    monitor = FIELD_MONITOR_2D
+
+    def read_beam_file(self, filename: str) -> tuple:
+        """Read in a Zemax Beam file
+        Adapted from https://community.zemax.com/code-exchange-10/python-class-to-open-zbf-file-4845
+
+        Parameters
+        ----------
+        beamfilename : string
+            the filename of the beam file to read
+
+        Returns
+        -------
+        version : integer
+            the file format version number
+        n : 2-tuple, (nx, ny)
+            the number of samples in the x and y directions
+        ispol : boolean
+            is the beam polarized?
+        units : integer (0 or 1 or 2 or 3)
+            the units of the beam, 0 = mm, 1 = cm, 2 = in, 3 for m
+        d : 2-tuple, (dx, dy)
+            the x and y grid spacing
+        zposition : 2-tuple, (zpositionx, zpositiony)
+            the x and y z position of the beam
+        rayleigh : 2-tuple, (rayleighx, rayleighy)
+            the x and y rayleigh ranges of the beam
+        waist : 2-tuple, (waistx, waisty)
+            the x and y waists of the beam
+        lamda : double
+            the wavelength of the beam
+        index : double
+            the index of refraction in the current medium
+        receiver_eff : double
+            the receiver efficiency. Zero if fiber coupling is not computed
+        system_eff : double
+            the system efficiency. Zero if fiber coupling is not computed.
+        grid_pos : 2-tuple of lists, (x_matrix, y_matrix)
+            lists of x and y positions of the grid defining the beam
+        efield : 4-tuple of 2D lists, (Ex_real, Ex_imag, Ey_real, Ey_imag)
+            a tuple containing two dimensional lists with the real and
+            imaginary parts of the x and y polarizations of the beam
+
+        """
+        # will need to get rid of during automation.
+        # beamfilename = self.open_zbf()
+        f = open(filename, "rb")
+        # zemax version number
+        version = unpack("i", f.read(4))[0]
+        nx = unpack("i", f.read(4))[0]
+        ny = unpack("i", f.read(4))[0]
+        ispol = unpack("i", f.read(4))[0]
+        units = unpack("i", f.read(4))[0]
+
+        if units == 0:
+            unit = "mm"
+        elif units == 1:
+            unit = "cm"
+        elif units == 2:
+            unit = "in"
+        elif units == 3:
+            unit = "m"
+
+        f.read(16)
+        dx = unpack("d", f.read(8))[0]
+        dy = unpack("d", f.read(8))[0]
+
+        # if version==1:
+        zposition_x = unpack("d", f.read(8))[0]
+        rayleigh_x = unpack("d", f.read(8))[0]
+        waist_x = unpack("d", f.read(8))[0]
+        # f.read(16)
+        zposition_y = unpack("d", f.read(8))[0]
+        rayleigh_y = unpack("d", f.read(8))[0]
+        waist_y = unpack("d", f.read(8))[0]
+        lamda = unpack("d", f.read(8))[0]
+        index = unpack("d", f.read(8))[0]
+        receiver_eff = unpack("d", f.read(8))[0]
+        system_eff = unpack("d", f.read(8))[0]
+        f.read(64)  # 8 empty doubles
+
+        rawx = [unpack("d", f.read(8))[0] for _ in range(2 * nx * ny)]
+        if ispol:
+            rawy = [unpack("d", f.read(8))[0] for _ in range(2 * nx * ny)]
+        f.close()
+
+        # Get Image array
+        # Ex Real Value for point 1;
+        # Ex Imaginary Value for point 1;
+        # Ex Real Value for point 2;
+        # Ex Imaginary Value for point 2, etc
+
+        Ex_real = np.asarray(rawx[0::2]).reshape(nx, ny)
+        Ex_imag = np.asarray(rawx[1::2]).reshape(nx, ny)
+        if ispol:
+            Ey_real = np.asarray(rawy[0::2]).reshape(nx, ny)
+            Ey_imag = np.asarray(rawy[1::2]).reshape(nx, ny)
+        else:
+            Ey_real = np.zeros((nx, ny))
+            Ey_imag = np.zeros((nx, ny))
+
+        return (
+            version,
+            nx,
+            ny,
+            ispol,
+            unit,
+            dx,
+            dy,
+            zposition_x,
+            zposition_y,
+            rayleigh_x,
+            rayleigh_y,
+            waist_x,
+            waist_y,
+            lamda,
+            index,
+            receiver_eff,
+            system_eff,
+            Ex_real,
+            Ex_imag,
+            Ey_real,
+            Ey_imag,
+        )
+
+    def make_data(self, coords: dict) -> td.components.data.data_array.DataArray:
+        """make a random DataArray out of supplied coordinates and data_type."""
+        # get shape of the data array
+        data_shape = [len(coords[k]) for k in td.ScalarFieldDataArray._dims]
+
+        # generate random data
+        np.random.seed(1)
+        data = np.random.random(data_shape) + 1j * np.random.random(data_shape)
+
+        # return as a td.ScalarFieldDataArray
+        data_array = td.ScalarFieldDataArray(data, coords=coords)
+        return data_array
+
+    @pytest.fixture(scope="class")
+    def field_data(self) -> td.FieldData:
+        """make a random FieldData from a FieldMonitor."""
+        # get the grid for the FieldData
+        grid = self.simulation.discretize_monitor(self.monitor)
+
+        # generate random complex field components
+        field_cmps = {}
+        for field_name in self.monitor.fields:
+            coords = get_spatial_coords_dict(self.simulation, self.monitor, field_name)
+            coords["f"] = list(self.monitor.freqs)
+            field_cmps[field_name] = self.make_data(coords=coords)
+
+        return td.FieldData(
+            monitor=self.monitor,
+            symmetry=(0, 0, 0),
+            symmetry_center=self.simulation.center,
+            grid_expanded=grid,
+            **field_cmps,
+        )
+
+    @pytest.mark.parametrize("background_index", [1, 2, 3])
+    @pytest.mark.parametrize("freq", list(monitor.freqs) + [None])
+    @pytest.mark.parametrize("num_samples", [2, 4, 8, 16, 32, 64, None])
+    def test_zbfread(self, tmp_path, field_data, background_index, freq, num_samples):
+        zbf_filename = tmp_path / "testzbf.zbf"
+
+        # write to zbf and then load it back in
+        ex, ey = field_data.to_zbf(
+            fname=zbf_filename,
+            background_refractive_index=background_index,
+            freq=freq,
+            num_samples=num_samples,
+        )
+        zbfdata = self.read_beam_file(zbf_filename)
+        (
+            _version,
+            nx,
+            ny,
+            _ispol,
+            _unit,
+            _dx,
+            _dy,
+            _zposition_x,
+            _zposition_y,
+            _rayleigh_x,
+            _rayleigh_y,
+            _waist_x,
+            _waist_y,
+            wl,
+            index,
+            _receiver_eff,
+            _system_eff,
+            Ex_real,
+            Ex_imag,
+            Ey_real,
+            Ey_imag,
+        ) = zbfdata
+
+        assert index == background_index
+
+        if freq is not None:
+            assert np.isclose(wl * 1e3, td.C_0 / freq)
+        else:
+            assert np.isclose(wl * 1e3, td.C_0 / self.monitor.freqs[len(self.monitor.freqs) // 2])
+
+        if num_samples is not None:
+            assert nx == num_samples
+            assert ny == num_samples
+        else:
+            assert nx == 32
+            assert ny == 64
+
+        # check that fields are close
+        assert np.allclose(ex.real.values, Ex_real)
+        assert np.allclose(ex.imag.values, Ex_imag)
+        assert np.allclose(ey.real.values, Ey_real)
+        assert np.allclose(ey.imag.values, Ey_imag)

tidy3d/components/data/monitor_data.py

@@ -2,10 +2,11 @@
 
 from __future__ import annotations
 
+import struct
 import warnings
 from abc import ABC
 from math import isclose
-from typing import Any, Callable, Dict, List, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import autograd.numpy as np
 import pydantic.v1 as pd
@@ -1224,6 +1225,127 @@ def make_adjoint_sources(
 
         return sources
 
+    def to_zbf(
+        self,
+        fname: str,
+        background_refractive_index: float,
+        freq: Optional[float] = None,
+        mode_index: Optional[int] = 0,
+        num_samples: Optional[int] = None,
+    ) -> Tuple[ScalarFieldDataArray, ScalarFieldDataArray]:
+        """For a 2D monitor, export the fields to a Zemax Beam File (zbf).
+
+        The mode area is used to approximate the beam waist, which is only valid
+        if the beam profile approximates a Gaussian beam.
+
+        Parameters
+        ----------
+        fname : str
+            Full path to the .zbf file to be written.
+        background_refractive_index : float
+            Refractive index of the medium surrounding the monitor.
+        freq : float
+            Field frequency selection. If ``None``, the central one is used.
+        mode_index : int
+            Mode index selection.
+        num_samples: int
+            Number of samples to use. Must be a power of 2
+
+        Returns
+        -------
+        Tuple[:class:`ScalarFieldDataArray`,:class:`ScalarFieldDataArray`]
+            The two E field components being exported to zbf.
+        """
+        log.warning(
+            "Exporting to ZBF is currently an experimental feature. If any issues are encountered, please contact Flexcompute support."
+        )
+
+        # Mode area calculation ensures all E components are present
+        mode_area = self.mode_area
+        dim1, dim2 = self._tangential_dims
+
+        # Using file-local coordinates x, y for the tangential components
+        e_x = self._tangential_fields["E" + dim1]
+        e_y = self._tangential_fields["E" + dim2]
+        x = e_x.coords[dim1].values
+        y = e_x.coords[dim2].values
+
+        # Use the central frequency if freq is not specified
+        if freq is None:
+            freq = e_x.coords["f"].values[e_x.coords["f"].size // 2]
+
+        mode_area = mode_area.sel(f=freq, drop=True)
+        e_x = e_x.sel(f=freq, drop=True)
+        e_y = e_y.sel(f=freq, drop=True)
+
+        if "mode_index" in e_x.coords:
+            mode_area = mode_area.isel(mode_index=mode_index, drop=True)
+            e_x = e_x.isel(mode_index=mode_index, drop=True)
+            e_y = e_y.isel(mode_index=mode_index, drop=True)
+
+        # Header info
+        version = 1
+        polarized = 1
+        units = 0  # "mm": 0, "cm": 1, "in": 2, "m": 3
+        lda = C_0 / freq * 1e-3
+
+        # Pilot (reference) beam waist: use the mode area to approximate the expected value
+        w_x = (mode_area.item() / np.pi) ** 0.5 * 1e-3
+        w_y = w_x
+
+        # Pilot beam Rayleigh distance (ignored on input)
+        r_x = 0
+        r_y = 0
+
+        # Pilot beam z position w.r.t. the waist
+        z_x = 0 * 1e-3
+        z_y = 0 * 1e-3
+
+        # Number of samples: n_x = 2 ** a, a >= 5 and a <= 13 (32-bit), or a <= 14 (64-bit)
+        if num_samples is not None:
+            # custom user defined number of samples
+            # Raise an error if the number of samples is not a power of 2
+            if (num_samples & (num_samples - 1)) != 0:
+                raise ValueError("num_samples must be a power of 2.")
+            n_x = num_samples
+            n_y = num_samples
+        else:
+            n_x = 2 ** min(13, max(5, int(np.log2(x.size) + 1)))
+            n_y = 2 ** min(13, max(5, int(np.log2(y.size) + 1)))
+
+        # Interpolating coordinates
+        x = np.linspace(x.min(), x.max(), n_x)
+        y = np.linspace(y.min(), y.max(), n_y)
+
+        # Interpolate fields
+        coords = {dim1: x, dim2: y}
+        e_x = e_x.interp(coords, assume_sorted=True)
+        e_y = e_y.interp(coords, assume_sorted=True)
+
+        # Sampling distance
+        d_x = np.mean(np.diff(x)) * 1e-3
+        d_y = np.mean(np.diff(y)) * 1e-3
+
+        # Receiver and system efficiencies
+        rec_efficiency = 0  # zero if fiber coupling is not computed
+        sys_efficiency = 0  # zero if fiber coupling is not computed
+
+        with open(fname, "wb") as fout:
+            fout.write(struct.pack("=5I", version, n_x, n_y, polarized, units))
+            fout.write(struct.pack("=4I", 0, 0, 0, 0))  # unused values
+            fout.write(struct.pack("=8d", d_x, d_y, z_x, r_x, w_x, z_y, r_y, w_y))
+            fout.write(
+                struct.pack("=4d", lda, background_refractive_index, rec_efficiency, sys_efficiency)
+            )
+            fout.write(struct.pack("=8d", 0, 0, 0, 0, 0, 0, 0, 0))  # unused values
+            for e in (e_x, e_y):
+                e_flat = e.values.flatten(order="C")
+                # Interweave real and imaginary parts
+                e_values = np.ravel(np.column_stack((e_flat.real, e_flat.imag)))
+                fout.write(struct.pack(f"={2 * n_x*n_y}d", *e_values))
+
+        return e_x, e_y
+
 
 class FieldTimeData(FieldTimeDataset, ElectromagneticFieldData):
     """