%**************************************************************************
%This function generates the field from an EPU quadrant in the coordinate
%system where z is in the beam direction, r is away from the magnets of the
%quadrant, and s is along the magnets perpendicular to r and z.
%r, s, z makes a right handed system.
%z=0 is at the first pole.
%
%Input:
%num_per, number of periods in the undulator, including ends
%lambda_u, undulator period (m)
%kr, parameter giving the radial exponential potential decrease (1/m)
%ks, parameter giving the azimuther potential roll off (1/m)
%phi_0, peak magnetic scalar potential (Tm)
%r, r-position to give the field at (m)
%s, s-position to give the field at (m)
%z, z-position to give the field at (m)
%
%Output:
%br, Br field at (r, s, z) (T)
%bs, Bs field at (r, s, z) (T)
%bz, Bz field at (r, s, z) (T)
%
%Zachary Wolf
%3/31/15

function [br, bs, bz] = gen_epu_quadrant_field(num_per, lambda_u, kr, ks, phi_0, r, s, z)

%Define the undulator parameters
ku = 2 * pi / lambda_u;

%First calculate the fields for all z, afterward put in ends
%phi = phi_0 cos(ks s) exp(-kr r) cos(ku z)
%b = grad(phi)
br = - kr * phi_0 * cos(ks * s) * exp(-kr * r) * cos(ku * z);
bs = - ks * phi_0 * sin(ks * s) * exp(-kr * r) * cos(ku * z);
bz = - ku * phi_0 * cos(ks * s) * exp(-kr * r) * sin(ku * z);

%Make the ends

%z before undulator
z_start_xy = -lambda_u / 4;
z_start_z = 0;
if (z < z_start_xy)
    br = 0;
    bs = 0;
end
if (z < z_start_z)
    bz = 0;
end

%z in first half period
if (z >= z_start_xy) && (z < z_start_xy + lambda_u / 2)
    br = 0.25 * br;
    bs = 0.25 * bs;
end
if (z >= z_start_z) && (z < z_start_z + lambda_u / 2)
    bz = 0.25 * bz;
end

%z in second half period
if (z >= z_start_xy + lambda_u / 2) && (z < z_start_xy + lambda_u)
    br = 0.75 * br;
    bs = 0.75 * bs;
end
if (z >= z_start_z + lambda_u / 2) && (z < z_start_z + lambda_u)
    bz = 0.75 * bz;
end

%z after undulator
z_end_xy = num_per * lambda_u - lambda_u / 4;
z_end_z = num_per * lambda_u;
if (z > z_end_xy)
    br = 0;
    bs = 0;
end
if (z > z_end_z)
    bz = 0;
end

%z in first half period from end
if (z <= z_end_xy) && (z > z_end_xy - lambda_u / 2)
    br = 0.25 * br;
    bs = 0.25 * bs;
end
if (z <= z_end_z) && (z > z_end_z - lambda_u / 2)
    bz = 0.25 * bz;
end

%z in second half period from end
if (z <= z_end_xy - lambda_u / 2) && (z > z_end_xy - lambda_u)
    br = 0.75 * br;
    bs = 0.75 * bs;
end
if (z <= z_end_z - lambda_u / 2) && (z > z_end_z - lambda_u)
    bz = 0.75 * bz;
end

end
