example code for Fig 7. of the paper(https://opg.optica.org/oe/fulltext.cfm?uri=oe-28-25-37773&id=444132)

test_fig7.cpp

#include <iostream>
#include <fstream>
#include <iomanip>
#include <functional>
#include <random>
#include <chrono>

#include "utils/timer.hpp"
#include "rcwa/WaveEquationCoeff.h"
#include "core/RCWAScatterMatrix.h"
#include "core/ScatterMatrixDerivative.hpp"
#include "rcwa/LayerStructure.h"
#include "rcwa/dispersion.h"
#include "opt/ConstrainedGradientDescentOptimizer.h"

using namespace std;
using namespace Eigen;

// the function to return the loss and gradient (grad) given a certain wavelength (lambda) and target amplitude (targetAmp)
// input x is the parameter to control the shape
scalar funcgrad_general_Amp(const VectorXs& x, VectorXs& grad, scalar lambda, scalar targetAmp)
{
    int N = 11;
    scalar wz = 1.4;
    scalar L = 0.66;
    scalar Lw = 0.6;
    scalar Lb = 0.5 * (L - Lw);

    WaveEquationCoeff coeff(N, N, L, L);

    VectorXs coordX(6);
    coordX << 0., Lb, 0., 0., L-Lb, L;
    coordX(2) = 0.5 * L - 0.5 * Lw * x(0);
    coordX(3) = 0.5 * L + 0.5 * Lw * x(0);

    VectorXs coordY(6);
    coordY << 0., Lb, 0., 0., L-Lb, L;
    coordY(2) = 0.5 * L - 0.5 * Lw * x(1);
    coordY(3) = 0.5 * L + 0.5 * Lw * x(1);

    cout << "coordX: " << coordX.transpose() << endl;
    cout << "coordY: " << coordY.transpose() << endl;

    int midx = (N - 1) / 2;
    int midy = (N - 1) / 2;
    int index = midy + midx * N;

    VectorXcs Ex(2*N*N);
    Ex.setZero();
    Ex(index) = 1.;

    scalex e_Si = permSi(lambda);
    scalex e_b = e_Vacuum;

    MatrixXcs eps(5, 5);
    eps << e_b, e_b, e_b, e_b, e_b,
            e_b, e_Si, e_Si, e_Si, e_b,
            e_b, e_Si, e_b, e_Si, e_b,
            e_b, e_Si, e_Si, e_Si, e_b,
            e_b, e_b, e_b, e_b, e_b;

    struct LayerStructure layer {coordX, coordY, eps};

    MatrixXcs P, Q, dPx1, dPy1, dQx1, dQy1;
    coeff.solve(lambda, layer, 2, 2, P, Q, &dPx1, &dPy1, &dQx1, &dQy1);

    MatrixXcs dPx2, dPy2, dQx2, dQy2;
    coeff.solve(lambda, layer, 3, 3, P, Q, &dPx2, &dPy2, &dQx2, &dQy2);
    
    MatrixXcs dPx = 0.5 * Lw * (dPx2 - dPx1);
    MatrixXcs dPy = 0.5 * Lw * (dPy2 - dPy1);
    MatrixXcs dQx = 0.5 * Lw * (dQx2 - dQx1);
    MatrixXcs dQy = 0.5 * Lw * (dQy2 - dQy1);

    dtmm::RCWAScatterMatrix rcwa(N, N, 0.66, 0.66);
    rcwa.compute(lambda, wz, P, Q);

    std::array<dtmm::ScatterMatrixDerivative<dtmm::RCWAScatterMatrix>, 2> deriv;
    deriv[0].compute(rcwa, P, Q, dPx, dQx);
    deriv[1].compute(rcwa, P, Q, dPy, dQy);

    scalex output = rcwa.Tdd()(index, index);
    scalar output2 = std::abs(output) * std::abs(output);

    cout << "output: " << std::abs(output) << ", " << std::arg(output) << endl;

    grad.resize(2);
    grad.setZero();

    for (int i = 0; i < 2; ++i)
    {
        grad(i) = 2 * (output2 - targetAmp*targetAmp) * 2. * std::real(deriv[i].dTdd()(index, index) * std::conj(output));
    }

    return (output2 - targetAmp*targetAmp) * (output2 - targetAmp*targetAmp);
}

// the function to return the loss and gradient (grad) given multiple wavelengths (lambdas) and the corresponding target amplitudes (targetAmps)
// input x is the parameter to control the shape
scalar funcgrad_general_Amp_multiple_frequencies(const VectorXs& x, VectorXs& grad, 
    const vector<scalar>& lambdas, const vector<scalar>& targetAmps)
{
    scalar sum = 0.;
    int Nx = x.size();
    grad.resize(Nx);
    grad.setZero();

    int n = lambdas.size();
    for (int i = 0; i < n; ++i)
    {
        scalar lambda = lambdas[i];
        scalar targetAmp = targetAmps[i];

        VectorXs temp(Nx);
        temp.setZero();

        scalar loss = funcgrad_general_Amp(x, temp, lambda, targetAmp);
        cout << "loss" << i << ": " << loss << ", ";
        sum += loss;
        grad += temp;
    }
    cout << endl;

    return sum;
}

// main function
void test(const vector<scalar>& lambdas, const vector<scalar>& targetAmps)
{
    // for a certain target, convert it into a standard function returning the 
    // gradient and loss
    auto funcgrad = bind(funcgrad_general_Amp_multiple_frequencies, 
        placeholders::_1, 
        placeholders::_2,
        lambdas, targetAmps);

    default_random_engine generator;
    auto time_point = chrono::high_resolution_clock::now();
    generator.seed(time_point.time_since_epoch().count());
    uniform_real_distribution<scalar> distribution(0., 1.);
    auto dice = bind(distribution, generator);
    
    // parameter to control the shape
    VectorXs x(2);
    for (int i = 0; i < x.size(); ++i)
    {
        x(i) = dice();
    }

    cout << "=============================================\n";
    cout << "initial x: " << x.transpose() << endl;
    cout << "=============================================\n";

    // any differentiable optimizer that you write
    ConstrainedGradientDescentOptimizer opt;
    opt.setStepSize(0.01);
    auto t1 = sploosh::now();
    opt.optimize(x, funcgrad, 40); // input: initial x, the function to evaluate gradient and loss, maximum number of iterations
    auto t2 = sploosh::now();
    cout << "time: " << sploosh::duration_milli_d(t1, t2) << endl;
    ofstream fout("../record_amp.txt", ios::app);
    for (auto a : targetAmps)
    {
        fout << a << "\t";
    }
    fout << x.transpose() << "\t" << opt.loss() << endl;

}

int main(int argc, char* argv[])
{
    for (auto A1 : {0.2, 0.4, 0.6, 0.8, 1.})
    {
        for (auto A2 : {0.2, 0.4, 0.6, 0.8, 1.})
        {
            test({1.2, 1.6}, {A1, A2});
        }
    }
    
    return 0;
}