[bug fix] Correction of fresnel_polarized for complex IORs (metallic)

[shinyoung-yi]

Fix the imaginary part of cosine refracted angles for complex IORs (metallic) in fresnel_polarized()

Description

1. Motivation

First of all, note that there are two conventions: $\exp\left(+i\omega t\right)$ and $\exp\left(-i\omega t\right)$ formulations to express waves. Mitsuba 3 uses the former one ( $\exp\left(+i\omega t\right)$ ), which yields $\eta = n - ik$ formulation for complex refractive indices.
I would not discuss this background since it has already been discussed in Issue #1113 and PR #1150 with @tizian .

This PR fixes the remained problem raised in Issue #1113 . The current Mitsuba 3 produces wrong sign of imaginary parts of $\cos\tilde{\theta}_t$ ($\tilde\theta_t$ denotes a complex refracted angle) for complex IORs.

Concretely, with the $\exp\left(+i\omega t\right)$ convention, my claim is that $\Im \cos \tilde\theta_t \le 0$, as best of my knowledge.
However current implementation of fresnel_polarized(Float cos_theta_i, dr::Complex<Float> eta) produces inconsistent signs for cos_theta_t depending on ranges of arguments.

While a technical reasoning for my claim $\Im \cos \tilde\theta_t \le 0$ is contained in Issue #1113, I would like to focus on concrete outputs for fresnel_polarized() here.

2. Discontinuous outputs for fresnel_polarized

import numpy as np
import matplotlib.pyplot as plt
import drjit as dr
import mitsuba as mi
mi.set_variant('cuda_ad_rgb')

print(f"{dr.__version__ = }")
print(f"{mi.__version__ = }")

print(mi.fresnel_polarized(0.5, 0.5))
print(mi.fresnel_polarized(0.5, 0.5-1e-8j))

As the above code, let's imagine Fresnel refectance values for slightly different two refractive indices.
It will be nature that two results would be almost equal, but the current Mitsuba 3 implementation produces sign flippling for a_s and a_p as follows.
BEFORE:

([[-0.3333333432674408 + 0.9428090453147888i]], [[-0.9393939971923828 + 0.34283965826034546i]], [0.0], [[0.5 + 0.0i]], [[2.0 - 0.0i]])
([[-0.3333333432674408 - 0.9428090453147888i]], [[-0.9393939971923828 - 0.34283965826034546i]], [0.0], [[0.5 - 9.99999993922529e-09i]], [[2.0 + 3.999999975690116e-08i]])

After bug fix by this PR, the results have changed as follows.
AFTER:

([[-0.3333333432674408 + 0.9428090453147888i]], [[-0.9393939971923828 + 0.34283965826034546i]], [0.0], [[0.5 + 0.0i]], [[2.0 - 0.0i]])
([[-0.3333333432674408 + 0.9428090453147888i]], [[-0.9393939971923828 + 0.34283962845802307i]], [0.0], [[0.5 - 9.99999993922529e-09i]], [[2.0 + 3.999999975690116e-08i]])

3. Reproducing plots in the documentation

I also provide a verification through reproducing plots of reflectance and phase delay in the Mitsuba 3 documentation.
Mitsuba function fresnel_polarized() does not returns complex values of cos_theta_t (it returns only the real part), so also implemented a short numpy code to show $\Im \cos\theta_t\le 0$ is correct and $\Im \cos\theta_t\ge 0$ is not relavant to the convention which Mitsuba 3 relies on.

N_theta = 91

theta_i = dr.linspace(mi.Float, 0, 90, N_theta)
cos_theta_i = dr.cos(dr.deg2rad(theta_i))
sinsq_i = 1 - cos_theta_i.numpy()**2

def fresnel_reimpl(cos_i, cos_t, eta):
    rs = (cos_i - eta*cos_t) / (cos_i + eta*cos_t)
    rp = (eta*cos_i - cos_t) / (eta*cos_i + cos_t)
    return rs, rp

def show_fresnel(eta, suptitle=None):
    cos_t = np.sqrt(1+0j - sinsq_i / (eta**2))
    cos_t_ipos = np.where(cos_t.imag < 0, -cos_t, cos_t)
    cos_t_ineg = np.where(cos_t.imag > 0, -cos_t, cos_t)
    assert (cos_t_ipos.imag >= 0).all()
    assert (cos_t_ineg.imag <= 0).all()

    configs = [r"$\Im \cos \theta_t > 0$", r"$\Im \cos \theta_t < 0$", "Mitsuba 3"]
    fig, axes = plt.subplots(3, len(configs), figsize=(9, 8))
    for i, label in enumerate(configs): # for each column
        if i == 0:
            rs, rp = fresnel_reimpl(cos_theta_i.numpy(), cos_t_ipos, eta)
        elif i == 1:
            rs, rp = fresnel_reimpl(cos_theta_i.numpy(), cos_t_ineg, eta)
        else:
            rs, rp, _, _, _ = mi.fresnel_polarized(cos_theta_i, eta)
            rs = rs.numpy()
            # print(f"{rs.shape = }")
            if dr.is_complex_v(rs):
                rs = rs[:,0] + rs[:,1] * 1j
            rp = rp.numpy()
            if dr.is_complex_v(rp):
                rp = rp[:,0] + rp[:,1] * 1j

        # ==================== First row ====================
        axes[0, i].plot(theta_i, rs.real, label=r"$\Re r_s$")
        axes[0, i].plot(theta_i, rs.imag, label=r"$\Im r_s$")
        axes[0, i].plot(theta_i, rp.real, label=r"$\Re r_p$")
        axes[0, i].plot(theta_i, rp.imag, label=r"$\Im r_p$")
        axes[0, i].set_xlabel(r"$\theta_i$")
        axes[0, i].set_ylabel("reflectance (complex amplitude)")
        axes[0, i].set_title(label)
        axes[0, i].legend()

        # ==================== Second row ====================
        Rs = np.abs(rs*rs)
        Rp = np.abs(rp*rp)

        axes[1, i].plot(theta_i, Rs, label=r"$R_s$")
        axes[1, i].plot(theta_i, Rp, label=r"$R_p$")
        axes[1, i].set_xlabel(r"$\theta_i$")
        axes[1, i].set_ylabel("reflectance (power)")
        axes[1, i].set_title(label)
        axes[1, i].legend()

        # ==================== Third row ====================
        delta = np.angle(rp/rs, deg=True)
        # print(f"{delta.shape = }")
        if delta[0] < -150 and delta[1] > 150:
            # For better visualization
            delta[0] += 360
        axes[2, i].plot(theta_i, delta)
        axes[2, i].set_xlabel(r"$\theta_i$")
        axes[2, i].set_ylabel(r"$\delta_p - \delta_s$")
        axes[2, i].set_title(label)
    if not suptitle is None:
        plt.suptitle(suptitle)
    plt.tight_layout()

eta = 1/1.5
show_fresnel(eta, "$\\eta=1/1.5$")

Figure 1.

Let's compare it with Figure 12 in the documentation. From the first and second column in Figure 1, we can see $\Im \cos\theta_t \le 0$ produces the same plots as the documentation. Current Mitsuba 3 implementaion of fresnel_polarized(Float, Float) also produces the same results as the second column, so I did not modified fresnel_polarized(Float, Float). Figure 1 does not changed by this PR.

On the other hand, for another overloaded function fresnel_polarized(Float, dr::Complex<Float>), it yields wrong values but this error is somewhat obscure since it produces correct value for 0.183-3.43j but incorrect one for 0.5-1e-8j.
I provide a diagram to show how the current implementation has a bug.

Figure 2.

(a) shows ranges of the complex IOR. By our convention, we focus on the region of orange and red colors.
(I consider we need a much more sophisticated careful discussion for signs for complex IORs and $\cos\theta_t$ with a negative real part of a IOR (Sec. 17.3.1 in [Modern Electrodynamics - Andrew Zangwill]). However, I cannot handle such cases due to my limited knowledge. Thus I would like to focus on cases with $\Re\eta \ge 0$, which covers all materials implemented in Mitsuba 3)
In (a) there are two markers for specific IORs: i) and ii). I will show plots like Figure 1 for cases of the two markers later.

(b) shows $\cos^2\theta_t$. Note that due to the range of $\sin\theta_i$ (0 to 1), specific IOR values i) and ii) have been changed into lines. For i) $\eta=0.5-10^{-8}i$, a circle marker indicates a specific value of $\cos\theta_i=0.5$, which is the case reported above (2. Discontinuous outputs for fresnel_polarized).

(c) indicates correct values of $\cos\theta_t$. Choosing the $\exp\left(+i\omega t\right)$ convention, we should choose the lower half-infinite red line on $\Re = 0$.

While directly applying sqrt (in Dr.Jit and most of other scientific computing library), the resulting value shown in (d) has wrong sign of the imaginary part when $\Re = 0$. It can be corrected by (f) (= this PR).

However, current Mitsuba 3 implementation uses (e), which produces incorrect values depending on arguments. For instance, the IOR ii) $\eta=0.183-3.43i$ (Figure 14 of the documentation) yields correct values.

eta = 0.183 - 3.43j
show_fresnel(eta, "$\\eta=0.183-3.43i$")
# plt.savefig(f"Before eta={eta:.2f}.png")
plt.savefig(f"After eta={eta:.2f}.png")

Figure 3. (Applying this PR still yields the same result)

However, the IOR ii) $\eta=0.5-10^{-8}i$ yields flipped signs in $\Im \cos\theta_t$.

eta = 0.5-1e-8j
show_fresnel(eta, "$\\eta=0.5-10^{-8}i$")

Figure 4. (Applying this commit changes the results in the third column. I manually attached the third column of the changed result as the forth column of this figure.)

Though this PR, fresnel_polarized will produces correct values of $\cos\theta_t$ with consistent negative signs for imaginary parts.

Checklist

• My code follows the style guidelines of this project
• My changes generate no new warnings
• My code also compiles for cuda_* and llvm_* variants. If you can't test this, please leave below
• I have commented my code
• I have made corresponding changes to the documentation
• I have added tests that prove my fix is effective or that my feature works
• I cleaned the commit history and removed any "Merge" commits
• I give permission that the Mitsuba 3 project may redistribute my contributions under the terms of its license