def generate_3D_animation(x,
y,
z,
sdir,
title="",
npts=360,
xlabel="",
ylabel="",
zlabel="",
mode='default',
get_2D=True):
tdir = sdir + "/tmp/"
mkdir_p(tdir)
if mode == 'default':
alim, elim = 360, 360
for ang in range(0, 360, int(360 / npts)):
scatter_3D(x,
y,
z,
sdir=tdir + "{:5d}.png".format(ang),
azim=ang % alim,
elev=ang % alim,
xlabel=xlabel,
ylabel=ylabel,
zlabel=zlabel)
animate(tdir, sdir, title + "_3D_animation", rmdir=True)
def generate_attractor_animation(data, sdir, xlabel="", ylabel="", title=""):
sns.set()
tdir = sdir + "/tmp/"
mkdir_p(tdir)
if data.ndim == 1:
for n in range(data.shape[-1] - 1):
fig, ax1 = plt.subplots(1, 1)
ax1.plot(data[:n], np.roll(data, 1)[:n])
ax1.set_xlabel(xlabel)
ax1.set_ylabel(ylabel)
ax1.set_title(title)
fig.savefig(tdir + "{:04d}".format(n))
elif data.ndim == 2:
for m in range(data.shape[-1]):
for n in range(data.shape[1] - 1):
fig, ax1 = plt.subplots()
ax1.plot(data[:n, m], np.roll(data, 1)[:n, m])
ax1.set_xlabel(xlabel)
ax1.set_ylabel(ylabel)
ax1.set_title(title)
fig.savefig(tdir + "{:4d}{:4d}".format(m, n))
animate(tdir, sdir, title + "_attractor", rmdir=True)
def extract_animation(ii, mesh, fname, sdir, mode='train'):
tdir = sdir + "/tmp/"
mkdir_p(tdir)
if mode == 'all':
for train in range(ii.shape[-1]):
for pulse in range(ii.shape[-2]):
basic_plot(
ii[:, :, pulse, train],
mesh,
sdir=tdir + "train_{}_pulse_{}.png".format(train, pulse),
label="train/pulse: {}/{}".format(train + 1, pulse + 1))
if mode == 'train':
ii = mean_intensity(ii, mode='train')
for train in range(ii.shape[-1]):
basic_plot(ii[:, :, train],
mesh,
sdir=tdir + "train_{}.png".format(train),
label='train: {}'.format(train + 1))
animate(indir=tdir, outdir=sdir, fname=fname, delay=0.06)
def propagate_ensemble(ensemble,
bl,
sdir=None,
scale_input=None,
VERBOSE=True):
"""
propagate an ensemble of wavefronts (stored in hdf5 files) in ensemble_dir
down a desired beamline.
Useful in the case when there is a lot of ensembles
:param ensemble: list of strings pointing to hdf5 wavefront files
:param scale: bool or list of scaling parameters [nx, ny, fov] for scaling
of input wavefield (to be updated)
"""
for file in ensemble:
wfr = load_wfr(file)
if scale_input is not None and type(scale) == list:
scale(wfr, scale[0], scale[1], scale[2])
elif scale:
scale(wfr)
bl.propagate(wfr)
if sdir is not None:
mkdir_p(sdir)
wfr.store_hdf5(sdir + file)
def propagate_NVE():
wfr_directory = sys.argv[1].replace("*", "/")
job_name = "NKB_4980eV_250pC_NVE_to_EHC"
python_command = propagate_NVE
input_directory = dCache + "/NanoKB-Pulse/NVE/"
save_directory = input_directory.replace("/NVE/", "/EHC/")
mkdir_p(save_directory)
log_directory = logs
focus = "nano"
analysis = False
filename = __file__
dt = datetime.now().__str__()
function = python_command.__name__
description = "Propagate NanoKB Pulses 4.98 keV, 250 pC from Undulator NVE mirror to the EHC Screen"
append = None
print("info")
print("wavefront directory: {}".format(wfr_directory))
print("save directory: {}".format(save_directory))
print("focus (i.e. beamline option): {}".format(focus))
print("analysis: {}".format(analysis))
print("datetime: {}".format(dt))
print("filename: {}".format(filename))
print("function: {}".format(function))
print("description: {}".format(description))
wfr = Wavefront()
wfr.load_hdf5(wfr_directory)
bl = Beamline()
bl.append(Drift(2.2+3.5), propagation_parameters(1/3,1,1/3,1,'quadratic'))
bl.propagate(wfr)
wfr.custom_fields['focus'] = focus
wfr.custom_fields['job name'] = job_name
wfr.custom_fields['input directory'] = wfr_directory
wfr.custom_fields['datetime'] = dt
wfr.custom_fields['function'] = function
wfr.custom_fields['filename'] = filename
wfr.custom_fields['description'] = description
#wfr.custom_fields['bl'] = bl.__str__
if analysis:
wfr.analysis()
wfr.store_hdf5(wfr_directory.replace("/NVE/", "/EHC/"))
def propagate_from_ensemble(wfr_file,
ensemble_dir,
bl,
ekev,
sdir=None,
scale_input=None,
VERBOSE=True):
"""
propagate a single wavefont from an ensemble of wavefronts (stored in hdf5 files)
down a desired beamline. This script can be used in conjunction with jobscheduler
:param wfr_file: llocation of .hdf5 file
:param scale: bool or list of scaling parameters [nx, ny, fov] for scaling
of input wavefield (to be updated)
"""
if VERBOSE:
print("loading wavefront")
print("ensemble dir: {}".format(ensemble_dir))
print("file: {}".format(wfr_file))
wfr = load_wfr(ensemble_dir + "/" + wfr_file)
srwlib.srwl.SetRepresElecField(wfr._srwl_wf, 'f')
if VERBOSE:
print("wavefront loaded")
print(wfr.srw_info())
if scale_input is not None and type(scale) == list:
scale(wfr, scale[0], scale[1], scale[2])
elif scale:
wfr = scale(wfr)
if VERBOSE:
print("wavefront scaled")
if VERBOSE:
print("propagating wavefront")
print(bl)
bl.propagate(wfr)
if VERBOSE:
print("wavefront propagated")
if sdir is not None:
mkdir_p(sdir)
wfr.store_hdf5(sdir + wfr_file)
def propagation_batch_launcher(input_directory, sdir, focus, analysis = False, crop = None, append = None, descriptor = "", VERBOSE = True):
"""
This part launches the jobs that run in main
"""
wfr_directory = input_directory
mkdir_p(sdir)
cwd = os.getcwd()
script = __file__
filename = script
function = inspect.currentframe().f_code.co_name
js = JobScheduler(cwd + "/" + script, logDir = logs + "/",
jobName = "NVE_4.96keV", partition = 'exfel', nodes = 2, jobType = 'array',
jobArray = wfr_directory, options = [sdir, focus, analysis, crop, append, descriptor])
js.run(test = False)
def preprocess_shimadzu_data(proposal, exp, run, px):
sdir = dCache + "/NKB_sensing/whitefield/" + run + "/"
print("Saving results to dCache: {}".format(sdir))
mkdir_p(sdir)
print("Directory exists: {}".format(os.path.exists(sdir)))
data = load_data(run, proposal, exp)
ii = data.get_array('SPB_EHD_HPVX2_1/CAM/CAMERA:daqOutput',
'data.image.pixels')
ii = shimadzu_reshape(ii)
mesh = get_mesh(ii, px, px)
print("Mesh Shape: {}".format(mesh.shape))
print("")
print("Shimadzu Data Pre-Processed")
return ii, mesh
return arr
if __name__ == '__main__':
array_dir = dCache + "/whitefield_data/cropped_intensity_r0047.npy"
train_no = int(sys.argv[1])
method = [get_normalised_difference, get_second_order_doc]
if type(method) == list:
for m in method:
array = np.load(array_dir).astype('float64')[:, :, 3:103:2, :]
array = np.delete(array, 24, axis=-2)
map_loc = dCache + "/tmp/"
sdir = dCache + "/whitefield_data/r0047/"
mkdir_p(sdir)
sdir += "{}/".format(m.__name__)
mkdir_p(map_loc)
mkdir_p(sdir)
get_intensity_autocorrelation_train(train_no,
array,
map_loc=map_loc,
sdir=sdir,
method=m)
from felpy.utils.os_utils import mkdir_p
from felpy.utils.job_utils import JobScheduler
<<<<<<< HEAD
from wpg.wavefront import Wavefront
from labwork.about import logs, dCache
=======
from felpy.model.wavefront import Wavefront
>>>>>>> 108cfb9b6fc97d3841ee1db54862523eee5b184e
indir = dCache + "NanoKB-Pulse/source/"
outdir = dCache + "NanoKB-Pulse/source_data/"
mkdir_p(outdir)
tmp_dir = outdir + "/tmp/"
mkdir_p(tmp_dir)
intensity_dir = tmp_dir + "/integrated_intensity/"
mkdir_p(intensity_dir)
complex_dir = tmp_dir + "/complex_wavefield/"
mkdir_p(complex_dir)
def launch():
"""
This part launches the jobs that run in main
"""
def propagate_NVE():
wfr_directory = sys.argv[1].replace("*", "/")
job_name = "NKB_4980eV_250pC_Source_to_NVE"
python_command = propagate_NVE
input_directory = dCache + "/NanoKB-Pulse/source/"
save_directory = input_directory.replace("/source/", "/NVE/")
mkdir_p(save_directory)
log_directory = logs
focus = "nano"
analysis = False
filename = __file__
dt = datetime.now().__str__()
function = python_command.__name__
description = "Propagate NanoKB Pulses 4.98 keV, 250 pC from Undulator-exit to the NVE"
crop = 'NVE'
append = None
print("info")
print("wavefront directory: {}".format(wfr_directory))
print("save directory: {}".format(save_directory))
print("focus (i.e. beamline option): {}".format(focus))
print("analysis: {}".format(analysis))
print("datetime: {}".format(dt))
print("filename: {}".format(filename))
print("function: {}".format(function))
print("description: {}".format(description))
wfr = Wavefront()
wfr.load_hdf5(wfr_directory)
wfr.set_electric_field_representation('frequency')
wfr = scale(wfr)
print("wfr domain: {}".format(wfr.params.wDomain))
bl = get_beamline_object(ekev=4.96,
options=focus,
crop=crop,
theta_KB=5e-03,
theta_HOM=3.5e-03)
wfr.custom_fields['focus'] = focus
wfr.custom_fields['job name'] = job_name
wfr.custom_fields['input directory'] = wfr_directory
wfr.custom_fields['datetime'] = dt
wfr.custom_fields['function'] = function.__str__()
wfr.custom_fields['filename'] = filename
wfr.custom_fields['description'] = description
#wfr.custom_fields['bl'] = bl.__str__
bl.propagate(wfr)
if analysis:
wfr.analysis()
wfr.store_hdf5(wfr_directory.replace("/source/", "/NVE/"))
nph = get_energy_statistics(wfr)[1]
print("Mirror Angle: {} mrad Complete".format(ang*1e3
))
return (ang, nph)
>>>>>>> Stashed changes
MIRRORS = ['HOM1', 'HOM2', 'NHE', 'NVE']
<<<<<<< Updated upstream
if len(sys.argv) > 1:
ekev = float(sys.argv[1])
else:
ekev = 5.0
edir = "./{}keV/".format(ekev)
mkdir_p(edir)
for MIRROR in MIRRORS:
sdir = edir + "/{}/".format(MIRROR)
mkdir_p(sdir)
spb = Instrument(VERBOSE = False)
pool = mp.Pool(PROCESSES)
r = pool.map(partial(core, mirror_name = MIRROR, ekev = ekev),
tqdm(np.linspace(spb.params[MIRROR]['ang_min'],
spb.params[MIRROR]['ang_max'],
N)))
#animate(sdir, edir, "{}_mirror_rotation".format(MIRROR))
np.save(edir + "{}_mirror_flux_data".format(MIRROR), r)
def __init__(self,
python_command,
job_name,
log_directory,
partition='exfel',
nodes=1,
job_type='single',
job_array=None,
n_spawns=1,
VERBOSE=True,
runtime="14-00:00:00",
email="*****@*****.**",
mailtype="ALL",
options=None,
rundir=None):
"""
:param job_type: options = spawn, array, single
"""
if type(python_command) == str:
self.command = "script"
elif type(python_command) == type(random_string):
self.command = "method"
self.python_command = python_command
self.job_name = job_name
self.partition = partition
self.nodes = nodes
self.job_type = job_type
self.job_array = job_array
self.n_spawns = n_spawns
self.VERBOSE = VERBOSE
if self.VERBOSE:
print("\nInitialising Job Scheduler\n")
if rundir:
self.rundir = rundir
else:
self.rundir = os.getcwd()
self.runTime = runtime
self.log_directory = log_directory + job_name + "/"
self.jobDir = log_directory + "jobs/" + job_name + "/"
self.outDir = log_directory + "out/" + job_name + "/"
self.errDir = log_directory + "error/" + job_name + "/"
self.email = email
self.mailtype = mailtype
self.options = options
if os.path.exists(self.jobDir):
shutil.rmtree(self.jobDir)
mkdir_p(self.jobDir)
if os.path.exists(self.outDir):
shutil.rmtree(self.outDir)
mkdir_p(self.outDir)
if os.path.exists(self.errDir):
shutil.rmtree(self.errDir)
mkdir_p(self.errDir)
if self.VERBOSE == True:
self.__str__()
nodes = sys.argv[5]
else:
nodes = 2
mode = input("batch of single?")
if mode == 'single':
print("Launching Single Job")
launch_single(in_directory,
out_directory,
job_name,
logs,
nodes)
elif mode == 'batch':
print("Launching Batch Job")
mkdir_p(out_directory)
for item in os.listdir(in_directory):
launch_single(in_directory + "{}".format(item),
out_directory + "{}".format(item),
job_name,
logs,
nodes)
if DEBUG:
print(in_directory)
print(out_directory)
print(job_name)
print(logs)
print("nodes: ",nodes)
bl = get_beamline_object(ekev, apertures = True, surface = True,
crop = ["d1", "HOM1"], theta_HOM = angle)
bl.propagate(wfr)
return wfr.get_intensity().sum()
if __name__ == '__main__':
from labwork.about import dCache
from felpy.utils.os_utils import mkdir_p
try:
SDIR = "./mirror_reflectivity/"
mkdir_p(SDIR)
except(FileNotFoundError):
SDIR = input("Save Directory: ")
mkdir_p(SDIR)
#ii = no_mirror()
energies = [5.0, 7.0, 9.0, 11.0, 12.0]
angles = np.linspace(0, 10e-03, 35)
noap = memory_map(SDIR + "mirror_refl_no_aperture", shape = (len(energies), len(angles),2))
ap = memory_map(SDIR + "mirror_refl_aperture", shape = (len(energies), len(angles),2))
cpus = mpi.cpu_count()//2