grating.lua

--(C) 2015 Steven Byrnes

--This is the lua script that S4 runs. Generally you don't run this yourself, you let Python call it
--(see grating.py)

pi = math.pi
degree = pi / 180
math.randomseed(os.time())

function almost_equal(a,b,tol)
    return math.abs(a-b) <= tol * (math.abs(a) + math.abs(b))
end

function str_from_complex(a)
    return string.format('%.4f + %.4f i', a[1], a[2])
end

function polar_str_from_complex(a)
    local phase = math.atan2(a[1], a[2])
    if phase < 0 then
        phase = phase + 2*pi
    end
    return string.format('amp:%.4f, phase:%.4fdeg', math.sqrt(a[1]^2 + a[2]^2), phase/deg)
end

function mergeTables(t1, t2)
    --for hash tables. With duplicate entries, t2 will overwrite.
    local out = {}
    for k,v in pairs(t1) do out[k] = v end
    for k,v in pairs(t2) do out[k] = v end
    return out
end

-- read parameters from configuration files (hopefully just now outputted by Python)
-- all lengths in microns
f = assert(io.open("temp/grating_setup.txt", "r"))
what_to_do = f:read("*line")
assert(what_to_do == '1' or what_to_do == '2')
if what_to_do == '1' then what_to_do = 'fom' end
if what_to_do == '2' then what_to_do = 'characterize' end

if what_to_do == 'fom' then
    grating_period = tonumber(f:read("*line")) * 1e6
    lateral_period = tonumber(f:read("*line")) * 1e6
    angle_in_air = tonumber(f:read("*line"))
    n_glass = tonumber(f:read("*line"))
    nTiO2 = tonumber(f:read("*line"))
    cyl_height = tonumber(f:read("*line")) * 1e6
    --num_G is the number of modes to include. Higher = slower but more accurate
    num_G = tonumber(f:read("*line"))
elseif what_to_do == 'characterize' then
    grating_period = tonumber(f:read("*line")) * 1e6
    lateral_period = tonumber(f:read("*line")) * 1e6
    n_glass = tonumber(f:read("*line"))
    nTiO2 = tonumber(f:read("*line"))
    cyl_height = tonumber(f:read("*line")) * 1e6
    num_G = tonumber(f:read("*line"))
    --ux,uy is direction cosine
    ux_min = tonumber(f:read("*line"))
    ux_max = tonumber(f:read("*line"))
    uy_min = tonumber(f:read("*line"))
    uy_max = tonumber(f:read("*line"))
    u_steps = tonumber(f:read("*line"))
    wavelength = tonumber(f:read("*line"))
end


--setting either nTiO2 or n_glass to 0 means "use tabulated data".
-- this table is generated by refractive_index.py
if nTiO2 == 0 then
    nTiO2_data =
      {[450]=2.5,
       [500]=2.433,
       [525]=2.41,
       [550]=2.391,
       [575]=2.375,
       [580]=2.372,
       [600]=2.362,
       [625]=2.351,
       [650]=2.341}
end
if n_glass == 0 then
    nSiO2_data =
      {[450]=1.466,
       [500]=1.462,
       [525]=1.461,
       [550]=1.46,
       [575]=1.459,
       [580]=1.459,
       [600]=1.458,
       [625]=1.457,
       [650]=1.457}
end

f = assert(io.open("temp/grating_xyrra_list.txt", "r"))

xyrra_list = {}
while f:read(0) == '' do
    local line = f:read("*line")
    --print('line' .. line)
    local t={}
    local i=1
    for str in string.gmatch(line, "([^%s]+)") do
                    t[i] = tonumber(str)
                    i = i + 1
    end
    if #t > 0 then xyrra_list[#xyrra_list+1] = t end
end

function set_up()
    --initial setup of a simulation

    local S = S4.NewSimulation()
    S:SetLattice({grating_period,0}, {0, lateral_period})
    S:SetNumG(num_G)
    ------- Materials -------
    --S:AddMaterial(name, {real part of epsilon, imag part of epsilon})
    --if nTiO2 or n_glass is 0, don't worry, we'll set it later in
    --set_wavelength_angle() when we know the wavelength.
    S:AddMaterial("Air", {1,0})
    S:AddMaterial("TiO2", {nTiO2^2,0})
    S:AddMaterial("Glass", {n_glass^2,0})

    ------- Layers -------
    -- S:AddLayer(name, thickness, default material)
    S:AddLayer('Air', 0, 'Air')
    S:AddLayer('Cylinders', cyl_height, 'Air')
    S:AddLayer('Substrate', 0, 'Glass')

    for _,xyrra in ipairs(xyrra_list) do
        x,y,rx,ry,angle = unpack(xyrra)
        --print(x,y,rx,ry,angle)
        --S:SetLayerPatternEllipse(layer, material, center, angle (CCW, in degrees), halfwidths)
        S:SetLayerPatternEllipse('Cylinders', 'TiO2', {x,y}, angle, {rx,ry})
    end
    return S
end

function set_wavelength_angle(arg)
    --set the wavelength and angle of incoming light
    --S is the result of set_up()
    local S = arg.S
    local pol = arg.pol
    assert((pol == 's') or (pol == 'p'))
    local wavelength = arg.wavelength
    local incident_theta = arg.incident_theta
    local incident_phi = arg.incident_phi

    if nTiO2 == 0 then
        local nTiO2_now = nTiO2_data[math.floor(wavelength*1000+0.5)]
        assert(nTiO2_now > 0)
        S:SetMaterial("TiO2", {nTiO2_now^2,0})
    end
    local n_glass_now
    if n_glass == 0 then
        n_glass_now = nSiO2_data[math.floor(wavelength*1000+0.5)]
        assert(n_glass_now > 0)
        S:SetMaterial("Glass", {n_glass_now^2,0})
    else
        n_glass_now = n_glass
    end

    local sAmplitude, pAmplitude
    if pol == 's' then
        sAmplitude = {1,0} -- amplitude and phase
        pAmplitude = {0,0}
    else
        sAmplitude = {0,0}
        pAmplitude = {1,0}
    end

    S:SetFrequency(1/wavelength)
    -- we set it up so that the incoming wave is the (0,0) diffraction order
    S:SetExcitationPlanewave(
        {incident_theta/degree,incident_phi/degree},  -- incidence angles (spherical coordinates: phi in [0,180], theta in [0,360])
        sAmplitude,  -- s-polarization amplitude and phase (in degrees)
        pAmplitude) -- p-polarization amplitude and phase

    --futz with these settings to make it more accurate for a given num_G
    S:UsePolarizationDecomposition()
    S:UseNormalVectorBasis()
    --S:UseJonesVectorBasis()
    --S:UseSubpixelSmoothing()
    --S:SetResolution(4)
    return {S, n_glass_now}
end

function fom(arg)
    -- Calculate a figure-of-merit, basically the amount of power going into the desired diffraction order.
    -- This is used for optimizing the gratings.

    --S is output from set_up()
    local S = arg.S
    local target_diffraction_order = arg.target_diffraction_order
    local incident_theta = arg.incident_theta
    local pol = arg.pol
    local wavelength = arg.wavelength
    local inphase = arg.inphase

    local S, n_glass_now = unpack(set_wavelength_angle{pol=pol, S=S, wavelength=wavelength,
                                                         incident_theta=incident_theta, incident_phi=0})

    local outgoing_order_index = S:GetDiffractionOrder(target_diffraction_order,0)
    local forw,_ = S:GetAmplitudes('Substrate',0)
    local complex_output_amplitude
    if pol == 's' then
        complex_output_amplitude = forw[outgoing_order_index]
        --at normal incidence (i.e. center of lens), it seems that s and p wind up with opposite phases due to some weird convention
        --(I think ... I still need to check explicitly). To keep
        --everything adding in phase, we need to keep the same s-vs-p phase relation throughout the lens
        complex_output_amplitude = {-complex_output_amplitude[1], -complex_output_amplitude[2]}
    else
        complex_output_amplitude = forw[outgoing_order_index + S:GetNumG()]
   end
    -- End result: Power in desired diffraction order
    if inphase then
        return math.abs(complex_output_amplitude[2]) * complex_output_amplitude[2] / n_glass_now / math.cos(incident_theta)
    else
        return (complex_output_amplitude[1]^2 + complex_output_amplitude[2]^2) / n_glass_now / math.cos(incident_theta)
    end
    -- We are looking at the imaginary part complex_output_amplitude[2], rather than the absolute
    -- value complex_output_amplitude[1]^2 + complex_output_amplitude[2]^2, because we want the
    -- output phases of different parts of the lens to agree.
    -- We are looking at the imaginary part complex_output_amplitude[2], rather than the real
    -- part complex_output_amplitude[1], for no particular reason, they would each work equally
    -- well as an optimization target.
    -- We are looking at math.abs(complex_output_amplitude[2]) * complex_output_amplitude[2], rather
    -- than complex_output_amplitude[2]^2, because we want to have consistent phase, no possible
    -- sign-flip.
    -- We are using S:GetAmplitudes() rather than S:GetPowerFluxByOrder() because the latter
    -- is calculated in the region where the different modes are all overlapping. Also, we want
    -- just the s --> s and p --> p polarization-conserved power because it's easier to think
    -- about that than the phases of two different polarizations at once. That means it's
    -- perhaps an underestimate of collimated power.

    --consistency check: The following two should match (at least approximately
    --return (complex_output_amplitude[1]^2 + complex_output_amplitude[2]^2) / n_glass_now / math.cos(incident_theta)
    -- ..... VS .....
    --[[
    local incident_power, total_reflected_power = S:GetPowerFlux('Air', 0)
    local total_transmitted_power, temp = S:GetPowerFlux('Substrate', 0)
    assert(temp==0)
    local transmitted_normal_power, back_r, forw_i, back_i = unpack(S:GetPowerFluxByOrder('Substrate',0)[outgoing_order_index])
    assert(almost_equal(-total_reflected_power + total_transmitted_power, incident_power, 1e-3))
    assert(back_i == 0 and back_r == 0 and forw_i == 0)
    return (transmitted_normal_power / incident_power)
    --]]

    -- This is the formula I was using earlier
    --local phase = math.atan2(complex_output_amplitude[1], complex_output_amplitude[2])
    -- note: I got the atan2 arguments backwards, so that's really 90 degrees minus the phase. Doesn't matter.
    --return (transmitted_normal_power / incident_power) * math.cos(phase)
end

function print_orders(arg)
    -- This function will print the complex amplitude for the requested diffraction
    -- orders. This is used for far-field calculations.

    local S = arg.S
    local pol = arg.pol
    assert((pol == 's') or (pol == 'p'))
    local wavelength = arg.wavelength
    -- orders should be an array like {{0,0},{1,0},...} -- diffraction orders to calculate
    local orders=arg.orders
    assert(orders ~= nil)
    --ux,uy are direction cosines. They are outputted to make the results easier to parse,
    -- but otherwise are not used
    local ux = arg.ux
    local uy = arg.uy

    local forw,_ = S:GetAmplitudes('Substrate',0)
    local _,back = S:GetAmplitudes('Air',0)

    for _,order in ipairs(orders) do
        local ox,oy = unpack(order)
        local index = S:GetDiffractionOrder(ox,oy)
        local G = S:GetNumG()
        local complex_output_amplitude_fy = forw[index]
        local complex_output_amplitude_fx = forw[index + G]
        local complex_output_amplitude_ry = back[index]
        local complex_output_amplitude_rx = back[index + G]
        print(math.floor(wavelength*1000+0.5), pol, arg.ux, arg.uy, ox, oy,
              complex_output_amplitude_fy[1], complex_output_amplitude_fy[2],
              complex_output_amplitude_fx[1], complex_output_amplitude_fx[2],
              complex_output_amplitude_ry[1], complex_output_amplitude_ry[2],
              complex_output_amplitude_rx[1], complex_output_amplitude_rx[2])
    end
end

function display_fom()
    --display figure of merit for optimizing gratings
    --order (short for target_diffraction_order) should be -1 for lens, 0 for light passing right through
    --inphase is whether or not we demand that the outgoing waves have consistent complex phase

    --[[
    wavelengths_weights_orders_inphase =
    {{0.580, 1, -1, true}}
    --]]

    ---[[
    wavelengths_weights_orders_inphase =
    {{0.580, 0.5, -1, true},
    {0.450, 0.5, 0, true}}
    --]]

    --[[
    wavelengths_weights_orders_inphase =
    {{0.500, 1, -1, false},
    {0.580, 1, -1, true},
    {0.650, 1, -1, false}}
    --]]

    score_so_far = 0
    weight_so_far = 0
    for i=1,#wavelengths_weights_orders_inphase,1 do
        local wavelength, weight, target_diffraction_order, inphase = unpack(wavelengths_weights_orders_inphase[i])
        --right now I'm assuming that if we want light to pass through, we only care about normal incidence
        if target_diffraction_order ~= 0 then incident_theta = angle_in_air else incident_theta=0 end
        local S = set_up()
        local s = fom{pol='s', S=S, wavelength=wavelength,target_diffraction_order=target_diffraction_order,
                      incident_theta=incident_theta, inphase=inphase}
        local p = fom{pol='p', S=S, wavelength=wavelength, target_diffraction_order=target_diffraction_order,
                      incident_theta=incident_theta, inphase=inphase}
        --print('lam:', wavelength, 's:', s, 'p:', p)
        score_so_far = score_so_far + (s+p)/2 * weight
        weight_so_far = weight_so_far + weight
    end
    print(score_so_far / weight_so_far)

    --S:OutputLayerPatternDescription('Cylinders', 'temp/grating_img.ps')

end

if what_to_do == 'fom' then
    display_fom()
    do return end
end

function epsilon_map(S)
    local f = assert(io.open('temp/grating_eps.txt', 'w'))
    local z = cyl_height/2
    for x = -grating_period/2,grating_period/2,grating_period/100 do
        for y = -lateral_period/2,lateral_period/2,lateral_period/100 do
            local eps_r, eps_i = S:GetEpsilon({x,y,z})
            f:write(x,' ',y,' ',z,' ',eps_r,' ',eps_i,'\n')
        end
    end
end

--epsilon_map(S)

function print_fields(S, z)
    local num_x = 20
    local num_y = 20
    for ix=1,num_x,1 do
        local x = -grating_period/2 + (ix / num_x) * grating_period
        for iy=1,num_y,1 do
            local y = -lateral_period/2 + (iy / num_y) * lateral_period
            Exr, Eyr, Ezr, Hxr, Hyr, Hzr, Exi, Eyi, Ezi, Hxi, Hyi, Hzi = S:GetFields({x, y, z})
            print(x, y, z, Exr, Eyr, Ezr, Hxr, Hyr, Hzr, Exi, Eyi, Ezi, Hxi, Hyi, Hzi)
        end
    end
end


function characterize(arg)
    --characterize a grating by calculating all the complex output amplitudes as a
    --function of incoming angle
    local S = set_up()
    local wavelength = arg.wavelength
    local kvac = 2*pi / wavelength
    local grating_kx = 2*pi / grating_period
    local grating_ky = 2*pi / lateral_period

    --ux,uy are the direction cosines, i.e. the x and y components of the unit vector in the direction of the incoming light
    local ux_min = arg.ux_min
    local ux_max = arg.ux_max
    local uy_min = arg.uy_min
    local uy_max = arg.uy_max
    num_steps = arg.num_steps

    for ix=0,num_steps-1,1 do
        local ux
        if num_steps == 1 then
            ux = (ux_min + ux_max)/2
        else
            ux = ux_min + ix * (ux_max - ux_min) / (num_steps-1)
        end
        local kx = kvac * ux
        for iy=0, num_steps-1,1 do
            local uy
            if num_steps == 1 then
                uy = (uy_min + uy_max)/2
            else
                uy = uy_min + iy * (uy_max - uy_min) / (num_steps-1)
            end
            local ky = kvac * uy
            if ux^2+uy^2 < 1 then
                local theta = math.acos((1-ux^2-uy^2)^0.5)
                local phi = math.atan2(uy,ux)
                --find the propagating diffraction orders (ox,oy), where (0,0) is by definition the
                --incoming wave. Exclude the light that propagates but
                --totally-internally-reflects, unless include_tir flag is true. For calculating the
                --far-field, we really don't need that stuff, but for near-field consistency checks
                --we do.
                local k_cutoff
                if arg.include_tir == true then
                    if n_glass > 0 then
                        k_cutoff = kvac * n_glass
                    else
                        k_cutoff = kvac * nSiO2_data[math.floor(wavelength*1000+0.5)]
                    end
                else
                    k_cutoff = kvac
                end
                local orders = {}
                for ox=-5,5,1 do
                    for oy=-5,5,1 do
                        if (kx + ox*grating_kx)^2 + (ky + oy*grating_ky)^2 < k_cutoff^2 then
                            table.insert(orders, {ox,oy})
                        end
                    end
                end
                for _, pol in ipairs({'s', 'p'}) do
                    set_wavelength_angle{pol=pol, S=S, wavelength=wavelength, incident_theta=theta, incident_phi=phi}
                    --ux and uy are sent only for display purposes, not used in the calculation
                    print_orders{pol=pol, S=S, wavelength=wavelength, orders=orders, ux=ux, uy=uy}
                end
            end
        end
    end


end

-- UNCOMMENT THESE LINES FOR S4conventions.py, where I'm figuring out the relation
-- that S4 uses to translate between complex amplitudes and real-space fields
--[[
num_G = 50
pol = 's'
theta = math.asin((ux_min^2 + uy_min^2)^0.5)
phi = math.atan2(uy_min, ux_min)
characterize{num_G=num_G, ux_min=ux_min, ux_max=ux_max, uy_min=uy_min,
    uy_max=uy_max, num_steps=u_steps, wavelength=wavelength, include_tir=true}
print('Fields')
S = set_up()
set_wavelength_angle{S=S, pol=pol, wavelength=wavelength, incident_theta=theta, incident_phi=phi}

print_fields(S,-10.0)

do return end
--]]


if what_to_do == 'characterize' then
    characterize{ux_min=ux_min, ux_max=ux_max, uy_min=uy_min,
            uy_max=uy_max, num_steps=u_steps, wavelength=wavelength}
    do return end
end