simulation.py

import numpy as np
import meep as mp


class Simulation:
    """Simulation of the components of the greens tensor."""

    def __init__(self, size, src_pos, fcen, df, res=10, pml=0.7,
                 eps=mp.Vector3(12, 12, 12)):
        """Initialize a Simulation instance.
        Args:
            size (mp.Vector3): size of the simulation box
            src_pos (mp.Vector3): position of the gaussian source current
            fcen (float): frequency of the source
            df (float): frequency width of the gaussian
            res (int): resolution of the simulation
            pml (float): size of perfectly matching layer
        Returns:
            None.
        """
        self.size = size
        self.src_pos = src_pos
        self.fcen = fcen
        self.df = df
        self.eps = eps
        self.geom = []
        self.level_set = np.zeros((int(size.y), int(size.x)))
        self.res = res
        self.pml = [mp.PML(pml)]
        self.gSource = mp.GaussianSource(self.fcen, fwidth=self.df)
        self.jw = self.gSource.fourier_transform(self.fcen)
        self.G = np.zeros((3, 3, int(self.size.y)*self.res,
                           int(self.size.x)*self.res), dtype=complex)
        self.r = self._r_cellmiddle()
        self.sim = 0

    def _sim_Efourier(self, comp):
        """Simulate the E-field from a point source with comp polarization.
        Args:
            comp (mp.Const): mp.Ex, mp.Ey or mp.Ey
        Returns:
            np.arrays: Ex, Ey, Ez: fourier transformed fields.
        """
        src = [mp.Source(self.gSource,
                         component=comp,
                         center=self.src_pos, amplitude=1.0)]
        sim = mp.Simulation(cell_size=self.size,
                            sources=src, resolution=self.res, boundary_layers=self.pml,
                            force_complex_fields=True, force_all_components=True,
                            geometry=self.geom, eps_averaging=False)
        dft_obj = sim.add_dft_fields(
            [mp.Ex, mp.Ey, mp.Ez], self.fcen, self.fcen, 1,
            center=mp.Vector3(), size=self.size)
        sim.run(until_after_sources=100)
        self.sim = sim
        ex = sim.get_dft_array(dft_obj, mp.Ex, 0).transpose()
        ey = sim.get_dft_array(dft_obj, mp.Ey, 0).transpose()
        ez = sim.get_dft_array(dft_obj, mp.Ez, 0).transpose()
        ex = self._interpolate_E(ex, mp.Ex)
        ey = self._interpolate_E(ey, mp.Ey)
        ez = self._interpolate_E(ez, mp.Ez)
        return ex, ey, ez

    def sim_to_file(self, direct, comp):
        """Simulate continuous source and write to a hdf5 file.
        Args:
            filename (str): filename
            comp (mp.Const): component of the field to simulate
        Returns:
            None.
        """
        if direct[-1] == '/':
            direct = direct[:-1]
        src = [mp.Source(self.gSource,
                         component=comp,
                         center=self.src_pos)]
        sim = mp.Simulation(cell_size=self.size,
                            sources=src, resolution=self.res, boundary_layers=self.pml,
                            force_complex_fields=True, force_all_components=True,
                            geometry=self.geom, eps_averaging=False)

        A = '{:}/analyze-eps-000000.00.h5'.format(direct)
        a = 'yarg:0.5'
        out = mp.output_png(
            comp, '-Zc dkbluered -m 0.5 -M 0.5 -A {:} -a {:}'.format(A, a))
        # sim.use_output_directory(dname=direct)
        # sim.run(mp.at_beginning(mp.output_epsilon),
        #         mp.at_every(0.3, out), until=200)
        sim.run(mp.at_beginning(mp.output_epsilon),
                mp.to_appended("ez", mp.at_every(0.6, mp.output_efield_z)),
                until=200)

    def simulate_G(self, comp):
        """Calculate the components of G generated by j_comp.
        Args:
            comp (mp.Const): mp.Ex, mp.Ey, mp.Ez
        Returns:
            None.
        """
        pol = [mp.Ex, mp.Ey, mp.Ez]
        ex, ey, ez = self._sim_Efourier(comp)
        self.G[0, pol.index(comp)] = ex/(1j*2*np.pi*self.fcen*self.jw)
        self.G[1, pol.index(comp)] = ey/(1j*2*np.pi*self.fcen*self.jw)
        self.G[2, pol.index(comp)] = ez/(1j*2*np.pi*self.fcen*self.jw)

    def _r_cellmiddle(self):
        """Calculate the position values for the middle of the cells.
        In the format r[y,x] = (x,y,z)."""
        r = np.zeros((int(self.size.y)*self.res,
                      int(self.size.x)*self.res, 3))
        for y in range(len(r)):
            for x in range(len(r[0])):
                r[y, x] = np.array([x/self.res-self.size.x/2+1/(2*self.res),
                                    y/self.res-self.size.y/2+1/(2*self.res), 0])
        return r

    def _interpolate_E(self, E, comp):
        """Linear interpolation of E to the middle of the cell.
        Args:
            comp (mp.Const): mp.Ex, mp.Ey, mp.Ez
        Returns:
            np.array: 2d array of the E-field
        """
        centE = np.zeros_like(E)
        if comp == mp.Ex:
            centE[0: -1, :] = (E[0: -1, :] + E[1:, :])/2
            m = (E[-1, :] - E[-2, :]) * self.res
            centE[-1, :] = m/(2*self.res)
        elif comp == mp.Ey:
            centE[:, 0: -1] = (E[:, 0: -1] + E[:, 1:])/2
            m = (E[:, -1] - E[:, -2])*self.res
            centE[:, -1] = m/(2*self.res)
        elif comp == mp.Ez:
            centE[: -1, : -1] = (E[: -1, : -1] + E[1:, 1:])/2
            E[-1, :] = E[-2, :]
            E[:, -1] = E[:, -2]
        return centE

    def update_geom(self, geom):
        """Append a geometry to the existing ones.
        Args:
            geom (mp.Geometry): the geometry
        Returns:
            None.
        """
        self.geom.append(geom)

    def phi_to_geom(self, phi):
        """Create a geometry from the level set function phi (array).
        Where phi is negative, matter will be placed.
        Args:
            phi (np.array): 2d array representing the level set function
        Returns:
            None.
        """
        geom = []
        size = mp.Vector3(1.0001/self.res, 1.0001/self.res, mp.inf)
        for y in range(len(phi)):
            for x in range(len(phi[0])):
                if phi[y, x] <= 0:
                    cent = mp.Vector3(self.r[y, x][0], self.r[y, x][1])
                    geom.append(mp.Block(size, center=cent,
                                         material=mp.Medium(epsilon_diag=self.eps)))
        self.geom = geom

    def get_G(self, i, j):
        """Return the ij component of the tensor."""
        if (np.max(self.G[i, j]) == 0 and np.min(self.G[i, j]) == 0):
            print("Component G_{:}{:} not yet calculated!".format(i, j))
            return None
        return self.G[i, j]

    def get_dielectric(self):
        """Return a numpy array of the value of the dielectric.
        Args:
            None
        Returns:
            np.array: 2d array.
        """
        eps = self.sim.get_array(center=mp.Vector3(), size=self.size,
                                 component=mp.Dielectric)
        return eps.transpose()