Automatic differentiation of the solvers

[aadityacs]

Hi @mfschubert ,

Thank you so much for these wrappers. I'm wondering how to go about extending your implementation to also encompass computing the gradients w.r.t eigenvalues.

This could be particularly useful for dispersion problems such as for instance in sigmund 2003 we compute an objective that depends on the eigenvalues to increase bandgap in the structure.

[mfschubert]

Hi @aadityacs , as an example I suggest the implementation in fmmax here: https://github.com/mfschubert/fmmax/blob/main/src/fmmax/eig.py

As noted in the code, the details of the vjp are handled differently throughout the literature, and so I am a bit hesitant to use a specific implementation in this package.

https://github.com/mfschubert/fmmax/blob/2d4397409e7e68f62d3e051da69605a8e6ffc9d0/src/fmmax/eig.py#L111-L141

    # The expression for gradient of matrix eigenvectors with respect to matrix values
    # contains a difference between eigenvalues in the denominator. This causes
    # numerical problems in situations where eigenvalues are degenerate or nearly
    # degenerate. This problem is addressed in a few different ways in existing rcwa
    # codes. Defining `delta = eigval_i - eigval_j`,
    #
    # - torcwa uses Lorentzian broadening of the form
    #      1 / delta -> delta.conj / (abs(delta)**2 + eps)
    #   with an eps value of 1e-10.
    #   (https://github.com/kch3782/torcwa/blob/main/torcwa/torch_eig.py#L29)
    #
    # - tf_rcwa uses an expression similar to torcwa, but with apparently different
    #   conjugation. Their value of eps is 1e-6. (We tested their implementation and
    #   found it fails some of our gradient validation tests.)
    #   (https://github.com/scolburn54/rcwa_tf/blob/master/src/tensor_utils.py#L89)
    #
    # - grcwa uses `1 / delta -> 1 / (delta + eps)` with an eps value of 1e-10.
    #   (https://github.com/weiliangjinca/grcwa/blob/master/grcwa/primitives.py#L44)
    #
    # Also, from the wider literature:
    #
    # - Liao et al. (https://arxiv.org/pdf/1903.09650 sec III.A.1) uses Lorentzian
    #   broadening with an eps value of 1e-12.
    #
    # One issue with Lorentzian broadening and a fixed eps is the independence from
    # matrix scale. While a given eps may be appropriate for a matrix, it would be
    # inappropriate for the same matrix scaled by e.g. 1e-6. (The gradient should
    # be identical, except scaled by 1 / 1e-6.)
    #
    # Therefore, we use Lorentzian broadening similar to torcwa, but with an eps
    # value that is computed in a way that considers the eigenvalue range.

[aadityacs]

Thank you!