def get_intensity_autocorrelation_train(train_no,
array,
map_loc,
method=get_normalised_difference,
sdir=None,
mpi=True):
array = array[:, :, 3:103, train_no]
map_loc = map_loc + "autocorrelation_map_{}".format(train_no)
mmap = memory_map(map_loc, shape=(*array.shape, array.shape[-1]))
pool = mp.Pool(mp.cpu_count() // 2)
func = partial(get_intensity_autocorrelation_pulse,
array=array,
map_loc=map_loc,
method=method)
pool.map(func, range(array.shape[-1]))
g = readMap(map_loc, (*array.shape, array.shape[-1]))
del mmap
if sdir is not None:
np.save(sdir + "autocorrelation_{}".format(train_no), g)
del g
print("Completed Autocorrelation Calculations for Train: {}".format(
train_no))
def get_coherence_time(cfr,
tStep,
mpi=False,
map_loc="/tmp/coherence_map",
bins=1,
VERBOSE=True):
"""
Calculate the coherence time of complex wavefield of shape
[nx, ny, nt].
Relevant for statistically stationary sources.
ref: Coherence properties of the radiation from X-ray free electron laser
:param cfr: complex wavefield
:param tstep: temporal step between slices
:returns tau: coherence time [s]
"""
mmap = memory_map(map_loc=map_loc, shape=cfr.shape, dtype='complex64')
nz0 = cfr.shape[-1]
if bins == 1:
nz1 = nz0
else:
cfr = binArray(cfr, axis=-1, binstep=nz0 // bins, binsize=1)
nz1 = cfr.shape[-1]
tStep *= (nz0 / nz1)
g = np.zeros([*cfr.shape], dtype='complex64')
if VERBOSE:
print("Calculating Coherence Time")
if mpi:
processes = mp.cpu_count() // 2
pool = mp.Pool(processes)
pool.map(partial(get_longitudinal_coherence, cfr=cfr, map_loc=map_loc),
range(cfr.shape[-1]))
g = readMap(map_loc, cfr.shape, dtype='complex64')
else:
for i in tqdm(range(cfr.shape[-1])):
g[:, :, i] = get_longitudinal_coherence(slice_no=i, cfr=cfr)
tau = (abs(g)**2).sum(axis=-1)[0, 0]
print("g", np.max(g))
if VERBOSE:
print("\n")
print(tau)
print("Time Step: {} fs".format(tStep * 1e15))
print("Coherence Time: {:.2e} fs".format(tau * 1e15 * tStep))
del mmap
os.remove(map_loc)
return tau * tStep, g
def get_longitudinal_coherence_new(slice_no,
cfr,
map_loc=None,
bins=1,
VERBOSE=True):
"""
Calculate the longitudinal correlation of each slice of a complex wavefront
of shape [nx, ny, nz] against a single slice of shape [nx,ny] at longitudinal
interval defined by the slice_no
:param cfr: complex wavefield
:param slice_no: longitudinal index [int]
:returns g: complex degree of coherence
"""
A = np.roll(cfr, -slice_no, axis=2)
B = np.repeat(cfr[:, :, slice_no][:, :, np.newaxis],
cfr.shape[-1],
axis=-1)
## DEGUB print(A[:,:,0] == wfr[:,:,i])
## DEBUG print([B[:,:,k] == wfr[:,:,i] for k in range(wfr.shape[-1])])
if map_loc is not None:
mmap = memory_map(map_loc, shape=cfr.shape, dtype='complex64')
mmap[:, :, slice_no] = ((A * B.conjugate()).mean(axis=-1)) / np.sqrt(
(abs(A)**2).mean(axis=-1) * (abs(B)**2).mean(axis=-1))
else:
return ((A * B.conjugate()).mean(axis=-1)) / np.sqrt(
(abs(A)**2).mean(axis=-1) * (abs(B)**2).mean(axis=-1))
def get_intensity_autocorrelation_pulse(pulse_no,
array,
map_loc,
method=get_normalised_difference):
"""
computes the time-lagged intensity-intensity autocorrelation for a single
slice in time (as denoted by slice_no)
"""
### create a memory map for this pulse-train
### send each pulse to its own processor
mmap = memory_map(map_loc, shape=(*array.shape, array.shape[-1]))
pulse = np.repeat(array[:, :, pulse_no][:, :, np.newaxis],
array.shape[-1],
axis=-1)
mmap[:, :, pulse_no, :] = method(pulse, array)
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
print("Distributing to {} cpus".format(cpus))
for a in tqdm(range(len(energies))):
pool = mpi.Pool(processes = cpus)
#ii = no_mirror(energy)
func_a = partial(no_aperture, ekev = energies[a])
res = pool.map(func_a, angles)
noap[a, :, :] = np.array([angles, res]).T