Detector Base Class

def draw_plot(
    self,
    state: dict[str, np.ndarray],
    progress: Progress | None = None,
) -> dict[str, Figure | str]:
    """Generates plots or videos from recorded detector data.

    Creates visualizations based on dimensionality of recorded data and detector
    configuration. Supports 1D line plots, 2D heatmaps, and video generation
    for time-varying data.

    Args:
        state (dict[str, np.ndarray]): Dictionary containing recorded field data arrays.
        progress (Progress | None, optional): Optional progress bar for video generation.

    Returns:
        dict[str, Figure | str]: Dictionary mapping plot names to either
            matplotlib Figure objects or paths to generated video files.
    """
    squeezed_arrs = {}
    squeezed_ndim = None
    for k, v in state.items():
        v_squeezed = v.squeeze()
        if self.inverse and self.if_inverse_plot_backwards and self.num_time_steps_recorded > 1:
            squeezed_arrs[k] = v_squeezed[::-1, ...]
        else:
            squeezed_arrs[k] = v_squeezed
        if squeezed_ndim is None:
            squeezed_ndim = len(v_squeezed.shape)
        else:
            if len(v_squeezed.shape) != squeezed_ndim:
                raise Exception("Cannot plot multiple arrays with different ndim")
    if squeezed_ndim is None:
        raise Exception(f"empty state: {state}")

    figs = {}
    if squeezed_ndim == 1 and self.num_time_steps_recorded > 1:
        # do line plot
        time_steps = np.where(np.asarray(self._is_on_at_time_step_arr))[0]
        time_steps = time_steps * self._config.time_step_duration
        for k, v in squeezed_arrs.items():
            fig = plot_line_over_time(arr=v, time_steps=time_steps.tolist(), metric_name=f"{self.name}: {k}")
            figs[k] = fig
    elif squeezed_ndim == 1 and self.num_time_steps_recorded == 1:
        SCALE = 10
        xlabel = None
        if self.grid_shape[0] > 1 and self.grid_shape[1] <= 1 and self.grid_shape[2] <= 1:
            xlabel = "X axis (μm)"
        elif self.grid_shape[0] <= 1 and self.grid_shape[1] > 1 and self.grid_shape[2] <= 1:
            xlabel = "Y axis (μm)"
        elif self.grid_shape[0] <= 1 and self.grid_shape[1] <= 1 and self.grid_shape[2] > 1:
            xlabel = "Z axis (μm)"
        assert xlabel is not None, "This should never happen"
        for k, v in squeezed_arrs.items():
            spatial_axis = np.arange(len(v)) / SCALE
            fig = plot_line_over_time(
                arr=v, time_steps=spatial_axis, metric_name=f"{self.name}: {k}", xlabel=xlabel
            )
            figs[k] = fig
    elif squeezed_ndim == 2 and self.num_time_steps_recorded > 1:
        # multiple time steps, 1d spatial data - visualize as 2D waterfall plot
        time_steps = np.where(np.asarray(self._is_on_at_time_step_arr))[0]
        time_steps = time_steps * self._config.time_step_duration

        # Determine spatial axis based on which dimension has size > 1
        SCALE = 10  # μm per grid point

        for k, v in squeezed_arrs.items():
            # Determine which dimension is spatial (not time)
            spatial_dim = 1 if v.shape[1] > 1 else 0
            if spatial_dim == 0:
                # Transpose if needed so time is always first dimension
                v = v.T

            # Create spatial axis in μm
            spatial_points = np.arange(v.shape[1]) / SCALE

            fig = plot_waterfall_over_time(
                arr=v,
                time_steps=time_steps,
                spatial_steps=spatial_points,
                metric_name=f"{self.name}: {k}",
                spatial_unit="μm",
            )
            figs[k] = fig
    elif squeezed_ndim == 2 and self.num_time_steps_recorded == 1:
        # single time step, 2d-plot  # TODO:
        if all([x in squeezed_arrs.keys() for x in ["XY Plane", "XZ Plane", "YZ Plane"]]):
            fig = plot_2d_from_slices(
                xy_slice=squeezed_arrs["XY Plane"],
                xz_slice=squeezed_arrs["XZ Plane"],
                yz_slice=squeezed_arrs["YZ Plane"],
                resolutions=(
                    self._config.resolution,
                    self._config.resolution,
                    self._config.resolution,
                ),
                plot_dpi=self.plot_dpi,
                plot_interpolation=self.plot_interpolation,
            )
            figs["sliced_plot"] = fig
        else:
            raise Exception(f"Cannot plot {squeezed_arrs.keys()}")
    elif squeezed_ndim == 3 and self.num_time_steps_recorded > 1:
        # multiple time steps, 2d-plots
        if all([x in squeezed_arrs.keys() for x in ["XY Plane", "XZ Plane", "YZ Plane"]]):
            path = generate_video_from_slices(
                plt_fn=plot_from_slices,
                xy_slice=squeezed_arrs["XY Plane"],
                xz_slice=squeezed_arrs["XZ Plane"],
                yz_slice=squeezed_arrs["YZ Plane"],
                progress=progress,
                num_worker=self.num_video_workers,
                resolutions=(
                    self._config.resolution,
                    self._config.resolution,
                    self._config.resolution,
                ),
                plot_dpi=self.plot_dpi,
                plot_interpolation=self.plot_interpolation,
            )
            figs["sliced_video"] = path
        else:
            raise Exception(
                f"Cannot plot {squeezed_arrs.keys()}. "
                f"Consider setting plot=False for Object {self.name} ({self.__class__=})"
            )
    elif squeezed_ndim == 3 and self.num_time_steps_recorded == 1:
        # single step, 3d-plot. # TODO: do as mean over planes
        for k, v in squeezed_arrs.items():
            xy_slice = squeezed_arrs[k].mean(axis=0)
            xz_slice = squeezed_arrs[k].mean(axis=1)
            yz_slice = squeezed_arrs[k].mean(axis=2)
            fig = plot_2d_from_slices(
                xy_slice=xy_slice,
                xz_slice=xz_slice,
                yz_slice=yz_slice,
                resolutions=(
                    self._config.resolution,
                    self._config.resolution,
                    self._config.resolution,
                ),
                plot_dpi=self.plot_dpi,
                plot_interpolation=self.plot_interpolation,
            )
            figs[k] = fig
    elif squeezed_ndim == 4 and self.num_time_steps_recorded > 1:
        # video with 3d-volume in each time step. plot as slices
        for k, v in squeezed_arrs.items():
            xy_slice = squeezed_arrs[k].mean(axis=1)
            xz_slice = squeezed_arrs[k].mean(axis=2)
            yz_slice = squeezed_arrs[k].mean(axis=3)
            path = generate_video_from_slices(
                plt_fn=plot_from_slices,
                xy_slice=xy_slice,
                xz_slice=xz_slice,
                yz_slice=yz_slice,
                progress=progress,
                num_worker=self.num_video_workers,
                resolutions=(
                    self._config.resolution,
                    self._config.resolution,
                    self._config.resolution,
                ),
                plot_dpi=self.plot_dpi,
                plot_interpolation=self.plot_interpolation,
            )
            figs[k] = path
    else:
        raise Exception("Cannot plot detector with more than three dimensions")
    return figs

@abstractmethod
def update(
    self,
    time_step: jax.Array,
    E: jax.Array,
    H: jax.Array,
    state: DetectorState,
    inv_permittivity: jax.Array,
    inv_permeability: jax.Array | float,
) -> DetectorState:
    """Updates detector state with current field values.

    Args:
        time_step (jax.Array): Current simulation time step.
        E (jax.Array): Electric field array.
        H (jax.Array): Magnetic field array.
        state (DetectorState): Current detector state.
        inv_permittivity (jax.Array): Inverse permittivity array.
        inv_permeability (jax.Array | float): Inverse permeability array.

    Returns:
        DetectorState: Updated detector state.
    """
    del (
        time_step,
        E,
        H,
        state,
        inv_permittivity,
        inv_permeability,
    )
    raise NotImplementedError()