from blond.beam.beam import Beam, Proton
from blond.beam.distributions import bigaussian
from blond.beam.profile import CutOptions, FitOptions, Profile
from blond.monitors.monitors import BunchMonitor
from blond.plots.plot import Plot
from blond.utils import bmath as bm
# from blond.utils.bmath import use_gpu
import os
print(sys.argv)
if (len(sys.argv) < 3):
print("... <num_of_particles> <dt_plt> <gpu>")
exit(0)
if (len(sys.argv) == 4):
bm.use_gpu()
print("Using Gpu")
this_directory = os.path.dirname(os.path.realpath(__file__)) + '/'
try:
os.mkdir(this_directory + '../output_files')
except:
pass
try:
os.mkdir(this_directory + '../output_files/EX_01_fig')
except:
pass
start = time.time()
# Simulation parameters -------------------------------------------------------
# Bunch parameters
if args.get('monitorfile', None):
filename = args['monitorfile']
else:
filename = 'monitorfiles/sps-t{}-p{}-b{}-sl{}-approx{}-prec{}-r{}-m{}-se{}-w{}'.format(
n_iterations, n_particles, n_bunches, n_slices, approx,
args['precision'], n_turns_reduce, args['monitor'], seed,
worker.workers)
slicesMonitor = MultiBunchMonitor(filename=filename,
n_turns=np.ceil(n_iterations /
args['monitor']),
profile=profile,
rf=rf_station,
Nbunches=n_bunches)
if args['gpu'] > 0:
bm.use_gpu(gpu_id=worker.gpu_id)
profile.use_gpu()
tracker.use_gpu()
phaseLoop.use_gpu()
bm.enable_gpucache()
print(
f'Glob rank: [{worker.rank}], Node rank: [{worker.noderank}], Intra rank: [{worker.intrarank}], GPU rank: [{worker.gpucommrank}], hasGPU: {worker.hasGPU}'
)
worker.initDLB(args['loadbalance'], n_iterations)
delta = 0
worker.sync()
timing.reset()
start_t = time.time()