###########################################################################
# Scriptfile: usr_custom_source.lsf
#
# Description:
# This is an example of how to create an arbitrary source using the
# usr_createsource.lsf script file.  In this example, we create a radially 
# or azumithally polarized source using either a simple paraxial beam profile
# definition, or a k-space beam profile definition.
# 
# The k-space beam profile definition is based on the method in the
# following reference:
# M. Mansuripur, "Distribution of light at and near the focus of high-numerical
# -aperture objectives," J. Opt. Soc. Am. A 3,2086-2093 (1986).
# 
###########################################################################

# clear all current variables
clear;
switchtolayout; 

# set desired beam polarization (radial=1, azimuthal=2)
pol = 1;

# define position vectors
x = linspace(-10e-6,10e-6,401);
y = linspace(-10e-6,10e-6,401);
z = 0;                             # this is z z-normal source
X = meshgridx(x,y);
Y = meshgridy(x,y);

lambda0 = 0.5e-6;
f = c/lambda0;
w = 2*pi*f;
k = 2*pi/lambda0;

# Two simple methods for defining the beam are provided.  
# 1: very simple near field definition
# 2: more complex K space definition.  
# NOTE: Both methods are approximate.
defineBeamMethod = 2;  


# define beam using a simple near field definition
if (defineBeamMethod==1) {

  ######################################
  # Specify beam envelope

  ## using FWHM
  # FWHM = 5e-6;  
  # sigma = FWHM / ( 2*sqrt(2*log(2)) );
  # envelope = exp( - (X^2)/(2*sigma^2) - (Y^2)/(2*sigma^2) );

  # using field intensity at 1/e^2
  w0 = 5e-6;  
  envelope = exp( - (X^2)/(w0^2) - (Y^2)/(w0^2) );


  #######################################
  # specify near fields
  phi=atan2(X,Y);  
  
  if(pol==1){    # if radial polarization
      Ex = sin(phi);
      Ey = cos(phi);
      Ez = 0;
  }else{  # if azimuthal polarization
      Ex = -cos(phi);
      Ey = sin(phi);
      Ez = 0;
  }
  Ex = Ex*envelope;
  Ey = Ey*envelope;
  Ez = Ez*envelope;

  # Package field data into the EM fields dataset that can be 
  # loaded into the Import source
  EM = rectilineardataset("EM fields",x,y,z);
  EM.addparameter("lambda",c/f,"f",f);   
  EM.addattribute("E",Ex,Ey,Ez);
}

# create a radial beam in k-space, with approximate NA
if (defineBeamMethod==2) {
  NA = 0.2;

  kx = linspace(-k,k,200);
  ky = linspace(-k,k,200);
  Kx = meshgridx(kx,ky);
  Ky = meshgridy(kx,ky);
  phi = atan2(Ky,Kx);
  theta = real(acos(sqrt(1-Kx^2/k^2-Ky^2/k^2)));
  envelope = exp(-0.5*(Kx^2 + Ky^2)/(NA*k)^2);
  
  if(pol==1){ # if radial polarization  
      Exk = cos(phi)*cos(theta)*envelope;
      Eyk = sin(phi)*cos(theta)*envelope;
      Ezk = sin(theta)*envelope;
      Ex = -1i*czt(Exk,kx,ky,x,y);
      Ey = -1i*czt(Eyk,kx,ky,x,y);
      Ez = -1i*czt(Ezk,kx,ky,x,y);
    
      Hxk = -sqrt(eps0/mu0)*sin(phi)*envelope;
      Hyk = sqrt(eps0/mu0)*cos(phi)*envelope;
      Hx = -1i*czt(Hxk,kx,ky,x,y);
      Hy = -1i*czt(Hyk,kx,ky,x,y);
      Hz = 0*Ex;
  } else{ # if azimuthal polarization
      Exk = -sin(phi)*envelope;
      Eyk = cos(phi)*envelope;
      Ezk = 0*Exk;
      Ex = -1i*czt(Exk,kx,ky,x,y);
      Ey = -1i*czt(Eyk,kx,ky,x,y);
      Ez = 0*Ex; #1i*czt(Ezk,kx,ky,x,y);
    
      Hxk = -sqrt(eps0/mu0)*cos(phi)*cos(theta)*envelope;
      Hyk = -sqrt(eps0/mu0)*sin(phi)*cos(theta)*envelope;
      Hzk = sqrt(eps0/mu0)*sin(theta)*envelope;
      Hx = -1i*czt(Hxk,kx,ky,x,y);
      Hy = -1i*czt(Hyk,kx,ky,x,y);
      Hz = -1i*czt(Hzk,kx,ky,x,y);
  }
  # scale field so |E|^2=1
  #   this is not required
  E2 = abs(Ex)^2+abs(Ey)^2+abs(Ez)^2;  
  scaleFactor = sqrt(max(E2));
  Ex = Ex/scaleFactor;
  Ey = Ey/scaleFactor;
  Ez = Ez/scaleFactor;
  Hx = Hx/scaleFactor;
  Hy = Hy/scaleFactor;
  Hz = Hz/scaleFactor;
              
  # Package field data into the EM fields dataset that can be 
  # loaded into the Import source
  EM = rectilineardataset("EM fields",x,y,z);
  EM.addparameter("lambda",c/f,"f",f);   
  EM.addattribute("E",Ex,Ey,Ez);
  EM.addattribute("H",Hx,Hy,Hz);
}


# directly load dataset into source
select("import_beam");
importdataset(EM);
set("center wavelength",lambda0);
set("wavelength span",0);

# save dataset to .mat file which can be imported at a later time
matlabsave("EM.mat",EM);

# run simulation;
run;


# plot field data
E_forward = getresult("T","E");
E_propagate = getresult("YZ","E");
T = getresult("T","T");
visualize(EM);
visualize(E_forward);
visualize(E_propagate);
visualize(T);

# plot E field data in vector plot
E_dataset = rectilineardataset("EM fields",x,y,z);
E_dataset.addparameter("lambda",c/f,"f",f);   
E_dataset.addattribute("E",Ex,Ey,Ez);

vectorplot(E_dataset);
vectorplot(E_forward);