error in incoherent example code

[RadPlanets]

When running example_inc_tmm.py, I get an error AssertionError: It's not clear which beam is incoming vs outgoing. Weird index maybe? related to the forward angle check.

The other example file (example_tmm) runs without any errors.

Full trace below:

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
File \\example_inc_tmm.py:46
     43 T[:, -1] = np.inf
     45 T = torch.from_numpy(T)
---> 46 O_fast = inc_tmm_fast(pol, M, T, mask, theta, wl, device='cpu')
     48 fig, ax = plt.subplots(2,1)
     49 cbar = ax[0].imshow(O_fast['R'][0].numpy(), aspect='auto')

File \\tmm_fast\vectorized_incoherent_tmm.py:152, in inc_vec_tmm_disp_lstack(pol, N, D, mask, theta, lambda_vacuum, device, timer)
    150 d = D[:, m_]
    151 d[:, 0] = d[:, -1] = np.inf
--> 152 forward = coh_tmm(
    153     pol, N_, d, snell_theta[0, :, m_[0], 0], lambda_vacuum, device
    154 )
    155 # the substack must be evaluated in both directions since we can have an incoming wave from the output side
    156 # (a reflection from an incoherent layer) and Reflectivit/Transmissivity can be different depending on the direction
    157 backward = coh_tmm(
    158     pol,
    159     N_.flip([1]),
   (...)
    163     device,
    164 )

File \\tmm_fast\vectorized_tmm_dispersive_multistack.py:134, in coh_vec_tmm_disp_mstack(pol, N, T, Theta, lambda_vacuum, device, timer)
    130    N = torch.tile(N, (num_wavelengths, 1)).T
    132 # SnellThetas is a tensor, for each stack and layer, the angle that the light travels
    133 # through the layer. Computed with Snell's law. Note that the "angles" may be complex!
--> 134 SnellThetas = SnellLaw_vectorized(N, Theta)
    137 theta = 2 * np.pi * torch.einsum('skij,sij->skij', torch.cos(SnellThetas), N)  # [theta,d, lambda]
    138 kz_list = torch.einsum('sijk,k->skij', theta, 1 / lambda_vacuum)  # [lambda, theta, d]

File \\tmm_fast\vectorized_tmm_dispersive_multistack.py:247, in SnellLaw_vectorized(n, th)
    238 # The first and last entry need to be the forward angle (the intermediate
    239 # layers don't matter, see https://arxiv.org/abs/1603.02720 Section 5)
    241 angles[:, :, 0] = torch.where(
    242     is_not_forward_angle(n[:, 0], angles[:, :, 0]).bool(),
    243     pi - angles[:, :, 0],
    244     angles[:, :, 0],
    245 )
    246 angles[:, :, -1] = torch.where(
--> 247     is_not_forward_angle(n[:, -1], angles[:, :, -1]).bool(),
    248     pi - angles[:, :, -1],
    249     angles[:, :, -1],
    250 )
    252 return angles

File \\tmm_fast\vectorized_tmm_dispersive_multistack.py:297, in is_not_forward_angle(n, theta)
    294 assert (ncostheta.real > -100 * EPSILON)[answer].all(), error_string
    295 assert ((n * torch.cos(torch.conj(theta))).real > -100 * EPSILON)[answer].all(), error_string
--> 297 assert (ncostheta.imag < 100 * EPSILON)[~answer].all(), error_string
    298 assert (ncostheta.real < 100 * EPSILON)[~answer].all(), error_string
    299 assert ((n * torch.cos(torch.conj(theta))).real < 100 * EPSILON)[~answer].all(), error_string

AssertionError: It's not clear which beam is incoming vs outgoing. Weird index maybe?
n: tensor([[2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j],
        [2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j,
         2.2000+0.0500j, 2.2000+0.0500j, 2.2000+0.0500j]],
       dtype=torch.complex128)   angle: tensor([[[0.0000+0.0000j, 0.0000+0.0000j, 0.0000+0.0000j,  ...,
          0.0000+0.0000j, 0.0000+0.0000j, 0.0000+0.0000j],
         [0.0160-0.0004j, 0.0160-0.0004j, 0.0160-0.0004j,  ...,
          0.0160-0.0004j, 0.0160-0.0004j, 0.0160-0.0004j],
         [0.0321-0.0007j, 0.0321-0.0007j, 0.0321-0.0007j,  ...,
          0.0321-0.0007j, 0.0321-0.0007j, 0.0321-0.0007j],
         ...,
         [0.4696-0.0115j, 0.4696-0.0115j, 0.4696-0.0115j,  ...,
          0.4696-0.0115j, 0.4696-0.0115j, 0.4696-0.0115j],
         [0.4709-0.0116j, 0.4709-0.0116j, 0.4709-0.0116j,  ...,
          0.4709-0.0116j, 0.4709-0.0116j, 0.4709-0.0116j],
         [0.4715-0.0116j, 0.4715-0.0116j, 0.4715-0.0116j,  ...,
          0.4715-0.0116j, 0.4715-0.0116j, 0.4715-0.0116j]],

        [[0.0000+0.0000j, 0.0000+0.0000j, 0.0000+0.0000j,  ...,
          0.0000+0.0000j, 0.0000+0.0000j, 0.0000+0.0000j],
         [0.0160-0.0004j, 0.0160-0.0004j, 0.0160-0.0004j,  ...,
          0.0160-0.0004j, 0.0160-0.0004j, 0.0160-0.0004j],
         [0.0321-0.0007j, 0.0321-0.0007j, 0.0321-0.0007j,  ...,
          0.0321-0.0007j, 0.0321-0.0007j, 0.0321-0.0007j],
         ...,
         [0.4696-0.0115j, 0.4696-0.0115j, 0.4696-0.0115j,  ...,
          0.4696-0.0115j, 0.4696-0.0115j, 0.4696-0.0115j],
         [0.4709-0.0116j, 0.4709-0.0116j, 0.4709-0.0116j,  ...,
          0.4709-0.0116j, 0.4709-0.0116j, 0.4709-0.0116j],
         [0.4715-0.0116j, 0.4715-0.0116j, 0.4715-0.0116j,  ...,
          0.4715-0.0116j, 0.4715-0.0116j, 0.4715-0.0116j]]],
       dtype=torch.complex128)

[qizheng-1998]

I'm getting the same error message. It seems like the imaginary number caused the problem? I have the older version that does not get this error message, but I'm getting the error message "Opacity warning. The imaginary part of the refractive index is clamped to 35i for numerical stability." Any fix on this?

[Nerrror]

Hey,
can you both try to run again with the develop branch checked out and see if it works? I think I worked on this bug a while ago but wasn't ready to merge this into main because of some other issue.

[qizheng-1998]

Hi, @Nerrror ,
I don't get the error AssertionError: It's not clear which beam is incoming vs outgoing. Weird index maybe? in the develop branch, but I'm still getting the warnging "Opacity warning. The imaginary part of the refractive index is clamped to 35i for numerical stability.". Could you please evaluate the purpose of this warning?

[RadPlanets]

The develop branch works for me as well. Thanks!

I get a similar warning (pasted below), but as I understand it, this does not affect the computations. Clamping the imaginary refractive index to a high value prevents divide by zero errors and other instabilities.

\\tmm_fast\vectorized_tmm_dispersive_multistack.py:153: UserWarning: Opacity warning. The imaginary part of the refractive index is clamped to 35i for numerical stability.
You might encounter problems with gradient computation...
  warn('Opacity warning. The imaginary part of the refractive index is clamped to 35i for numerical stability.\n'+
\\tmm_fast\vectorized_incoherent_tmm.py:256: UserWarning: Casting complex values to real discards the imaginary part (Triggered internally at C:\cb\pytorch_1000000000000\work\aten\src\ATen\native\Copy.cpp:300.)
  P_[..., 0, 0] = 1/P

[Nerrror]

Hi, @Nerrror , I don't get the error AssertionError: It's not clear which beam is incoming vs outgoing. Weird index maybe? in the develop branch, but I'm still getting the warnging "Opacity warning. The imaginary part of the refractive index is clamped to 35i for numerical stability.". Could you please evaluate the purpose of this warning?

@qizheng-1998 basically, exactly what @RadPlanets describes. The intensity in this layer becomes so low that it creates instabilities when you take this up to e to propagate the wave in the transfer matrix.
Technically, what is clamped (ie. capped at the value of 35i) is the accrued phase k*T through an absorbing medium which is called delta in the code. We clamp delta because it's not a good idea to clamp the imaginary part of the material directly. What we want to avoid is to get numerically instable low intensity. You can achieve that by limiting the amount of power that is getting absorbed when a ray passes through the medium. This is achieved by clamping delta.
However, this clampings effect on any results should be negligible and is just intendend to make you aware of this intervention