Macroscopic Maxwell Solver

Introduction

This Python 3 package enables solving the macroscopic Maxwell equations in complex dielectric materials.

The material properties are defined on a rectangular grid (1D, 2D, or 3D) for which each voxel defines an isotropic or anisotropic permittivity. Optionally, a heterogeneous (anisotropic) permeability as well as bi-anisotropic coupling factors may be specified (e.g. for chiral media). The source, such as an incident laser field, is specified as an oscillating current-density distribution.

The method iteratively corrects an estimated solution for the electric field (default: all zero). Its memory requirements are on the order of the storage requirements for the material properties and the electric field within the calculation volume. Full details can be found in the open-access manuscript "Calculating coherent light-wave propagation in large heterogeneous media". When the machine learning library PyTorch is detected, the wave equations can also be solved on the cloud or a local GPU, as described in the paper doi:10.34133/icomputing.0098.

Examples of usage can be found in the examples/ sub-folder. The Complete MacroMax Documentation can be found at https://macromax.readthedocs.io. All source code is available on GitHub under the MIT License: https://opensource.org/license/MIT

Installation

Prerequisites

This library requires Python 3 with the numpy and scipy packages for the main calculations. These modules will be automatically installed. The modules multiprocessing, torch, pyfftw, and mkl-fft (Intel(R) CPU specific) can significantly speed up the calculations.

The examples require matplotlib for displaying the results. The code has been tested on Python 3.11 and 3.12, though it is expected to work on versions 3.10 and above.

Installing

Installing the macromax package and its mandatory dependencies is as straightforward as running the following command in a terminal:

pip install -U macromax

The -U option is to request an upgrade if macromax is already installed. While this is sufficient to get started, optional packages are useful to display the results and to speed-up the calculations.

Maximizing execution speed on a CPU

The most straightforward way to increase the computation speed is by installing PyTorch for your system.

pip install -U macromax torch

This will drastically reduce the computation time, both with CPU and GPU. Do make sure to install the correct version of PyTorch that makes full use of your hardware (see below). MacroMax will fall back to using NumPy on your CPU if it cannot find a working PyTorch installation.

Alternatively, the calculation time on the CPU can be improved by ensuring that you have the fastest libraries installed for your system. In particular the FFTW library can easily halve the calculation time on a CPU:

pip install -U macromax pyFFTW

On some systems the pyFFTW Python package requires the separate installation of the FFTW library; however, it is easy to install it using Anaconda with the commands: conda install fftw, or on Debian-based systems with sudo apt-get install fftw.

Another option is the mkl-fft package that is available for Intel(R) CPUs. However, this may require compilation or relying on the Anaconda or Intel Python distributions:

conda install -c intel intelpython

Leveraging a machine learning framework for GPU and cloud-based calculations

The calculation time can be reduced by several orders of magnitude using the PyTorch machine learning library. This can be as straightforward as using the appropriate runtime on Google Colab. The MacroMax library can here be installed by prepending the command with an exclamation mark as follows:

!pip install --upgrade macromax

For more details, check out the Google Colab deployment example.

Local GPUs can also be used provided that PyTorch has a compatible implementation. At the time of writing, these includes NVidia's CUDA-enabled GPU as well AMD's ROCm-enabled GPUs (on Linux). Prior to installing the PyTorch module following the PyTorch Guide, install the appropriate CUDA or ROCm drivers for your GPU. Note that for PyTorch to work correctly, Nvidia drivers need to be up-to-date and match the installed CUDA version. At the time of writing, for CUDA version 13.2, PyTorch can be installed as follows using pip:

pip install -U torch --extra-index-url https://download.pytorch.org/whl/cu132

Specifics for your CUDA version and operating system are listed on PyTorch Guide.

When PyTorch and a compatible GPU are detected, these will be used by default. If not, FFTW and mkl-fft will be used if available. Otherwise, NumPy and SciPy will be used as a fallback. The default backend can be set at the start of your code, or by creating a text file named backend_config.json in the current working directory with contents as:

[
  {"type": "torch", "device": "cuda"},
  {"type": "torch", "device": "cpu"},
  {"type": "numpy"}
]

to choose PyTorch when a compatible GPU is available, and NumPy otherwise. Although this machine learning library can be used without a hardware accelerator, we found that NumPy (with FFTW) can be faster when no GPU is available. This backend selection rule is used by default.

Additional packages

The package comes with a submodule containing example code that should run as-is on most desktop installations of Python. Some systems may require the installation of the ubiquitous matplotlib graphics library:

pip install -U matplotlib

The output logs can be colored by installing the coloredlogs packaged:

pip install -U coloredlogs

Building and distributing the library may require further packages as indicated below.

Usage

The basic calculation procedure consists of the following steps:

1. define the material

2. define the coherent light source

3. call solution = macromax.solve(...)

4. display the solution

The macromax package must be imported to be able to use the solve function. The package also contains several utility functions that may help in defining the property and source distributions.

Examples can be found in the examples package in the examples/ folder. Ensure that the entire examples/ folder is downloaded, including the __init__.py file with general definitions. Run the examples from the parent folder using e.g. python -m examples.air_glass_air_1D.

The complete functionality is described in the Library Reference Documentation at [https://macromax.readthedocs.io].

Loading the Python 3 package

The macromax package can be imported using:

import macromax

Optional: If the package is installed without a package manager, it may not be on Python's search path. If necessary, add the library to Python's search path, e.g. using:

import sys
import os
sys.path.append(os.path.dirname(os.getcwd()))

Reminder: this library requires Python 3, numpy, and scipy. Optionally, pyfftw can be used to speed up the calculations. The examples also require matplotlib.

Specifying the material

Defining the sampling grid

The material properties are sampled on a plaid uniform rectangular grid of voxels. The sample points are defined by one or more linearly increasing coordinate ranges, one range per dimensions. The coordinates must be specified in meters, e.g.:

import numpy as np
x_range = 50e-9 * np.arange(1000)

Ranges for multiple dimensions can be passed to solve(...) as a tuple of ranges: ranges = (x_range, y_range), though it is often easier to use a Grid object. The latter can be used as follows:

data_shape = (200, 400)
sample_pitch = 50e-9  # or (50e-9, 50e-9)
grid = macromax.Grid(data_shape, sample_pitch)

This defines a uniformly spaced plaid grid, centered around the origin, unless specified otherwise.

Defining the material property distributions

The material properties are defined by ndarrays of 2+N dimensions, where N can be up to 3 for three-dimensional samples. In each sample point, or voxel, a complex 3x3 matrix defines the anisotropy at that point in the sample volume. The first two dimensions of the ndarray are used to store the 3x3 matrix, the following dimensions are the spatial indices x, y, and z. Optionally, four complex ndarrays can be specified: epsilon, mu, xi, and zeta. These ndarrays represent the permittivity, permeability, and the two coupling factors, respectively.

When the first two dimensions of a property are found to be both a singleton, i.e. 1x1, that property is assumed to be isotropic. Similarly, singleton spatial dimensions are interpreted as homogeneity in that property. The default permeability mu is 1, and the coupling constants are zero by default.

Boundary conditions

Most calculations make some assumption about the electromagnetic field outside the calculation grid. Often the grid is assumed to sitting in an infinitely-large free space. However, the algorithm implicitly uses periodic boundary conditions. While periodic boundaries are useful for some situations, infinite free space would require that all waves travel outward from the calculation grid. This can be approximated by adding a boundary layer that gradually absorbs the emanating waves whilst minimizing any potential back reflections.

from macromax.bound import LinearBound

bound = LinearBound(grid, thickness=5e-6)

The thickness can also be specified as an N-by-2-matrix with the boundaries on all 2N sides of the N-dimensional calculation grid in the order ([top, bottom], [left, right], [front, back]). If you would like one or more of the boundaries to be periodic, just set its thickness to 0. The Matrix class in macromax.matrix uses periodic boundaries in the transverse axes to model plane waves with an infinite extent. The LinearBound adds a linearly-increasing extinction coefficient to the default refractive index of 1. At the edges of the calculation grid, the extinction coefficient increases to 0.25 by default. Custom boundary conditions can be implemented by inheriting from one of the Bound classes in the macromax.bound module.

Alternative boundary conditions can be implemented by surrounding the calculation area with absorbing (or reflective) layers. Back reflections can be suppressed by e.g. linearly increasing the imaginary part of the permittivity with depth into a boundary with a thickness of a few wavelengths.

Defining the source

The coherent source is defined by as a spatially-variant free current density. Although the current density may be non-zero in all of space, it is more common to define a source at one of the edges of the volume, to model e.g. an incident laser beam; or even as a single voxel, to simulate a dipole emitter. The source density can be specified as a complex number, indicating the phase and amplitude of the current at each point. If an extended source is defined, care should be taken so that the source currents constructively interfere in the desired direction. I.e. the current density at neighboring voxels should have a phase difference matching the k-vector in the background medium. Optionally, instead of a current density, the internally-used source distribution may be specified directly. It is related to the current density as follows: S = i omega mu_0 J with units of rad s⁻¹ H m⁻¹ A m⁻² = rad V m⁻³, where omega is the angular frequency, and mu_0 is the vacuum permeability, mu₀.

The source distribution is stored as a complex ndarray with 1+N dimensions. The first dimension contains the current 3D direction and amplitude for each voxel. The complex argument indicates the relative phase at each voxel.

Calculating the electromagnetic light field

Once the macromax module is imported, the solution satisfying the macroscopic Maxwell's equations is calculated by calling:

solution = macromax.solve(...)

The function arguments to macromax.solve(...) can be the following:

• grid|x_range: A Grid object, a vector (1D), or tuple of vectors (2D, or 3D) indicating the spatial coordinates of the sample points. Each vector must be a uniformly increasing array of coordinates, sufficiently dense to avoid aliasing artefacts.

• vacuum_wavelength|wave_number|anguler_frequency: The wavelength in vacuum of the coherent illumination in units of meters.

• current_density or source_distribution: An ndarray of complex values indicating the source value and direction at each sample point. The source values define the free current density in the sample. The first dimension contains the vector index, the following dimensions contain the spatial dimensions. If the source distribution is not specified, it is calculated as :math:-i c k0 mu_0 J, where i is the imaginary constant, c, k0, and mu_0, the light-speed, wavenumber, and permeability in vacuum. Finally, J is the free current density (excluding the movement of bound charges in a dielectric), specified as the input argument current_density. These input arguments should be numpy.ndarrays with a shape as specified by the grid input argument, or have one extra dimension on the left to indicate the polarization. If polarization is not specified the solution to the scalar wave equation is calculated. However, when polarization is specified the vectorial problem is solved. The returned macromax.Solution object has the property vectorial to indicate whether polarization is accounted for or not.

• refractive_index: A complex numpy.ndarray of a shape as indicated by the grid argument. Each value indicates the refractive at the corresponding spatial grid point. Real values indicate a loss-less material. A positive imaginary part indicates the absorption coefficient, :math:\kappa. This input argument is not required if the permittivity, epsilon is specified.

• epsilon: (optional, default: :math:n^2) A complex numpy.ndarray of a shape as indicated by the grid argument for isotropic media, or a shape with two extra dimensions on the left to indicate anisotropy/birefringence. The array values indicate the relative permittivity at all sample points in space. The optional two first (left-most) dimensions may contain a 3x3 matrix at each spatial location to indicate the anisotropy/birefringence. By default the 3x3 identity matrix is assumed, scaled by the scalar value of the array without the first two dimensions. Real values indicate loss-less permittivity. This input argument is unit-less, it is relative to the vacuum permittivity.

Optionally one can also specify magnetic and coupling factors:

• mu: A complex ndarray that defines the 3x3 permeability matrix at all sample points. The first two dimensions contain the matrix indices, the following dimensions contain the spatial dimensions.

• xi and zeta: Complex ndarray that define the 3x3 coupling matrices at all sample points. This may be useful to model chiral materials. The first two dimensions contain the matrix indices, the following dimensions contain the spatial dimensions.

• initial_field: If an initial estimate for the electric-field is known, this can potentially reduce the number of iterations required. This could be useful when calculating the field for a structure that is very similar to a structure for which the field is known. Do bear in mind that small phase shifts can cause significant changes in the field values that can negate the speed advantage.

It is often useful to also specify a callback function that tracks progress. This can be done by defining the callback-argument as a function that takes an intermediate solution as argument. This user-defined callback function can display the intermediate solution and check if the convergence is adequate. The callback function should return True if more iterations are required, and False otherwise. E.g.:

callback=lambda s: s.residue > 0.01 and s.iteration < 1000

will iterate until the residue is at most 1% or until the number of iterations reaches 1,000.

The solution object (of the Solution class) fully defines the state of the iteration and the current solution as described below.

The macromax.solve(...) function returns a solution object. This object contains the electric field vector distribution as well as diagnostic information such as the number of iterations used and the magnitude of the correction applied in the last iteration. It can also calculate the displacement, magnetizing, and magnetic fields on demand. These fields can be queried as follows:

• solution.E: Returns the electric field distribution.
• solution.H: Returns the magnetizing field distribution.
• solution.D: Returns the electric displacement field distribution.
• solution.B: Returns the magnetic flux density distribution.
• solution.S: The Poynting vector distribution in the sample.

The field distributions are returned as complex numpy ndarrays in which the first dimensions is the polarization or direction index. The following dimensions are the spatial dimensions of the problem, e.g. x, y, and z, for three-dimensional problems.

The solution object also keeps track of the iteration itself. It has the following diagnostic properties:

• solution.iteration: The number of iterations performed.
• solution.residue: The relative magnitude of the correction during the previous iteration. and it can be used as a Python iterator.

Further information can be found in the examples and the signatures of each function and class.

Complete Example

The following code loads the library, defines the material and light source, calculates the result, and displays it. To keep this example as simple as possible, the calculation is limited to one dimension. Higher dimensional calculations simply require the definition of the material and light source in 2D or 3D.

The first section of the code loads the macromax library module as well as its utils submodule. More

import numpy as np
import matplotlib.pyplot as plt
import macromax
from macromax.bound import LinearBound

# %matplotlib notebook  # Uncomment this line in an iPython Jupyter notebook

#
# Define the material properties
#
wavelength = 500e-9  # [ m ] In SI units as everything else here
source_polarization = np.array([0, 1, 0])  # y-polarized
boundary_thickness = 5e-6  # Set this to 0 to get periodic boundary conditions.
# The thickness can also be specified as an Nx2 array: ([top, bottom], [left, right], [front, back])

# Set the sampling grid
nb_samples = 1024
sample_pitch = wavelength / 10  # [ m ]  # Oversample just for display
x_range = sample_pitch * np.arange(nb_samples) - boundary_thickness  # in meters. You can also use macromax.Grid for convenience.

# Define the medium as a spatially-variant permittivity. Glass has a refractive index of about 1.51:
refractive_index = 1 + (1.51 - 1) * ((x_range >= 20e-6) & (x_range < 30e-6))

#
# Define the illumination source
#
# point source at x = 0. Add the polarization vector on the left to do vectorial calculations.
current_density = source_polarization[:, np.newaxis] * (np.abs(x_range) < sample_pitch/4)

#
# Solve Maxwell's equations
#
# (the actual work is done in this line)
solution = macromax.solve(grid=x_range, vacuum_wavelength=wavelength, current_density=current_density,
                          refractive_index=refractive_index, 
                          bound=LinearBound(x_range, thickness=boundary_thickness)
                          )

#
# Display the results
#
fig, ax = plt.subplots(2, 1, frameon=False, figsize=(8, 6))

x_range = solution.grid[0]  # coordinates
E = solution.E[1, :]  # Electric field in y
H = solution.H[2, :]  # Magnetizing field in z
S = solution.S[0, :]  # Poynting vector in x
f = solution.f[0, :]  # Optical force in x
# Display the field for the polarization dimension
field_to_display = E
max_val_to_display = np.amax(np.abs(field_to_display))
poynting_normalization = np.amax(np.abs(S)) / max_val_to_display
ax[0].plot(x_range / 1e-6, np.abs(field_to_display) ** 2 / max_val_to_display, color=[0, 0, 0])
ax[0].plot(x_range / 1e-6, S.real / poynting_normalization, color=[1, 0, 1])
ax[0].plot(x_range / 1e-6, field_to_display.real, color=[0, 0.7, 0])
ax[0].plot(x_range / 1e-6, field_to_display.imag, color=[1, 0, 0])
figure_title = f'Iteration {solution.iteration}, '
ax[0].set_title(figure_title)
ax[0].set_xlabel(r'x  [$\mu$m]')
ax[0].set_ylabel('I, E  [a.u., V/m]')
ax[0].set_xlim(x_range[[0, -1]] * 1e6)

ax[1].plot(x_range[-1] * 2e6, 0, color=[0, 0, 0], label='I')
ax[1].plot(x_range[-1] * 2e6, 0, color=[1, 0, 1], label='$S_{real}$')
ax[1].plot(x_range[-1] * 2e6, 0, color=[0, 0.7, 0], label='$E_{real}$')
ax[1].plot(x_range[-1] * 2e6, 0, color=[1, 0, 0], label='$E_{imag}$')
ax[1].plot(x_range / 1e-6, refractive_index.real, color=[0, 0, 1], label=r'$\epsilon_{real}$')
ax[1].plot(x_range / 1e-6, refractive_index.imag, color=[0, 0.5, 0.5], label=r'$\epsilon_{imag}$')
ax[1].set_xlabel(r'x  [$\mu$m]')
ax[1].set_ylabel(r'$\epsilon$')
ax[1].set_xlim(x_range[[0, -1]] * 1e6)
ax[1].legend(loc='upper right')

plt.show(block=True)  # Not needed for iPython Jupyter notebook

Optimization of time and memory efficiency

Electromagnetic calculations tend to test the limits of the hardware. Two factors should be considered when optimizing the calculation: computation and memory. Naturally, the number of operations and the duration of each operation should be minimized. However, the latter is often dominated by memory accesses and copying of arrays. The memory usage therefore does not only affect the size of the problems that can be solved, it also tends to have an important impact on the total calculation time.

A straightforward method to reduce memory usage is to switch from 128-bit precision complex numbers to 64-bit. By default, the precision of the source_density is used, which is typically np.complex128 or its real equivalent. The Solution's default dtype can be overridden by specifying it as solve(... dtype=np.complex64). Halving the storage requirements can eliminate additional copies between the main memory and CPU cache. In extreme cases it can also avoid swapping. Lower precision math also executes faster on many architectures.

While oversampling to less than 1/10th of the wavelength may aid visualization, it is often sufficient to sample at a quarter of the wavelength. The sample solution represents a sinc-interpolated continuous function. The final result can be visualized with arbitrary resolution using interpolation.

The number of operations can be kept to a minimum by:

• using non-magnetic and non-chiral materials,
• using isotropic materials,
• limiting the largest difference in permittivity (including the absorbing boundary), and
• using a scalar approximation whenever possible.

Optimization of the implementation is another route to consider. Potentially areas of improvement are:

• Profiling of memory usage and elimination of redundant temporary copies
• In-place fast-Fourier transforms. When available, the FFTW library is used; however, the drop-in fft and ifft replacements are used at the moment.
• Moving the calculations to a GPU or a cloud-computing environment. Since the copying-overheads may quickly become a bottleneck, it is important to consider the memory requirements for the problem you want to solve.

Modifying and extending the library

The CONTRIBUTING page provided an overview of the source code and how to build it and its documentation. We welcome any contributions, suggestions, bug reports, or better yet, pull requests. If you like to get involved, please get in touch.