#!/usr/bin/env python# -*- coding: utf-8 -*-# Author: Benjamin Vial# This file is part of nannos# License: GPLv3# See the documentation at nannos.gitlab.io__all__=["Lattice"]importfunctoolsfrom.importbackendasbkfrom.importget_backend,get_typesfrom.constantsimportpifrom.geometryimport*from.layersimportLayerfrom.utilsimportis_scalarFLOAT,COMPLEX=get_types()
[docs]classLattice:"""A lattice object defining the unit cell. Parameters ---------- basis_vectors : tuple The lattice vectors :math:`((u_x,u_y),(v_x,v_y))`. For mono-periodic gratings, specify the x-periodicity with a float `a`. discretization : int or tuple of length 2 Spatial discretization of the lattice. If given an integer N, the discretization will be (N, N). truncation : str The truncation method, available values are "circular" and "parallelogrammic" (the default is "circular"). This has no effect for mono-periodic gratings. harmonics_array : array of shape (2, nh) Array of harmonics. If specified, this is used instead of harmonics generated by the lattice object and nh is deduced from it. """def__init__(self,basis_vectors,discretization=(2**8,2**8),truncation="circular",harmonics_array=None,):ifis_scalar(discretization):discretization=[discretization,discretization]else:discretization=list(discretization)iftruncationnotin["circular","parallelogrammic"]:raiseValueError(f"Unknown truncation method '{truncation}', please choose between 'circular' and 'parallelogrammic'.")self.is_1D=is_scalar(basis_vectors)ifself.is_1D:self.truncation="1D"self.basis_vectors=(basis_vectors,0),(0,1)self.discretization=(discretization[0],1)else:self.truncation=truncationself.basis_vectors=basis_vectorsself.discretization=tuple(discretization)self.harmonics_array=harmonics_array@propertydefarea(self):ifself.is_1D:returnself.basis_vectors[0][0]v=self.basis_vectorsreturnbk.linalg.norm(bk.cross(v[0],v[1]))@propertydefmatrix(self):"""Basis matrix. Returns ------- array like Matrix containing the lattice vectors (L1,L2). """returnbk.array(self.basis_vectors,dtype=FLOAT).T@propertydefreciprocal(self):"""Reciprocal matrix. Returns ------- array like Matrix containing the lattice vectors (K1,K2) in reciprocal space. """return2*pi*bk.linalg.inv(self.matrix).T
[docs]defget_harmonics(self,nh,sort=False):"""Get the harmonics with a given truncation method. Parameters ---------- nh : int Number of harmonics. Returns ------- G : list of tuple of integers of length 2 Harmonics (i1, i2). nh : int The number of harmonics after truncation. """harmonics,nh=self._get_harmonics(nh)ifsort:srt=bk.argsort(harmonics[0])returnharmonics[:,srt],nhreturnharmonics,nh
def_get_harmonics(self,nh):ifself.harmonics_arrayisnotNone:nh=bk.shape(self.harmonics_array)[1]returnself.harmonics_array,nhifint(nh)!=nh:raiseValueError("nh must be integer.")ifself.truncation=="circular":return_circular_truncation(nh,self.reciprocal)elifself.truncation=="parallelogrammic":return_parallelogramic_truncation(nh,self.reciprocal)else:return_one_dim_truncation(nh)defno1d(func):@functools.wraps(func)definner(self,*args,**kwargs):ifself.is_1D:raiseValueError(f"Cannot use method {func.__name__} for 1D gratings, please use stripe")returnfunc(self,*args,**kwargs)returninner@propertydefunit_grid(self):"""Unit grid in cartesian coordinates. Returns ------- array like The unit grid of size equal to the attribute `discretization`. """Nx,Ny=self.discretizationx0=bk.array(bk.linspace(0,1.0,Nx))y0=bk.array(bk.linspace(0,1.0,Ny))x_,y_=bk.meshgrid(x0,y0,indexing="ij")returnbk.stack([x_,y_])@propertydefgrid(self):"""Grid in lattice vectors basis. Returns ------- array like The grid of size equal to the attribute `discretization`. """returnself.transform(self.unit_grid)
[docs]deftransform(self,grid):"""Transform from cartesian to lattice coordinates. Parameters ---------- grid : tuple of array like Input grid, typically obtained by meshgrid. Returns ------- array like Transformed grid in lattice vectors basis. """ifget_backend()=="torch":returnbk.tensordot(self.matrix,grid.double(),dims=([1],[0]))else:returnbk.tensordot(self.matrix,grid,axes=(1,0))
[docs]defones(self):"""Return a new array filled with ones. Returns ------- array like A uniform complex 2D array with shape ``self.discretization``. """returnbk.ones(self.discretization,dtype=COMPLEX)
[docs]defzeros(self):"""Return a new array filled with zeros. Returns ------- array like A uniform complex 2D array with shape ``self.discretization``. """returnbk.zeros(self.discretization,dtype=COMPLEX)
[docs]defconstant(self,value):"""Return a new array filled with value. Returns ------- array like A uniform complex 2D array with shape ``self.discretization``. """returnself.ones()*value
[docs]defgeometry_mask(self,geom):"""Return a geametry boolean mask discretized on the lattice grid. Parameters ---------- geom : Shapely object The geometry. Returns ------- array of bool The shape mask. """returngeometry_mask(geom,self,*self.discretization)
[docs]@no1ddefpolygon(self,vertices):"""Return a boolean mask for a polygon. Parameters ---------- vertices : array like The vertices of the polygon. Returns ------- array of bool The shape mask. """returnpolygon(vertices,self,*self.discretization)
[docs]@no1ddefcircle(self,center,radius):"""Return a boolean mask for a circle. Parameters ---------- center : tuple of float The center of the circle. radius : float The radius of the circle. Returns ------- array of bool The shape mask. """returncircle(center,radius,self,*self.discretization)
[docs]@no1ddefellipse(self,center,radii,rotate=0):"""Return a boolean mask for an ellipse. Parameters ---------- center : tuple of float The center of the ellipse. radii : tuple of float The radii of the ellipse. rotate : float The rotation of the ellipse in degrees. Returns ------- array of bool The shape mask. """returnellipse(center,radii,self,*self.discretization,rotate=rotate)
[docs]@no1ddefsquare(self,center,width,rotate=0):"""Return a boolean mask for a square. Parameters ---------- center : tuple of float The center of the square. width : float The width of the square. rotate : float, optional The rotation angle in degrees. Defaults to 0. Returns ------- array of bool The shape mask. """returnsquare(center,width,self,*self.discretization,rotate=rotate)
[docs]@no1ddefrectangle(self,center,widths,rotate=0):"""Return a boolean mask for a rectangle. Parameters ---------- center : tuple of float The center of the rectangle. widths : tuple of float The widths of the rectangle. Returns ------- array of bool The shape mask. """returnrectangle(center,widths,self,*self.discretization,rotate=rotate)
[docs]defstripe(self,center,width):"""Return a boolean mask for a stripe along the x-axis. Parameters ---------- center : float The center of the stripe. width : float The width of the stripe. Returns ------- array of bool A boolean mask of the stripe. """returnbk.abs(self.grid[0]-center)<=width/2
[docs]defLayer(self,name="layer",thickness=0,epsilon=1,mu=1,lattice=None,tangent_field=None,tangent_field_type="opt",):""" Return a new Layer object. Parameters ---------- name : str, optional Name of the layer. The default is "layer". thickness : float, optional Thickness of the layer. The default is 0. epsilon : complex or array_like, optional Relative permittivity. The default is 1. mu : complex or array_like, optional Relative permeability. The default is 1. lattice : nannos.Lattice, optional The lattice object. The default is None. tangent_field : array_like, optional The tangent field. The default is None. tangent_field_type : str, optional Type of tangent field ('opt' or 'fft'). The default is "opt". Returns ------- nannos.Layer The new layer object. """returnLayer(name,thickness,epsilon,mu,self,tangent_field,tangent_field_type,)
