test_constraint_function.py

"""                                                                             
Checks that the constraint function with constraint L applied to a circle
with diameter d has a return value based on the convention that a violation
(L > d) is >=0 and a non-violation is ~0. 
"""

from typing import Tuple

import matplotlib

import matplotlib.pyplot as plt
import meep.adjoint as mpa
import numpy as np
import nlopt

# properties of design region
res = 100
sx, sy = 5.0, 5.0
Nx, Ny = int(sx * res), int(sy * res)

# general experiment properties
beta = 32  # projection bias
eta_i = 0.5  # blueprint design field thresholding point in [0, 1]
eta_e = 0.75  # erosion design field thresholding point in [0, 1]
eta_d = 1 - eta_e  # dilation design field thresholding point in [0, 1]

# Projected field diameter
projected_diam = 1.0


def constraint_solid_void(
    min_length_scale: float, d: np.ndarray
) -> Tuple[float, float]:
    """Returns the value of the constraint function for the solid and void
       regions of the design weights.

    Args:
        min_length_scale: the minimum length scale (line width and spacing) of
            the design region to check for a violation.
        d: weights of the design region as a flattened 1d array.
    """

    filter_radius = mpa.get_conic_radius_from_eta_e(min_length_scale, eta_e)

    # filter radius should be equivalent to minimum length scale
    assert filter_radius == min_length_scale

    threshold_f = lambda a: mpa.tanh_projection(a, beta, eta_i)

    filter_f = lambda a: mpa.conic_filter(
        a,
        filter_radius,
        sx,
        sy,
        res,
    )

    prefac = 1e7  # hyperparameter 1
    pwr = 4.0  # hyperparameter 2
    c0 = 800#prefac * (filter_radius * 1 / res) ** pwr

    M1 = lambda a: mpa.constraint_solid(
        a,
        c0,
        eta_e,
        filter_f,
        threshold_f,
        res,
    )
    M2 = lambda a: mpa.constraint_void(
        a,
        c0,
        eta_d,
        filter_f,
        threshold_f,
        res,
    )
    return M1(d), M2(d)


def gen_latent_field(diam: float=projected_diam)->np.ndarray:
    """Generate the latent design field from the circle diameter.
    """
    x = np.linspace(-0.5 * sx, 0.5 * sx, Nx)
    y = np.linspace(-0.5 * sy, 0.5 * sy, Ny)
    xv, yv = np.meshgrid(x, y)

    cx, cy = 0.0, 0.0
    d = np.where(
        np.abs(xv - cx) ** 2 + np.abs(yv - cy) ** 2 < (0.5 * diam) ** 2,
        1.0,
        0.0,
    )

    return d

def find_latent_radius(min_length_scale: float)->float:
    """For a given lengthscale, find the latent diameter of the corresponding circle.
    
    """
    filter_radius = mpa.get_conic_radius_from_eta_e(min_length_scale, eta_e)
    threshold_f = lambda a: mpa.tanh_projection(a, beta, eta_i)

    filter_f = lambda a: mpa.conic_filter(
        a,
        filter_radius,
        sx,
        sy,
        res,
    )

    def J(diam, grad):
        x_expected = gen_latent_field(projected_diam)
        x = threshold_f(filter_f(gen_latent_field(diam)))
        err = np.mean(np.abs(x_expected-x)**2) 
        return err
    
    opt = nlopt.opt(nlopt.LN_NELDERMEAD, 1)
    opt.set_min_objective(J)
    opt.set_xtol_rel(1e-4)
    opt.set_maxeval(100)
    diam = opt.optimize([3.0])

    return diam


def plot_design_weights(d: np.ndarray, fname: str):
    """Saves a plot of the design weights.

    Args:
        d: the design weights as a flattened 1d array.
        fname: filename of the PNG image for saving.
    """
    fig, ax = plt.subplots(ncols=2)
    print(ax)
    extent = [x[0], x[-1], y[0], y[-1]]
    ax[0].imshow(d.reshape(Nx, Ny), extent=extent, cmap="binary")
    ax[0].set_xlabel("x")
    ax[0].set_ylabel("y")
    ax[0].set_title("solid")
    fig.savefig(fname + ".png", bbox_inches="tight", dpi=150)


if __name__ == "__main__":
    length_constraint = np.linspace(0.1, 2.0, 30)
    #length_constraint = [2.0]
    constraint_val = []
    for lc in length_constraint:
        diam = find_latent_radius(lc)
        d = gen_latent_field(diam)
        c_solid, c_void = constraint_solid_void(lc, d)
        constraint_val.append(c_solid)
        print(f"data:, {lc:.4f}, {c_solid:.6f}, {c_void:.6f}")
    
    plt.figure()
    plt.semilogy(length_constraint,constraint_val,'-o')
    plt.xlabel("lengthscale constraint $L$")
    plt.ylabel("constraint function value")
    plt.tight_layout()
    plt.savefig("opt_results.png")
    plt.show()