Table of Contents

• torchrdit.viz
  • plot2d
  • plot_layer
  • display_fitted_permittivity

torchrdit.viz

Visualization module for TorchRDIT electromagnetic simulations.

This module provides functions for visualizing simulation results, material properties, and field distributions in TorchRDIT. It includes tools for:

1. Plotting 2D field distributions (electric fields, magnetic fields)
2. Visualizing material property distributions in layers
3. Displaying dispersive material properties
4. Analyzing fitted dispersive material models

The visualization functions in this module work with matplotlib to create publication-quality figures for analyzing simulation results and understanding the behavior of electromagnetic fields in complex structures.

Key functions:

• plot2d: General-purpose function for plotting 2D data with customizable options
• plot_layer: Visualize the permittivity distribution in a specific layer
• display_fitted_permittivity: Plot fitted dispersive material properties in a simulation
• display_dispersive_profile: Visualize raw dispersive material data

All functions return matplotlib objects that can be further customized as needed.

Keywords: visualization, plotting, matplotlib, field distribution, permittivity, material properties, dispersive materials, electromagnetic simulation

plot2d

def plot2d(data: Any,
           layout: Any,
           func="abs",
           outline=None,
           fig_ax=None,
           cbar=True,
           cmap="magma",
           outline_alpha=0.5,
           title="",
           labels=("x", "y"))

Create a 2D visualization of electromagnetic field data or material distributions.

This function generates a 2D plot of field data or material properties (like permittivity) using matplotlib's pcolormesh. It handles both real and complex data, applying the specified function (abs, real, imag) for visualization. The function can also overlay an outline on the plot to highlight structural boundaries.

Arguments:

• data - 2D data to plot (torch.Tensor or numpy.ndarray). Can be complex-valued. For field data, this is typically the electric or magnetic field. For material data, this is typically the permittivity distribution.
• layout - Tuple of mesh grid coordinates (X, Y) defining the spatial coordinates for each point in the data. Should be obtained from sim.get_layout() or similar function.
• func - Function to apply to the data for visualization if the data is complex. Options are 'abs' (magnitude), 'real' (real part), 'imag' (imaginary part), 'log10' (log scale), or other numpy/torch functions. Default is 'abs'.
• outline - Optional binary mask (0s and 1s) to overlay as an outline on the plot. This is useful for highlighting boundaries between different materials. Default is None (no outline).
• fig_ax - Matplotlib figure and axis to use for plotting. If None, a new figure and axis will be created. Default is None.
• cbar - Whether to include a colorbar with the plot. Default is True.
• cmap - Matplotlib colormap to use for the data visualization. Default is 'magma'.
• outline_alpha - Transparency of the outline overlay (0-1). Default is 0.5.
• title - Title for the plot. Default is '' (no title).
• labels - Tuple of (x-axis label, y-axis label). Default is ('x', 'y').

Returns:

• matplotlib.axes.Axes - The matplotlib axes object containing the plot, allowing for further customization if needed.

Examples:

from torchrdit.solver import create_solver
from torchrdit.viz import plot2d
import matplotlib.pyplot as plt

# Create a solver
solver = create_solver()
# Plot the real part of the electric field distribution
x_grid, y_grid = solver.get_layout()
results = solver.solve(solver.add_source(theta=0, phi=0, pte=1.0, ptm=0.0))
e_field = results['tx']  # Transmission field x-component

fig, ax = plt.subplots(figsize=(8, 6))
plot2d(data=e_field,
        layout=(x_grid, y_grid),
        func='real',
        cmap='RdBu_r',
        title='Transmission Field (X) Distribution',
        labels=('x (μm)', 'y (μm)'))
plt.show()

Keywords: visualization, plotting, field distribution, color map, electromagnetic, pcolormesh, complex data, contour, outline, colorbar, 2D plot

plot_layer

def plot_layer(sim,
               layer_index: int = 0,
               frequency_index: int = 0,
               fig_ax=None,
               func="real",
               cmap="BuGn",
               labels=("x", "y"),
               title="")

Visualize the permittivity distribution of a specific layer in the simulation.

This function creates a 2D plot of the permittivity (ε) distribution within a specified layer of the simulation structure. For dispersive materials, the permittivity at a specific frequency is plotted.

Arguments:

• sim - A simulation object (typically a FourierBaseSolver instance) that contains the layer structure to be visualized.
• layer_index - Index of the layer to plot (0-based indexing). Default is 0 (first layer).
• frequency_index - Index of the frequency to use for dispersive materials. Only relevant for frequency-dependent (dispersive) materials. Default is 0 (first frequency).
• fig_ax - Matplotlib figure and axis to use for plotting. If None, a new figure and axis will be created. Default is None.
• func - Function to apply to the permittivity data before plotting. Options include 'real', 'imag', 'abs', 'log10', etc. Default is 'real' (plot real part of permittivity).
• cmap - Matplotlib colormap to use for the plot. Default is 'BuGn'.
• labels - Tuple of (x-axis label, y-axis label). Default is ('x', 'y').
• title - Title for the plot. If empty, a default title with layer information will be used. Default is ''.

Returns:

• matplotlib.axes.Axes - The matplotlib axes object containing the plot, allowing for further customization if needed.

Examples:

# Create a simulation with multiple layers
solver = create_solver(algorithm=Algorithm.RDIT)
solver.add_layer(material_name='silicon', thickness=torch.tensor(0.2))
solver.add_layer(material_name='sio2', thickness=torch.tensor(0.3))

# Plot the permittivity distribution of the second layer
fig, ax = plt.subplots(figsize=(8, 6))
plot_layer(solver, layer_index=1, fig_ax=ax,
            title='SiO2 Layer', labels=('x (μm)', 'y (μm)'))
plt.show()

Keywords: permittivity, layer visualization, material distribution, RCWA, R-DIT, electromagnetic simulation, layer structure, dispersive materials

display_fitted_permittivity

def display_fitted_permittivity(sim, fig_ax=None)

Visualize fitted permittivity for dispersive materials in the simulation.

This function plots the real and imaginary parts of the permittivity for dispersive materials in the simulation, showing both the original data points and the fitted curves used in the simulation. This is useful for verifying the quality of the polynomial fit to the material dispersion data.

The plot includes:

• Original data points for the real part (ε')
• Fitted curve for the real part at simulation wavelengths
• Original data points for the imaginary part (ε")
• Fitted curve for the imaginary part at simulation wavelengths

Arguments:

• sim - A simulation object (typically a FourierBaseSolver instance) that contains the dispersive materials to be visualized.
• fig_ax - Matplotlib axis to use for plotting. If None, a new figure and axis will be created. Default is None.

Returns:

• tuple - A tuple containing (primary_axis, secondary_axis) where:
  • primary_axis: The main matplotlib axis (for real part)
  • secondary_axis: The secondary y-axis (for imaginary part)

Notes:

This function automatically detects all dispersive materials in the simulation and plots their fitted permittivity data. If no dispersive materials are found, it prints a message indicating this.

Examples:

from torchrdit.solver import create_solver
from torchrdit.constants import Algorithm
from torchrdit.utils import create_material
from torchrdit.viz import display_fitted_permittivity
import torch
import matplotlib.pyplot as plt

# Create a simulation with a dispersive material (e.g., gold)
solver = create_solver(
    algorithm=Algorithm.RDIT,
    lam0=[1.55],
    lengthunit='um',
    rdim=[512, 512],
    kdim=[5, 5],
    device='cpu'
)
gold = create_material(name='gold', dielectric_dispersion=True,
                        user_dielectric_file='gold_data.txt')
solver.add_materials([gold])
solver.add_layer(material_name='gold', thickness=torch.tensor(0.1))

# Visualize the fitted permittivity
fig, ax = plt.subplots(2, 1, figsize=(10, 6))
ax1, ax2 = display_fitted_permittivity(solver, fig_ax=ax)
ax1.set_xlabel('Wavelength (μm)')
ax1.set_ylabel('Real Part (ε')')
ax2.set_ylabel('Imaginary Part (ε")')
plt.title('Gold Permittivity vs Wavelength')
plt.show()

Keywords: dispersive material, permittivity, wavelength dependence, polynomial fitting, material characterization, complex permittivity, optical properties, Drude-Lorentz model, material dispersion