Add 1D mode and FFT updates

docs/usage.rst

@@ -12,7 +12,11 @@ To enable autograd::
 
 To initialize the RCWA::
 
-  obj = grcwa.obj(nG,L1,L2,freq,theta,phi,verbose=0) # verbose=1 for output the actual nG
+  obj = grcwa.obj(nG, L1, L2, freq, theta, phi, verbose=0)
+
+The constructor accepts ``mode='2D'`` by default.  Set ``mode='1D'`` when one
+lattice vector is zero and grid layers use ``Nx=1`` or ``Ny=1`` to model
+1‑D gratings.
 
 To add layers, the order of adding will determine the layer order (1st added layer is 0-th layer, 2nd to be 1st layer, and so forth)::
 

example/ex1d.py

@@ -0,0 +1,50 @@
+"""Demonstrate 1-D binary gratings with Ny=1 and Nx=1."""
+import grcwa
+import numpy as np
+
+nG = 21
+freq = 1.0
+theta = 0.0
+phi = 0.0
+
+planewave = {'p_amp':1,'s_amp':0,'p_phase':0,'s_phase':0}
+
+# Variation along x (Ny=1)
+L1 = [0.5, 0]
+L2 = [0, 0]
+Nx = 40
+Ny = 1
+
+obj = grcwa.obj(nG, L1, L2, freq, theta, phi, verbose=1, mode='1D')
+obj.Add_LayerUniform(1., 1.)
+obj.Add_LayerGrid(0.2, Nx, Ny)
+obj.Add_LayerUniform(1., 1.)
+obj.Init_Setup()
+obj.MakeExcitationPlanewave(**planewave, order=0)
+
+ep = np.ones((Nx, Ny)) * 4.
+mask = np.arange(Nx) < Nx//2
+ep[mask,0] = 1.
+obj.GridLayer_geteps(ep.flatten())
+R, T = obj.RT_Solve(normalize=1)
+print('Ny=1: R=', R, ' T=', T)
+
+# Variation along y (Nx=1)
+L1 = [0, 0]
+L2 = [0, 0.5]
+Nx = 1
+Ny = 40
+
+obj = grcwa.obj(nG, L1, L2, freq, theta, phi, verbose=1, mode='1D')
+obj.Add_LayerUniform(1., 1.)
+obj.Add_LayerGrid(0.2, Nx, Ny)
+obj.Add_LayerUniform(1., 1.)
+obj.Init_Setup()
+obj.MakeExcitationPlanewave(**planewave, order=0)
+
+ep = np.ones((Nx, Ny)) * 4.
+mask = np.arange(Ny) < Ny//2
+ep[0,mask] = 1.
+obj.GridLayer_geteps(ep.flatten())
+R, T = obj.RT_Solve(normalize=1)
+print('Nx=1: R=', R, ' T=', T)

grcwa/backend.py

@@ -37,6 +37,8 @@ class NumpyBackend():
     trace = staticmethod(np.trace)
     fft2 = staticmethod(np.fft.fft2)
     ifft2 = staticmethod(np.fft.ifft2)
+    fft = staticmethod(np.fft.fft)
+    ifft = staticmethod(np.fft.ifft)
 
     sin = staticmethod(np.sin)
     cos = staticmethod(np.cos)
@@ -77,6 +79,8 @@ class AutogradBackend():
         trace = staticmethod(npa.trace)
         fft2 = staticmethod(npa.fft.fft2)
         ifft2 = staticmethod(npa.fft.ifft2)
+        fft = staticmethod(npa.fft.fft)
+        ifft = staticmethod(npa.fft.ifft)
 
         sin = staticmethod(npa.sin)
         cos = staticmethod(npa.cos)

grcwa/fft_funs.py

@@ -37,11 +37,16 @@ def get_conv(dN,s_in,G):
     G: shape (nG,2), 2 for Lk1,Lk2
     s_out: 1/N sum a_m exp(-2pi i mk/n), shape (nGx*nGy)
     '''
-    sfft = bd.fft2(s_in) * dN
-
-    gi = G[:, 0][:, None] - G[:, 0]
-    gj = G[:, 1][:, None] - G[:, 1]
-    s_out = sfft[gi, gj]
+    Nx, Ny = s_in.shape
+    if Ny == 1 or Nx == 1 or G.ndim == 1:
+        sfft = bd.fft(s_in.reshape(Nx * Ny)) * dN
+        gi = G[:, None] - G[None, :]
+        s_out = sfft[gi]
+    else:
+        sfft = bd.fft2(s_in) * dN
+        gi = G[:, 0][:, None] - G[:, 0]
+        gj = G[:, 1][:, None] - G[:, 1]
+        s_out = sfft[gi, gj]
     return s_out
 
 def get_fft(dN,s_in,G):
@@ -53,8 +58,13 @@ def get_fft(dN,s_in,G):
     s_out: 1/N sum a_m exp(-2pi i mk/n), shape (nGx*nGy)
     '''
 
-    sfft = bd.fft2(s_in)*dN
-    return sfft[G[:,0],G[:,1]]
+    Nx, Ny = s_in.shape
+    if Ny == 1 or Nx == 1 or G.ndim == 1:
+        sfft = bd.fft(s_in.reshape(Nx * Ny)) * dN
+        return sfft[G]
+    else:
+        sfft = bd.fft2(s_in) * dN
+        return sfft[G[:,0],G[:,1]]
 
 
 def get_ifft(Nx,Ny,s_in,G):
@@ -63,11 +73,14 @@ def get_ifft(Nx,Ny,s_in,G):
     '''
     dN = 1.0 / Nx / Ny
 
-    # Directly assign each Fourier coefficient to its corresponding
-    # location in the frequency domain array.  This is equivalent to
-    # the previous explicit loop but vectorised for efficiency.
-    s0 = bd.zeros((Nx, Ny), dtype=complex)
-    s0[G[:, 0], G[:, 1]] = s_in
-
-    s_out = bd.ifft2(s0)/dN
-    return s_out
+    if Ny == 1 or Nx == 1 or G.ndim == 1:
+        N = Nx * Ny
+        s0 = bd.zeros(N, dtype=complex)
+        s0[G] = s_in
+        s_out = bd.ifft(s0) / dN
+        return bd.reshape(s_out, (Nx, Ny))
+    else:
+        s0 = bd.zeros((Nx, Ny), dtype=complex)
+        s0[G[:, 0], G[:, 1]] = s_in
+        s_out = bd.ifft2(s0) / dN
+        return s_out

grcwa/kbloch.py

@@ -27,12 +27,15 @@ def Lattice_getG(nG,Lk1,Lk2,method=0):
     if not isinstance(nG, (int, np.integer)):
         raise TypeError('nG must be an integer')
 
-    if method == 0:
-        G,nG = Gsel_circular(nG, Lk1, Lk2)
-    elif method == 1:
-        G,nG = Gsel_parallelogramic(nG, Lk1, Lk2)
+    if np.linalg.norm(Lk1)==0 or np.linalg.norm(Lk2)==0:
+        G,nG = Gsel_1d(nG,method)
     else:
-        raise Exception('Truncation scheme is not included')
+        if method == 0:
+            G,nG = Gsel_circular(nG, Lk1, Lk2)
+        elif method == 1:
+            G,nG = Gsel_parallelogramic(nG, Lk1, Lk2)
+        else:
+            raise Exception('Truncation scheme is not included')
 
     return G,nG
 
@@ -42,12 +45,34 @@ def Lattice_SetKs(G, kx0, ky0, Lk1, Lk2):
     2pi factor is now included in the returned kx,ky
     '''
 
-    kx = kx0 + 2*np.pi*(Lk1[0]*G[:,0]+Lk2[0]*G[:,1])
-    ky = ky0 + 2*np.pi*(Lk1[1]*G[:,0]+Lk2[1]*G[:,1])
+    if G.ndim == 1:
+        if np.linalg.norm(Lk1) != 0:
+            kx = kx0 + 2*np.pi*Lk1[0]*G
+            ky = ky0 + 2*np.pi*Lk1[1]*G
+        else:
+            kx = kx0 + 2*np.pi*Lk2[0]*G
+            ky = ky0 + 2*np.pi*Lk2[1]*G
+    else:
+        kx = kx0 + 2*np.pi*(Lk1[0]*G[:,0]+Lk2[0]*G[:,1])
+        ky = ky0 + 2*np.pi*(Lk1[1]*G[:,0]+Lk2[1]*G[:,1])
 
     return kx,ky
 
 
+def Gsel_1d(nG, method):
+    """Generate 1-D G orders."""
+    m = nG // 2
+    if method == 0:
+        G = np.arange(-m, m + 1, dtype=int)
+    elif method == 1:
+        pos = np.arange(1, m + 1, dtype=int)
+        G = np.concatenate(([0], pos, -pos))
+    else:
+        raise Exception('Truncation scheme is not included')
+    nG = len(G)
+    return G, nG
+
+
 def Gsel_parallelogramic(nG, Lk1, Lk2):
     ''' From Liu's gsel.c'''
     u = np.linalg.norm(Lk1)

grcwa/rcwa.py

@@ -3,7 +3,7 @@
 from .kbloch import Lattice_Reciprocate,Lattice_getG,Lattice_SetKs
 
 class obj:
-    def __init__(self,nG,L1,L2,freq,theta,phi,verbose=1):
+    def __init__(self,nG,L1,L2,freq,theta,phi,verbose=1,mode="2D"):
         '''The time harmonic convention is exp(-i omega t), speed of light = 1
 
         Two kinds of layers are currently supported: uniform layer,
@@ -23,6 +23,11 @@ def __init__(self,nG,L1,L2,freq,theta,phi,verbose=1):
         self.theta = theta
         self.nG = nG
         self.verbose = verbose
+        if mode not in ["1D", "2D"]:
+            raise ValueError("mode must be '1D' or '2D'")
+        self.mode = mode
+        if self.mode == "1D" and not ((L1[0]==0 and L1[1]==0) or (L2[0]==0 and L2[1]==0)):
+            raise ValueError("One lattice vector must be zero in 1D mode")
         self.Layer_N = 0  # total number of layers
 
         # the length of the following variables = number of total layers
@@ -62,6 +67,8 @@ def Add_LayerUniform(self,thickness,epsilon):
         self.Uniform_N += 1
 
     def Add_LayerGrid(self,thickness,Nx,Ny):
+        if self.mode == "1D" and not (Nx == 1 or Ny == 1):
+            raise ValueError("Nx or Ny must be 1 in 1D mode")
         self.thickness_list.append(thickness)
         self.GridLayer_Nxy_list.append([Nx,Ny])
         self.id_list.append([1,self.Layer_N,self.Patterned_N,self.GridLayer_N])
@@ -90,7 +97,19 @@ def Init_Setup(self,Pscale=1.,Gmethod=0):
         ky0 = self.omega*bd.sin(self.theta)*bd.sin(self.phi)*bd.sqrt(self.Uniform_ep_list[0])
 
         # set up reciprocal lattice
-        self.Lk1, self.Lk2 = Lattice_Reciprocate(self.L1,self.L2)
+        if self.mode == "1D":
+            if self.L1[0]==0 and self.L1[1]==0:
+                self.Lk1 = bd.array([0.,0.])
+                self.Lk2 = bd.array([
+                    1./self.L2[0] if self.L2[0]!=0 else 0.,
+                    1./self.L2[1] if self.L2[1]!=0 else 0.])
+            else:
+                self.Lk2 = bd.array([0.,0.])
+                self.Lk1 = bd.array([
+                    1./self.L1[0] if self.L1[0]!=0 else 0.,
+                    1./self.L1[1] if self.L1[1]!=0 else 0.])
+        else:
+            self.Lk1, self.Lk2 = Lattice_Reciprocate(self.L1,self.L2)
         self.G,self.nG = Lattice_getG(self.nG,self.Lk1,self.Lk2,method=Gmethod)
 
         self.Lk1 = self.Lk1/Pscale

tests/test_kbloch.py

@@ -99,3 +99,35 @@ def old_get_conv(dN, s_in, G):
         ref = old_get_conv(dN, s, G)
         grcwa.set_backend('autograd')
         assert np.allclose(ref, new)
+
+    def test_getG_1d():
+        Lk1 = np.array([1.0, 0.0])
+        Lk2 = np.array([0.0, 0.0])
+        G1, n1 = grcwa.Lattice_getG(5, Lk1, Lk2, method=0)
+        assert G1.ndim == 1
+        assert n1 == len(G1)
+        assert G1[0] == -2 and G1[-1] == 2
+
+    def test_fft_1d():
+        N = 16
+        dN1 = 1.0 / N
+        G1 = np.arange(-2, 3)
+        s = np.random.random((N, 1))
+        grcwa.set_backend('numpy')
+        new = grcwa.get_fft(dN1, s, G1)
+        ref = np.fft.fft(s.flatten()) * dN1
+        grcwa.set_backend('autograd')
+        assert np.allclose(ref[G1], new)
+
+    def test_ifft_1d():
+        N = 8
+        dN1 = 1.0 / N
+        G1 = np.arange(-2, 3)
+        x = np.random.random(len(G1))
+        grcwa.set_backend('numpy')
+        new = grcwa.get_ifft(1, N, x, G1)
+        ref_full = np.zeros(N, dtype=complex)
+        ref_full[G1] = x
+        ref = np.fft.ifft(ref_full) / dN1
+        grcwa.set_backend('autograd')
+        assert np.allclose(ref.reshape(1, N), new)

tests/test_rcwa.py

@@ -1,5 +1,6 @@
 import numpy as np
 import grcwa
+import pytest
 from .utils import t_grad
 
 try:
@@ -61,6 +62,19 @@ def rcwa_assembly(epgrid,freq,theta,phi,planewave,pthick,Pscale=1.):
 
     return obj
 
+
+def rcwa_assembly_1d(Nx, Ny):
+    obj = grcwa.obj(21, [0.5, 0] if Ny==1 else [0,0], [0,0.5] if Nx==1 else [0,0],
+                    freq, 0, 0, verbose=0, mode='1D')
+    obj.Add_LayerUniform(1., 1.)
+    obj.Add_LayerGrid(0.2, Nx, Ny)
+    obj.Add_LayerUniform(1., 1.)
+    obj.Init_Setup()
+    obj.MakeExcitationPlanewave(1,0,0,0,order=0)
+    ep = np.ones((Nx, Ny))
+    obj.GridLayer_geteps(ep.flatten())
+    return obj
+
 
 def test_rcwa():
     ## compared to S4
@@ -88,6 +102,21 @@ def test_rcwa():
     Tx,Ty,Tz = obj.Solve_ZStressTensorIntegral(0)
     assert Tz<0
 
+def test_addlayergrid_error():
+    obj = grcwa.obj(11, [0.5,0], [0,0], freq, 0, 0, verbose=0, mode='1D')
+    with pytest.raises(ValueError):
+        obj.Add_LayerGrid(0.1, 3, 3)
+
+def test_rcwa_1d_x():
+    obj = rcwa_assembly_1d(40,1)
+    R,T = obj.RT_Solve(normalize=1)
+    assert np.isfinite(R)
+
+def test_rcwa_1d_y():
+    obj = rcwa_assembly_1d(1,40)
+    R,T = obj.RT_Solve(normalize=1)
+    assert np.isfinite(R)
+
 if AG_AVAILABLE:
     grcwa.set_backend('autograd')
     def test_epsgrad():