Cannot compute gradients with respect to lattice parameters

[gevero]

Wherever i try and compute the derivative with respect to the lattice parameter I get:

Iteration 1: x = [ nan 0.24950846 0.2501126 ], f(x) = 0.9302539825439453 Traceback (most recent call last): File "/home/giovi/projects/poetry/sel_optimization/180925/torcwa_minimal_gradient_example.py", line 181, in <module> y = FoM(x, **FoM_kwargs) ^^^^^^^^^^^^^^^^^^^^ File "/home/giovi/projects/poetry/sel_optimization/180925/torcwa_minimal_gradient_example.py", line 138, in FoM sim.set_incident_angle(inc_ang=inc_ang, azi_ang=azi_ang) File "/home/giovi/projects/poetry/sel_optimization/.venv/lib/python3.11/site-packages/torcwa/rcwa.py", line 144, in set_incident_angle self._kvectors() File "/home/giovi/projects/poetry/sel_optimization/.venv/lib/python3.11/site-packages/torcwa/rcwa.py", line 1157, in _kvectors Vtmp1 = torch.linalg.inv(self.Vf+self.Vi) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ torch._C._LinAlgError: linalg.inv: The diagonal element 2 is zero, the inversion could not be completed because the input matrix is singular.

Is it even possible to compute this derivative within Torcwa, or there are inherent problems that make this forbidden? Thank you for the great code.

Best

Giovanni

P.S. Below the minimal code example to reproduce the error.

# ------ Load libraries ------
import numpy as np
import torch


# ------ Load and setup torcwa ------
import torcwa

# setup floats and devices
torch.backends.cuda.matmul.allow_tf32 = False
sim_dtype = torch.complex64
geo_dtype = torch.float32
device = torch.device("cuda")

# ------ keyword args ------
FoM_kwargs = {
    "dtype": geo_dtype,
    "device": device,
    "edge_sharpness": 500,
    "broadening_parameter": 1e-10,
    "nx": 1500,
    "ny": 1500,
    "n_sup": 1.0,
    "n_sub": 1.45,
    "n_ms": 3.0,
    "wavelength": 1.5,
    "theta": 0.0,
    "phi": 0.0,
    "order_N": 21,
}


def FoM(
    x,
    dtype=geo_dtype,
    device=device,
    edge_sharpness=500,
    broadening_parameter=1e-45,
    nx=1500,
    ny=1500,
    n_sup=1.0,
    n_sub=1.45,
    n_ms=3.0,
    wavelength=1.5,
    theta=1e-10,
    phi=0.0,
    order_N=21,
):
    """
    This function calculates the transmission efficiency of light through a
    metasurface structure consisting of a hexagonal lattice of circular holes
    in a layered system. The simulation uses the RCWA (Rigorous Coupled Wave Analysis)
    method to solve the electromagnetic problem.

        Vector containing [a, h, r] where:
        - a: lattice constant (period)
        - h: thickness of metasurface layer
        - r: radius of circular holes
    dtype : torch.dtype, optional
        Data type for computations (default: geo_dtype)
    device : torch.device, optional
        Device for computations (default: device)
    edge_sharpness : int, optional
        Sharpness parameter for edge discretization (default: 500)
    broadening_parameter : float, optional
        Broadening parameter for numerical stability (default: 1e-45)
    nx : int, optional
        Number of discretization points along x-axis (default: 1500)
    ny : int, optional
        Number of discretization points along y-axis (default: 1500)
    n_sup : float, optional
        Refractive index of superstrate material (default: 1.0)
    n_sub : float, optional
        Refractive index of substrate material (default: 1.45)
    n_ms : float, optional
        Refractive index of metasurface material (default: 3.0)
        Wavelength of incident light in micrometers (default: 1.5)
    theta : float, optional
        Polar incidence angle in degrees (default: 1e-10)
    phi : float, optional
        Azimuthal incidence angle in degrees (default: 0.0)
    order_N : int, optional
        Number of Fourier harmonics to include (default: 21)

    T : torch.Tensor
        Transmission efficiency (sum of squared magnitudes of transmission
        coefficients for all diffraction orders)
    """
    a, h, r = x

    # initialize geometry
    L = [a, a]
    torcwa.rcwa_geo.dtype = dtype
    torcwa.rcwa_geo.device = device
    torcwa.rcwa_geo.Lx = L[0]
    torcwa.rcwa_geo.Ly = L[1]
    torcwa.rcwa_geo.nx = nx
    torcwa.rcwa_geo.ny = ny
    torcwa.rcwa_geo.grid()
    torcwa.rcwa_geo.edge_sharpness = edge_sharpness
    torcwa.Eig.broadening_parameter = broadening_parameter

    # ------ layers ------
    # hex as face centered rectangular
    circle = torcwa.rcwa_geo.circle(R=r, Cx=L[0] / 2.0, Cy=L[1] / 2.0)
    layer0_geometry = circle
    layer0_thickness = h

    # Generate an array with all diffraction orders
    order = [order_N, order_N]
    order_range = np.arange(-order_N, order_N + 1)
    all_orders = np.stack(np.meshgrid(order_range, order_range), -1).reshape(-1, 2)

    # define ref index
    n_sup = n_sup
    n_sub = n_sub
    n_ms = n_ms

    # inc light
    inc_ang = theta * (np.pi / 180)  # radian
    azi_ang = phi * (np.pi / 180)  # radian

    # setup simulation,  input and output layers and incident angle
    sim = torcwa.rcwa(
        freq=1 / (wavelength), order=order, L=L, dtype=sim_dtype, device=device
    )
    sim.add_input_layer(eps=n_sup**2)
    sim.add_output_layer(eps=n_sub**2)
    sim.set_incident_angle(inc_ang=inc_ang, azi_ang=azi_ang)

    # build ms layer
    layer0_eps = layer0_geometry + (1.0 - layer0_geometry) * (n_ms**2)
    sim.add_layer(thickness=layer0_thickness, eps=layer0_eps)

    # solve problem
    sim.solve_global_smatrix()

    tss = sim.S_parameters(
        orders=all_orders,
        direction="f",
        port="t",
        polarization="ss",
    )

    T = torch.sum(torch.abs(tss) ** 2)

    return T


# 2. Initialize the parameter (the variable we want to optimize)
x = torch.tensor([2.0, 0.25, 0.25], requires_grad=True)

# 3. Define the optimizer
# We'll use Stochastic Gradient Descent (SGD)
learning_rate = 3.3e-4
optimizer = torch.optim.SGD(
    [x], lr=learning_rate
)  # Pass a list of parameters to optimize

# 4. Iterative Optimization (Gradient Descent Loop)
num_iterations = 1000

for i in range(num_iterations):
    # Zero the gradients (important for each iteration)
    optimizer.zero_grad()

    # Compute the function val  ue
    y = FoM(x, **FoM_kwargs)

    # Compute the gradient
    y.backward()

    # Update the parameter using the optimizer
    optimizer.step()

    print(f"Iteration {i+1}: x = {x.detach().numpy()}, f(x) = {y.item()}")