<jupyter_start><jupyter_text>Simple testing of FBT in Warp. Just transform beam in a drift. No solenoid included and no inverse transform.<jupyter_code>%matplotlib notebook
import sys
del sys.argv[1:]
from warp import *
from warp.data_dumping.openpmd_diag import particle_diag
import numpy as np
import os
from copy import deepcopy
import matplotlib.pyplot as plt
diagDir = 'diags/test/hdf5'

def cleanupPrevious(outputDirectory = diagDir):
    if os.path.exists(outputDirectory):
        files = os.listdir(outputDirectory)
        for file in files:
            if file.endswith('.h5'):
                os.remove(os.path.join(outputDirectory,file))

cleanupPrevious()

def setup():
	pass
##########################################
### Create Beam and Set its Parameters ###
##########################################

top.lrelativ = True
top.relativity = 1

beam = Species(type=Electron, name='Electron') 
beam.ekin = 55e6 #KE = 2.5 MeV
derivqty() #Sets addition derived parameters (such as beam.vbeam)

top.emitx    = 4.0*800e-6 / top.gammabar # geometric emittance: emit_full = 4 * emit_rms
top.emity    = 4.0*1e-6 / top.gammabar
beam.a0 = sqrt(top.emitx * 5.0)
beam.b0 = sqrt(top.emity * 5.0)
beam.ap0 = -1 * top.emitx * 0.0 / beam.a0
beam.bp0 = -1 * top.emity * 0.0 / beam.b0


beam.vthz = 0 #Sets the longitudinal thermal velocity (see iop_lin_002)
beam.ibeam = 0 # beam.ibeam/(top.gammabar**2) #Set correct current for relativity (see iop_lin_002)


top.npmax = 10000

w3d.distrbtn = "Gaussian0" 

w3d.cylinder = True  #Set True if running without envelope solver
#####################
### Setup Lattice ###
#####################


turnLength = 2.0e-3 #39.9682297148
steps = 2 #8000.

top.zlatstrt = 0#0.   #  z of lattice start (added to element z's on generate).  
top.zlatperi = 10.0#turnLength # Lattice periodicity

top.dt = turnLength / steps / beam.vbeam  

start = Marker()
drift1 = Drft(l=1e-3)
transform = Marker()
drift2 = Drft(l=1e-3)
end = Marker()

transformLine = start + drift1 + transform + drift2 + end

madtowarp(transformLine)
def FRBT(beta=5.0, alpha=0.0):
    """ 
        Transforms a matched flat beam to a round 'magnetized' beam.
    """

    gamma = (1. - alpha**2) / beta

    R = np.zeros([6,6],dtype='float64')
    R[0,0] = 1. + alpha
    R[0,1] = beta
    R[0,2] = 1. - alpha
    R[0,3] = -beta

    R[1,0] = -gamma
    R[1,1] = 1. - alpha
    R[1,2] = gamma
    R[1,3] = 1. + alpha

    R[2,0] = 1. - alpha
    R[2,1] = -beta
    R[2,2] = 1. + alpha
    R[2,3] = beta

    R[3,0] = gamma
    R[3,1] = 1. + alpha
    R[3,2] = -gamma
    R[3,3] = 1. - alpha

    R[4,4] = 2.
    R[5,5] = 2.

    R = 0.5 * R
    
    x = {}

    norm = {}
    for i in range(6):
        for j in range(6):
            norm[i,j] = 1.0
    norm[0,1] = norm[0,3] = norm[2,1] = norm[2,3] = 1./top.pgroup.uzp
    norm[1,0] = norm[1,2] = top.pgroup.uzp
    norm[3,0] = norm[3,2] = top.pgroup.uzp

    x = {}
    x[0] = np.copy(top.pgroup.xp)
    x[1] = np.copy(top.pgroup.uxp)
    x[2] = np.copy(top.pgroup.yp)
    x[3] = np.copy(top.pgroup.uyp)
    x[4] = np.copy(top.pgroup.zp)
    x[5] = np.copy(top.pgroup.uzp)

    print x[0].shape

    holding = []

    for i in range(6):
        val = 0
        for j in range(6):
            val += R[i,j] * x[j] * norm[i,j]
        
        holding.append(val)


    top.pgroup.xp = holding[0]
    top.pgroup.uxp = holding[1]
    top.pgroup.yp = holding[2]
    top.pgroup.uyp = holding[3]
    top.pgroup.zp = holding[4]
    top.pgroup.uzp = holding[5]
    
#     print "Transform!"
################################
### 3D Simulation Parameters ###
################################


top.prwall = pr1 = 0.14

#Set cells
w3d.nx = 128
w3d.ny = 128
w3d.nz = 1



#Set boundaries
w3d.xmmin = -0.10
w3d.xmmax =  0.10 
w3d.ymmin = -0.10 
w3d.ymmax =  0.10 
w3d.zmmin = -2e-3
w3d.zmmax =  2e-3

top.pboundxy = 0            # Absorbing Boundary for particles

top.ibpush   = 2            # set type of pusher to  vXB push without tan corrections
                            ## 0:off, 1:fast, 2:accurate

top.fstype = -1
############################
### Particle Diagnostics ###
############################

diagP0 = particle_diag.ParticleDiagnostic( period=1, top=top, w3d=w3d,
        species= { species.name : species for species in listofallspecies },
        comm_world=comm_world, lparallel_output=False, write_dir = diagDir[:-4] )

diagP = particle_diag.ParticleDiagnostic( period=1, top=top, w3d=w3d,
        species= { species.name : species for species in listofallspecies },
        comm_world=comm_world, lparallel_output=False, write_dir = diagDir[:-4] )

installbeforestep( diagP0.write )
installafterstep( diagP.write )
#################################
### Generate and Run PIC Code ###
#################################

package("wxy")
generate()
fieldsolve()


#installafterstep(thin_lens_lattice)
#Execute First Step

installbeforestep(FRBT)

step(1)
def readparticles(filename):
    """
    Reads in openPMD compliant particle file generated by Warp's ParticleDiagnostic class.

    Parameters:
        filename (str): Path to a ParticleDiagnostic output file.
    Returns:
        particle_arrays (dict): Dictionary with entry for each species in the file that contains an array
         of the 6D particle coordinates.
    """

    dims = ['momentum/x', 'position/y', 'momentum/y', 'position/z', 'momentum/z']
    particle_arrays = {}

    f = h5.File(filename, 'r')

    if f.attrs.get('openPMD') is None:
        print "Warning!: Not an openPMD file. This may not work."

    step = f['data'].keys()[0]
    species_list = f['data/%s/particles' % step].keys()

    for species in species_list:
        parray = f['data/%s/particles/%s/position/x' % (step, species)]
        for dim in dims:
            parray = np.column_stack((parray, f['data/%s/particles/%s/' % (step, species) + dim]))

        particle_arrays[species] = parray

    return particle_arrays

def convertunits(particlearray):
    """
    Putting particle coordinate data in good ol'fashioned accelerator units:
        x: m
        x': ux/uz
        y: m
        y': uy/uz
        z: m
        p: MeV/c

    """
    dat = deepcopy(particlearray)  # Don't copy by reference
    dat[:, 1] = dat[:, 1] / dat[:, 5]
    dat[:, 3] = dat[:, 3] / dat[:, 5]
    dat[:, 5] = dat[:, 5] / 5.344286E-22

    return dat

def svecplot(array):
    fig = plt.figure(figsize = (8,8))
    Q = plt.quiver(array[:,0],array[:,2],array[:,1],array[:,3])
    plt.quiverkey(Q,0.0, 0.92, 0.002, r'$2', labelpos='W')
    xmax = np.max(array[:,0])
    xmin = np.min(array[:,0])
    plt.xlim(1.5*xmin,1.5*xmax)
    plt.ylim(1.5*xmin,1.5*xmax)
    plt.show()
init = convertunits(readparticles('diags/test/hdf5/data00000000.h5')['Electron'])
fin = convertunits(readparticles('diags/test/hdf5/data00000001.h5')['Electron'])
svecplot(init)
plt.title("Initial Flat Beam")
plt.xlabel("x (m)")
plt.ylabel("y (m)")
svecplot(fin)
plt.title("Magnetized Beam after FRBT")
plt.xlabel("x (m)")
plt.ylabel("y (m)")
def vortex_check(init):
    beta = 5.0
    alpha = 0
    gamma = (1 - alpha**2) / beta
    x1 = ((1+alpha) * init[0,0] + (beta) * init[0,1] + (1-alpha) * init[0,2] + (-beta) * init[0,3]) * 0.5
    x2 = ((-gamma) * init[0,0] + (1-alpha) * init[0,1] + (gamma) * init[0,2] + (1+alpha) * init[0,3]) * 0.5
    y1 = ((1-alpha) * init[0,0] + (-beta) * init[0,1] + (1+alpha) * init[0,2] + (beta) * init[0,3]) * 0.5
    y2 = ((gamma) * init[0,0] + (1+alpha) * init[0,1] + (-gamma) * init[0,2] + (1-alpha) * init[0,3]) * 0.5
    print x1, fin[0,0]
    print x2, fin[0,1]
    print y1, fin[0,2]
    print y2, fin[0,3]
def calc_emittance(array):
    xemit = np.sqrt(np.average(array[:,0]**2) * np.average(array[:,1]**2) - np.average(array[:,0] * array[:,1])**2 )
    yemit = np.sqrt(np.average(array[:,2]**2) * np.average(array[:,3]**2) - np.average(array[:,2] * array[:,3])**2 )
    
    return xemit,yemit
epsx0,epsy0 = calc_emittance(init)
epsxf,epsyf = calc_emittance(fin)
print "Initial:\n x-emit: %s  Initial y-emit: %s" % (epsx0,epsy0)
print "After Transform:\n x-emit: %s   y-emit: %s" % (epsxf,epsyf)<jupyter_output>Initial:
 x-emit: 7.31384527417e-06  Initial y-emit: 9.17185992956e-09
After Transform:
 x-emit: 3.65810665997e-06   y-emit: 3.66491195757e-06
{'fraction_non_alphanumeric': 0.1980162090238297, 'fraction_numerical': 0.05370751179387928, 'mean_word_length': 3.037109375, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 2, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '7160', 'license': 'Apache-2.0', 'n_tokens_mistral': 3464, 'n_tokens_neox': 3222, 'n_words': 721, 'path': 'rscooler/FBT/warp/FRBT_test.ipynb', 'repo_name': 'radiasoft/mcool', 'size': '8267'}