FDFD/inverse_design_FDFD-epsstart-eps1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using StatsBase: mean\n",
    "using NLopt\n",
    "using ChainRules\n",
    "using PyPlot\n",
    "using LinearAlgebra\n",
    "using SparseArrays\n",
    "using Zygote\n",
    "using DelimitedFiles\n",
    "using BenchmarkTools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function Maxwell_2d(Lx, Ly, ϵ, ω, dpml, resolution;\n",
    "    Rpml=1e-20)\n",
    "\n",
    "    nx = n_func(Lx, dpml, resolution) #nb point in x\n",
    "    ny = n_func(Ly, dpml, resolution) #nb points in y\n",
    "    npml = round(Int, dpml*resolution)\n",
    "    δ = 1/resolution\n",
    "    σ0 = -log(Rpml) / (4dpml^3/3)\n",
    "\n",
    "    # coordinates centered in (0,0)\n",
    "    x = (1-npml:nx-npml) * δ\n",
    "    y = (1-npml:ny-npml) * δ\n",
    "    \n",
    "    #define the laplacian operator in x direction\n",
    "    o = ones(nx)/δ\n",
    "    \n",
    "    σ = Float64[ξ>Lx ? σ0 * (ξ-Lx)^2 : ξ<0 ? σ0 * ξ^2 : 0.0 for ξ in x]\n",
    "    Σ = spdiagm(0 => 1.0 ./(1 .+ (im/ω)*σ))\n",
    "    \n",
    "    Imat, J, V = SparseArrays.spdiagm_internal(-1 => -o[1:end-1], 0 => o);\n",
    "    D = sparse(Imat, J, V, nx, nx)\n",
    "    D[1,end] = -1/δ #periodic boundary condition in x direction\n",
    "    ∇2x = Σ * transpose(D) * Σ * D\n",
    "\n",
    "    #define the laplacian operator in y direction\n",
    "    o = ones(ny) / δ\n",
    "    y′=((-npml:ny-npml) .+ 0.5) * δ\n",
    "\n",
    "    σ = Float64[ξ>Ly ? σ0 * (ξ-Ly)^2 : ξ<0 ? σ0 * ξ^2 : 0.0 for ξ in y]\n",
    "    Σ = spdiagm(0 => 1.0 ./(1 .+ (im/ω)*σ))\n",
    "    σ′ = Float64[ξ>Ly ? σ0 * (ξ-Ly)^2 : ξ<0 ? σ0 * ξ^2 : 0.0 for ξ in y′]\n",
    "    Σ′ = spdiagm(0 => 1.0 ./(1 .+ (im/ω)*σ′))\n",
    "    \n",
    "    Imat, J, V = SparseArrays.spdiagm_internal(-1 => -o, 0 => o);\n",
    "    D = sparse(Imat, J, V, ny+1, ny)\n",
    "    ∇2y = Σ * transpose(D) * Σ′ * D\n",
    "\n",
    "    #get 2d laplacian using kronecker product\n",
    "    Ix = sparse(1.0I, nx, nx)\n",
    "    Iy = sparse(1.0I, ny, ny)\n",
    "    ∇2d = (kron(Ix, ∇2y) + kron(∇2x, Iy))\n",
    "\n",
    "    if isa(ϵ, Function)\n",
    "        geometry = ComplexF64[ϵ(ξ, ζ) for ζ in y, ξ in x]\n",
    "    else\n",
    "        geometry = ϵ \n",
    "    end\n",
    "\n",
    "    return (∇2d - spdiagm(0 => reshape(ω^2 * geometry, length(x)*length(y))),\n",
    "    nx, ny, x, y)\n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Problem definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "changevarx(x) = x-2\n",
    "changevary(y) = y-2\n",
    "\n",
    "const dpml = 1\n",
    "const Lx = 4\n",
    "const Ly = 5\n",
    "const ω = 2pi\n",
    "const resolution = 40\n",
    "\n",
    "const dsource = 1.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_func(Lx, dpml, resolution) = round(Int, (Lx + 2*dpml) * resolution) #nb point in x\n",
    "\n",
    "function x_func(Ly, dpml, resolution)\n",
    "    ny = n_func(Ly, dpml, resolution) #nb points in y\n",
    "    npml = round(Int, dpml*resolution)\n",
    "    δ = 1/resolution\n",
    "    return (1-npml:ny-npml) * δ\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function create_geom_nopml(ϵ, ny, nx; defaulteps = 0)\n",
    "    geom = ones((ny, nx))\n",
    "    if defaulteps!=0\n",
    "        geom[design_domain, :] .= defaulteps\n",
    "    end\n",
    "    geom[design_domain, non_pml_x] = ϵ\n",
    "    return geom\n",
    "end\n",
    "\n",
    "function create_geom(ϵ, ny, nx)\n",
    "    geom = ones((ny, nx))\n",
    "    geom[design_domain, :] = ϵ\n",
    "    return geom\n",
    "end\n",
    "\n",
    "function create_target(Lx, Ly, dpml, resolution, target_domain)\n",
    "    \n",
    "    nx = n_func(Lx, dpml, resolution)\n",
    "    ny = n_func(Ly, dpml, resolution)\n",
    "    target = zeros((ny, nx))\n",
    "    \n",
    "    x = x_func(Lx, dpml, resolution)\n",
    "    y = x_func(Ly, dpml, resolution)\n",
    "    mask_x = -0.5 .<= changevarx.(x) .<= 0.5\n",
    "    mask_y = 1 .<= changevary.(y) .<= 2\n",
    "    \n",
    "    target[mask_y, mask_x] .= 1\n",
    "    \n",
    "    return target[target_domain, non_pml_x]\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = x_func(Ly, dpml, resolution)\n",
    "x = x_func(Lx, dpml, resolution)\n",
    "\n",
    "const design_domain = -1 .<= changevary.(y) .<= 0;\n",
    "const target_domain = 0 .<= changevary.(y) .<= 3;\n",
    "const non_pml_x = -2 .<= changevary.(x) .<= 2;\n",
    "\n",
    "const target = create_target(Lx, Ly, dpml, resolution, target_domain); # square [-0.5, 1], [0.5, 2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imshow(target)\n",
    "colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "const ub = 12\n",
    "const lb = 1;\n",
    "\n",
    "# const defaulteps = (lb+ub)/2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function solve_field(ϵ)\n",
    "    \n",
    "    nx = n_func(Lx, dpml, resolution) #nb point in x\n",
    "    ny = n_func(Ly, dpml, resolution) #nb points in y\n",
    "#     geom = create_geom_nopml(ϵ, ny, nx, defaulteps=defaulteps)\n",
    "    geom = create_geom(ϵ, ny, nx)\n",
    "    \n",
    "    A, nx, ny, x, y=Maxwell_2d(Lx, Ly, geom, ω, dpml, resolution);\n",
    "\n",
    "    J = zeros(ComplexF64, (ny, nx))\n",
    "    J[Integer(dsource * resolution), :]  = 1im  * ω * ones(nx) * resolution #plane source\n",
    "    Ez = reshape(A \\ J[:], (ny, nx));\n",
    "    return A, Ez, ny, nx, x, y\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimization problem setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## define gradient using adjoint method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "`MSE = MSE_func(ϵ)`\n",
    "this function returns the MSE between the intensity (defined as the abs2 of the electric field) and the target field\n",
    "\n",
    "ϵ is the permitivity of the design domain\n",
    "\"\"\"\n",
    "function MSE_func(ϵ)\n",
    "    A, Ez, ny, nx, x, y = solve_field(ϵ)\n",
    "    return mean((abs2.(Ez[target_domain, non_pml_x]) .- target).^2)\n",
    "end\n",
    "\n",
    "function ChainRules.rrule(::typeof(MSE_func), ϵ)\n",
    "    A, Ez, ny, nx, x, y = solve_field(ϵ)\n",
    "    \n",
    "    target_points = zeros(ComplexF64, (ny, nx))\n",
    "    target_points[target_domain,non_pml_x] = 2/length(target) .*(abs2.(Ez[target_domain, non_pml_x]) .- target).*conj.(Ez[target_domain, non_pml_x])\n",
    "#     gradient = 2 * real(ω^2 .*reshape(conj.(A' \\ conj.(target_points[:])) .* Ez[:], (ny, nx))[design_domain, non_pml_x]);\n",
    "    gradient = 2 * real(ω^2 .*reshape(conj.(A' \\ conj.(target_points[:])) .* Ez[:], (ny, nx))[design_domain, :]);\n",
    "\n",
    "        \n",
    "    function pullback(Δ)\n",
    "        return (NO_FIELDS, gradient .* Δ)\n",
    "    end\n",
    "    \n",
    "    return mean((abs2.(Ez[target_domain, non_pml_x]) .- target).^2), pullback    \n",
    "end\n",
    "\n",
    "Zygote.refresh()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cur_design = [lb for y_ in y[design_domain], x_ in x]; #initial guess with pml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MSE_func(cur_design)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MSE_func'(cur_design)[:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## inverse design using nlopt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "it = 0\n",
    "f_evals = []\n",
    "\n",
    "function myfunc(x::Vector, grad::Vector)\n",
    "    \n",
    "    global it, f_evals\n",
    "    cur_design = reshape(x, (:, n_func(Lx, dpml, resolution)))\n",
    "    \n",
    "    if length(grad) > 0\n",
    "        grad[:] = MSE_func'(cur_design)[:]\n",
    "    end\n",
    "    \n",
    "    f = MSE_func(cur_design)\n",
    "    println(\"f_$it = $f\")\n",
    "    it+=1\n",
    "    push!(f_evals, f)\n",
    "    return f\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "cur_design = [lb for y_ in y[design_domain], x_ in x]; #initial guess with pml\n",
    "\n",
    "opt = Opt(:LD_MMA, length(cur_design))\n",
    "\n",
    "opt.lower_bounds = lb .* ones(length(cur_design))\n",
    "opt.upper_bounds = ub .* ones(length(cur_design))\n",
    "opt.xtol_rel = 1e-8\n",
    "\n",
    "opt.min_objective = myfunc\n",
    "\n",
    "@time (minf,minx,ret) = optimize(opt, cur_design[:])\n",
    "numevals = opt.numevals # the number of function evaluations\n",
    "println(\"got $minf after $numevals iterations (returned $ret)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A, Ez, ny, nx, x, y = solve_field(minx);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loglog(f_evals)\n",
    "ylabel(\"MSE\")\n",
    "xlabel(\"# optimization step\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimal_design = reshape(minx, (:, n_func(Lx, dpml, resolution)))\n",
    "\n",
    "writedlm(\"optimal_design_eps1_.csv\", optimal_design, ',')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generating figures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimal_design = readdlm(\"optimal_design_eps1.csv\", ',')\n",
    "\n",
    "A, Ez, ny, nx, x, y = solve_field(optimal_design);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func = abs2\n",
    "cmapname = \"viridis\"\n",
    "\n",
    "subplot(1,2,1)\n",
    "title(\"$func(Ez)\")\n",
    "contourf(changevarx.(x),changevary.(y),func.(Ez), cmap=cmapname, levels=500)\n",
    "xlabel(\"x\")\n",
    "ylabel(\"y\")\n",
    "colorbar()\n",
    "\n",
    "subplot(1,2,2)\n",
    "title(\"geometry\")\n",
    "epsilonvals = create_geom(optimal_design, ny, nx)\n",
    "sourcecolor = maximum(epsilonvals)\n",
    "contourf(changevarx.(x),changevary.(y), epsilonvals, \n",
    "    cmap=\"viridis\", levels=100)\n",
    "xlabel(\"x\")\n",
    "ylabel(\"y\")\n",
    "colorbar()\n",
    "\n",
    "tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@btime MSE_func'(optimal_design)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@btime MSE_func(optimal_design)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "writed(\"intensity_FDFD.txt\", abs2.(Ez), ' ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "writedlm(\"epsilonvals_FDFD.txt\", epsilonvals, ' ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "writedlm(\"x_axis_FDFD.txt\", changevarx.(x), ' ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "writedlm(\"y_axis_FDFD.txt\", changevary.(y), ' ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.5.0",
   "language": "julia",
   "name": "julia-1.5"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.5.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}