deepmie/qxnew.py

import snlay as sn
import numpy as np
from scipy import interpolate
from noise import pnoise1
import random
from tmm import coh_tmm
import multiprocessing as mp
import makeqx as mkq
#import pyximport; pyximport.install(pyimport=True)
#import mtmm as mtmm


# make a materials dictionary
matsdict = {
    1: "./materials/gold.dat",
    2: "./materials/silicon.dat",
    3: "./materials/silica.dat",
    4: "./materials/tio2.dat",
    5: "./materials/silver.dat",
}


def get_nk(datafile, wavelengths):
    """Reads the given file and returns the n+ik complex at
    the given wavelength after suitable interpolation
    :datafile: TODO
    :wavelength: TODO
    :returns: TODO
    """
    rawdisp = np.loadtxt(datafile)
    f_r = interpolate.interp1d(rawdisp[:, 0], rawdisp[:, 1])
    f_i = interpolate.interp1d(rawdisp[:, 0], rawdisp[:, 2])
    return f_r(wavelengths) + 1j * f_i(wavelengths)


def nk_2_eps(n):
    """TODO: Docstring for nk_2_epsnk_2_eps.
    :returns: complex epsilon given n and kappa 
    """
    eps_r = n.real ** 2 - n.imag ** 2
    eps_i = 2 * n.real * n.imag
    return eps_r + 1j * eps_i


def eps_2_nk(eps):
    """TODO: Docstring for nk_2_epsnk_2_eps.
    :returns: complex epsilon given n and kappa 
    """
    modeps = np.abs(eps)
    n_r = np.sqrt(0.5 * (modeps + eps.real))
    n_i = np.sqrt(0.5 * (modeps - eps.real))
    return n_r + 1j * n_i


def LL_mixing(fH, n_H, n_L):
    """TODO: Docstring for brugg_mixingbrugg_mixing.
    Given the volumne fraction of the higher index material, give the effective
    index of the layer
    :fH: volumne fraction from 0 to 1 
    :n_H: ri  of the higher index material 
    :n_L: ri of the lower index material 
    :returns: TODO
    """
    eH = nk_2_eps(n_H)
    eL = nk_2_eps(n_L)
    bigK = fH * (eH - 1) / (eH + 2) + (1 - fH) * (eL - 1) / (eL + 2)
    e_eff = (1 + 2 * bigK) / (1 - bigK)
    return eps_2_nk(e_eff)


def make_qxn(num_layers, dlevels=0):
    """TODO: Docstring for make_qxn.

    :num_layers: TODO
    :returns: TODO

    """
    uni = np.linspace(0, 2*num_layers, 2*num_layers, endpoint=False)
    rwalk = np.zeros_like(uni).reshape((num_layers,2))
    rwalk[:,0] = mkq.make_qx(num_layers)
    if dlevels != 0:
        rwalk[:,1] = mkq.digitize_qx(mkq.make_qx(num_layers), dlevels)
    else:
        rwalk[:,1] = mkq.make_qx(num_layers)
    return rwalk.flatten()



def make_nxdx(qx, d_min, d_max, wavelen=550, n_substrate=1.52):
    """TODO: Docstring for make_nxdx.
    :qx: TODO
    :n_substrate: TODO
    :layer_thickness: TODO
    :returns: TODO
    """
    num_layers = int(qx.size/2)
    qx = qx.reshape((num_layers, 2))
    #d_x = d_min + (d_max - d_min)*qx[:,0] 
    d_x = qx[:,0]
    d_x = [np.inf] + d_x.tolist() + [np.inf]
    sde = LL_mixing(
            fH=qx[:,1], n_H=2.58, n_L=1.45  # get_nk(datafile=matsdict[3], wavelengths=wavelen),
    )  # get_nk(datafile=matsdict[4], wavelengths=wavelen
    #   ))
    n_x = [1.0] + sde.tolist() + [n_substrate]
    return d_x, n_x


# def tmm_eval2(
#     qx,
#     cthick,
#     lam_low=400,  # in nm
#     lam_high=1200,
#     lam_pts=100,
#     ang_low=0,  # degrees
#     ang_high=90,
#     ang_pts=25,
#     n_subs=1.52,
# ):
#     """TODO: Docstring for tmm_eval.
#     :qx: TODO
#     :lam_low: TODO
#     :#in nm 
#         lam_high: TODO
#     :lam_pts: TODO
#     :ang_low: TODO
#     :#degrees
#         ang_high: TODO
#     :ang_pts: TODO
#     :returns: TODO

#     """
#     degree = np.pi / 180
#     lams = np.linspace(lam_low, lam_high, endpoint=True, num=lam_pts)
#     thetas = np.linspace(ang_high, ang_low, endpoint=True, num=ang_pts)
#     Rs = np.zeros((thetas.size, lams.size))
#     Rp = np.zeros((thetas.size, lams.size))

#     for tidx, theta in enumerate(thetas):
#         for lidx, lam in enumerate(lams):
#             d_x, n_x = make_nxdx(qx=qx, cthick=cthick, wavelen=lam, n_substrate=n_subs)
#             Rs[tidx, lidx] = 100 * coh_tmm(
#                 "s", n_x, d_x, th_0=theta * degree, lam_vac=lam
#             )
#             Rp[tidx, lidx] = 100 * coh_tmm(
#                 "p", n_x, d_x, th_0=theta * degree, lam_vac=lam
#             )
#     return Rs, Rp


# def tmm_eval(
#     qx,
#     cthick,
#     lam_low=400,  # in nm
#     lam_high=1200,
#     lam_pts=100,
#     ang_low=0,  # degrees
#     ang_high=90,
#     ang_pts=25,
#     n_subs=1.52,
# ):
#     """TODO: Docstring for tmm_eval.
#     :qx: TODO
#     :lam_low: TODO
#     :#in nm 
#         lam_high: TODO
#     :lam_pts: TODO
#     :ang_low: TODO
#     :#degrees
#         ang_high: TODO
#     :ang_pts: TODO
#     :returns: TODO

#     """
#     degree = np.pi / 180
#     lams = np.linspace(lam_low, lam_high, endpoint=True, num=lam_pts)
#     thetas = np.linspace(ang_high, ang_low, endpoint=True, num=ang_pts)
#     Rs = np.zeros((thetas.size, lams.size))
#     Rp = np.zeros((thetas.size, lams.size))

#     for tidx, theta in enumerate(thetas):
#         for lidx, lam in enumerate(lams):
#             d_x, n_x = make_nxdx(qx=qx, cthick=cthick, wavelen=lam, n_substrate=n_subs)
#             Rs[tidx, lidx] = 100 * coh_tmm(
#                 "s", n_x, d_x, th_0=theta * degree, lam_vac=lam
#             )
#             Rp[tidx, lidx] = 100 * coh_tmm(
#                 "p", n_x, d_x, th_0=theta * degree, lam_vac=lam
#             )
#     return Rs, Rp


def tmm_eval_wbk(qx, d_min , d_max, inc_ang, lam, n_subs=1.52):
    """TODO: Docstring for tmm_eval.
    :qx: TODO
    :returns: TODO
    """
    degree = np.pi / 180
    d_x, n_x = make_nxdx(qx=qx, d_min=d_min, d_max=d_max, wavelen=lam, n_substrate=n_subs)
    Rs = 100 * coh_tmm("s", n_x, d_x, th_0=inc_ang * degree, lam_vac=lam)
    #     Rp = 100*coh_tmm('p',n_x,d_x, th_0=inc_ang*degree,lam_vac=lam)
    return Rs


def tmm_eval_wsweep(
    qx, d_min , d_max, inc_ang, lam_low=400, lam_high=800, lam_pts=50, n_subs=1.52  # in nm
):
    """TODO: Docstring for tmm_eval.
    :qx: TODO
    :lam_low: TODO
     lam_high: TODO
    :lam_pts: TODO
    :returns: TODO
    """
    degree = np.pi / 180
    lams = np.linspace(lam_low, lam_high, endpoint=True, num=lam_pts)
    Rs = np.zeros(lams.size)
    # Rp = np.zeros(lams.size)
    for lidx, lam in enumerate(lams):
        d_x, n_x = make_nxdx(qx=qx, d_min=0, d_max=1, wavelen=lam, n_substrate=n_subs)
        Rs[lidx] = 100 * coh_tmm("s", n_x, d_x, th_0=inc_ang * degree, lam_vac=lam)
        # Rp[lidx] = 100*coh_tmm('p',n_x,d_x, th_0=inc_ang*degree,lam_vac=lam)
    return Rs


# def tmm_wrapper2(arg):
#     args, kwargs = arg
#     return tmm_eval_wbk(*args, **kwargs)


# def tmm_lam_parallel(
#     q_x, cthick, inc_ang, n_par=12, lam_low=400, lam_high=1200, lam_pts=100, **kwargs
# ):

#     jobs = []
#     pool = mp.Pool(n_par)
#     lams = np.linspace(lam_low, lam_high, endpoint=True, num=lam_pts)
#     for lam in lams:
#         jobs.append((q_x, cthick, inc_ang, lam))
#     arg = [(j, kwargs) for j in jobs]
#     answ = np.array(pool.map(tmm_wrapper2, arg))
#     pool.close()
#     return answ[:, 0], answ[:, 1]


# def tmm_wrapper(arg):
#     args, kwargs = arg
#     return tmm_eval_wsweep(*args, **kwargs)


# def tmm_eval_parallel(q_x, cthick, n_ang=25, n_par=10, **kwargs):
#     jobs = []
#     pool = mp.Pool(n_par)
#     angs = np.linspace(90, 0, endpoint=True, num=n_ang)
#     for ang in angs:
#         jobs.append((q_x, cthick, ang))
#     arg = [(j, kwargs) for j in jobs]
#     answ = np.array(pool.map(tmm_wrapper, arg))
#     pool.close()
#     return answ[:, 0, :], answ[:, 1, :]