def __init__(self, experiment, diffraction_rings=None, log_scale=False):
self.experiment = experiment
self.center = self.experiment.det.geometry.point_coord_indexes((0, 0))
# ! Wavenumber definition != beam's.
self.wavenumber = np.linalg.norm(self.experiment.beam.get_wavevector())
self.distance = self.experiment.det.distance
pixel_width = asnumpy(self.experiment.det.pixel_width)
# Cupy doesn't have median yet.
self.pix_width = np.median(pixel_width)
recidet = ReciprocalDetector(self.experiment.det, self.experiment.beam)
self.q_max = asnumpy(
np.min(( # Max inscribed radius
xp.max(recidet.pixel_position_reciprocal[..., 1]),
-xp.min(recidet.pixel_position_reciprocal[..., 1]),
xp.max(recidet.pixel_position_reciprocal[..., 0]),
-xp.min(recidet.pixel_position_reciprocal[..., 0]))))
self._auto_rings = False
if diffraction_rings is not None:
if diffraction_rings.lower() == "auto":
self._auto_rings = True
else:
raise ValueError(
"Unrecognized value '{}' for argument diffraction_rings"
"".format(diffraction_rings))
if log_scale:
self._norm = LogNorm()
else:
self._norm = None
def calculate_atomic_factor(particle, q_space, pixel_num):
"""
Calculate the atomic form factor for each atom at each momentum
:param particle: The particle object
:param q_space: The reciprocal to calculate
:param pixel_num: The number of pixels.
:return:
"""
q_space = asnumpy(q_space) # CubicSpline is not compatible with Cupy
f_hkl = np.zeros((particle.num_atom_types, pixel_num))
q_space_1d = np.reshape(q_space, [
pixel_num,
])
if particle.num_atom_types == 1:
cs = CubicSpline(particle.q_sample,
particle.ff_table[:]) # Use cubic spline
f_hkl[0, :] = cs(q_space_1d) # interpolate
else:
for atm in range(particle.num_atom_types):
cs = CubicSpline(particle.q_sample,
particle.ff_table[atm, :]) # Use cubic spline
f_hkl[atm, :] = cs(q_space_1d) # interpolate
f_hkl = np.reshape(f_hkl, (particle.num_atom_types, ) + q_space.shape)
return xp.asarray(f_hkl)
def imshow(self, img):
img = asnumpy(img)
plt.imshow(img, norm=self._norm)
plt.colorbar()
plt.xlabel('Y')
plt.ylabel('X')
if self._auto_rings:
self.add_diffraction_rings()
def initialize(self, geom, beam):
"""
Initialize the detector with user defined parameters
:param geom: The dictionary containing all the necessary information to initialized the
detector.
:param beam: The beam object
:return: None
"""
"""
Doc:
To use this class, the user has to provide the necessary information to initialize
the detector.
All the necessary entries are listed in the example notebook.
"""
# 'detector distance': detector distance in (m)
##########################################################################################
# Extract necessary information
##########################################################################################
# Define the hierarchy system. For simplicity, we only use two-layer structure.
for key in {'panel number', 'panel pixel num x', 'panel pixel num y'}:
if key not in geom:
raise KeyError("Missing required '{}' key.".format(key))
self.panel_num = int(geom['panel number'])
self.panel_pixel_num_x = int(geom['panel pixel num x'])
self.panel_pixel_num_y = int(geom['panel pixel num y'])
self._shape = (self.panel_num, self.panel_pixel_num_x,
self.panel_pixel_num_y)
# Define all properties the detector should have
self._distance = None
if 'pixel center z' in geom:
if 'detector distance' in geom:
raise ValueError("Please provide one of "
"'pixel center z' or 'detector distance'.")
self.center_z = xp.asarray(geom['pixel center z'],
dtype=xp.float64)
self._distance = float(self.center_z.mean())
else:
if 'detector distance' not in geom:
KeyError("Missing required 'detector distance' key.")
self._distance = float(geom['detector distance'])
self.center_z = self._distance * xp.ones(self._shape,
dtype=xp.float64)
# Below: [panel number, pixel num x, pixel num y] in (m)
# Change dtype and make numpy/cupy array
self.pixel_width = xp.asarray(geom['pixel width'], dtype=xp.float64)
self.pixel_height = xp.asarray(geom['pixel height'], dtype=xp.float64)
self.center_x = xp.asarray(geom['pixel center x'], dtype=xp.float64)
self.center_y = xp.asarray(geom['pixel center y'], dtype=xp.float64)
self.orientation = np.array([0, 0, 1])
# construct the the pixel position array
self.pixel_position = xp.zeros(self._shape + (3, ))
self.pixel_position[..., 0] = self.center_x
self.pixel_position[..., 1] = self.center_y
self.pixel_position[..., 2] = self.center_z
# Pixel map
if 'pixel map' in geom:
# [panel number, pixel num x, pixel num y]
self.pixel_index_map = xp.asarray(geom['pixel map'],
dtype=xp.int64)
# Detector pixel number info
self.detector_pixel_num_x = asnumpy(
xp.max(self.pixel_index_map[..., 0]) + 1)
self.detector_pixel_num_y = asnumpy(
xp.max(self.pixel_index_map[..., 1]) + 1)
# Panel pixel number info
# number of pixels in each panel in x/y direction
self.panel_pixel_num_x = self.pixel_index_map.shape[1]
self.panel_pixel_num_y = self.pixel_index_map.shape[2]
# total number of pixels (px*py)
self.pixel_num_total = np.prod(self._shape)
###########################################################################################
# Do necessary calculation to finishes the initialization
###########################################################################################
# self.geometry currently only work for the pre-defined detectors
self.geometry = geom
# Calculate the pixel area
self.pixel_area = xp.multiply(self.pixel_height, self.pixel_width)
# Get reciprocal space configurations and corrections.
self.initialize_pixels_with_beam(beam=beam)
##########################################################################################
# Do necessary calculation to finishes the initialization
##########################################################################################
# Detector effects
if 'pedestal' in geom:
self._pedestal = xp.asarray(geom['pedestal'], dtype=xp.float64)
else:
self._pedestal = xp.zeros((self.panel_num, self.panel_pixel_num_x,
self.panel_pixel_num_y))
if 'pixel rms' in geom:
self._pixel_rms = xp.asarray(geom['pixel rms'], dtype=xp.float64)
else:
self._pixel_rms = xp.zeros((self.panel_num, self.panel_pixel_num_x,
self.panel_pixel_num_y))
if 'pixel bkgd' in geom:
self._pixel_bkgd = xp.asarray(geom['pixel bkgd'], dtype=xp.float64)
else:
self._pixel_bkgd = xp.zeros(
(self.panel_num, self.panel_pixel_num_x,
self.panel_pixel_num_y))
if 'pixel status' in geom:
self._pixel_status = xp.asarray(geom['pixel status'],
dtype=xp.float64)
else:
self._pixel_status = xp.zeros(
(self.panel_num, self.panel_pixel_num_x,
self.panel_pixel_num_y))
if 'pixel mask' in geom:
self._pixel_mask = xp.asarray(geom['pixel mask'], dtype=xp.float64)
else:
self._pixel_mask = xp.zeros(
(self.panel_num, self.panel_pixel_num_x,
self.panel_pixel_num_y))
if 'pixel gain' in geom:
self._pixel_gain = xp.asarray(geom['pixel gain'], dtype=xp.float64)
else:
self._pixel_gain = xp.ones((self.panel_num, self.panel_pixel_num_x,
self.panel_pixel_num_y))
def initialize(self, geom, run_num=0, cframe=0):
"""
Initialize the detector
:param geom: The *-end.data file which characterizes the geometry profile.
:param run_num: The run_num containing the background, rms and gain and the other
pixel pixel properties.
:param cframe: The desired coordinate frame, 0 for psana and 1 for lab conventions.
:return: None
"""
# Redirect the output stream
old_stdout = sys.stdout
f = six.StringIO()
# f = open('Detector_initialization.log', 'w')
sys.stdout = f
###########################################################################################
# Initialize the geometry configuration
############################################################################################
self.geometry = GeometryAccess(geom, cframe=cframe)
self.run_num = run_num
# Set coordinate in real space (convert to m)
temp = [
xp.asarray(t) * 1e-6
for t in self.geometry.get_pixel_coords(cframe=cframe)
]
temp_index = [
xp.asarray(t)
for t in self.geometry.get_pixel_coord_indexes(cframe=cframe)
]
self.panel_num = np.prod(temp[0].shape[:-2])
self._distance = float(temp[2].mean())
self._shape = (self.panel_num, temp[0].shape[-2], temp[0].shape[-1])
self.pixel_position = xp.zeros(self._shape + (3, ))
self.pixel_index_map = xp.zeros(self._shape + (2, ))
for n in range(3):
self.pixel_position[..., n] = temp[n].reshape(self._shape)
for n in range(2):
self.pixel_index_map[..., n] = temp_index[n].reshape(self._shape)
self.pixel_index_map = self.pixel_index_map.astype(xp.int64)
# Get the range of the pixel index
self.detector_pixel_num_x = asnumpy(
xp.max(self.pixel_index_map[..., 0]) + 1)
self.detector_pixel_num_y = asnumpy(
xp.max(self.pixel_index_map[..., 1]) + 1)
self.panel_pixel_num_x = np.array([
self.pixel_index_map.shape[1],
] * self.panel_num)
self.panel_pixel_num_y = np.array([
self.pixel_index_map.shape[2],
] * self.panel_num)
self.pixel_num_total = np.sum(
np.multiply(self.panel_pixel_num_x, self.panel_pixel_num_y))
tmp = float(self.geometry.get_pixel_scale_size() *
1e-6) # Convert to m
self.pixel_width = xp.ones((self.panel_num, self.panel_pixel_num_x[0],
self.panel_pixel_num_y[0])) * tmp
self.pixel_height = xp.ones((self.panel_num, self.panel_pixel_num_x[0],
self.panel_pixel_num_y[0])) * tmp
# Calculate the pixel area
self.pixel_area = xp.multiply(self.pixel_height, self.pixel_width)
###########################################################################################
# Initialize the pixel effects
###########################################################################################
# first we should parse the path
parsed_path = geom.split('/')
self.exp = parsed_path[-5]
if self.exp == 'calib':
self.exp = parsed_path[-6]
self.group = parsed_path[-4]
self.source = parsed_path[-3]
self._pedestals = None
self._pixel_rms = None
self._pixel_mask = None
self._pixel_bkgd = None
self._pixel_status = None
self._pixel_gain = None
if six.PY2:
try:
cbase = self._get_cbase()
self.calibdir = '/'.join(parsed_path[:-4])
pbits = 255
gcp = GenericCalibPars(cbase, self.calibdir, self.group,
self.source, run_num, pbits)
self._pedestals = gcp.pedestals()
self._pixel_rms = gcp.pixel_rms()
self._pixel_mask = gcp.pixel_mask()
self._pixel_bkgd = gcp.pixel_bkgd()
self._pixel_status = gcp.pixel_status()
self._pixel_gain = gcp.pixel_gain()
except NotImplementedError:
# No GenericCalibPars information.
pass
else:
try:
self.det = self._get_det_id(self.group)
except NotImplementedError:
# No GenericCalibPars information.
self.det = None
# Redirect the output stream
sys.stdout = old_stdout
def main():
# Parse user input for config file and dataset name
user_input = parse_input_arguments(sys.argv)
config_file = user_input['config']
dataset_name = user_input['dataset']
# Get the Config file parameters
with open(config_file) as config_file:
config_params = json.load(config_file)
# Check if dataset in Config file
if dataset_name not in config_params:
raise Exception("Dataset {} not in Config file.".format(dataset_name))
# Get the dataset parameters from Config file parameters
dataset_params = config_params[dataset_name]
# Get the input dataset parameters
pdb_file = dataset_params["pdb"]
beam_file = dataset_params["beam"]
beam_fluence_increase_factor = dataset_params["beamFluenceIncreaseFactor"]
geom_file = dataset_params["geom"]
dataset_size = dataset_params["numPatterns"]
# Divide up the task of creating the dataset to be executed simultaneously by multiple ranks
batch_size = dataset_params["batchSize"]
# Get the output dataset parameters
img_dir = dataset_params["imgDir"]
output_dir = dataset_params["outDir"]
# raise exception if batch_size does not divide into dataset_size
if dataset_size % batch_size != 0:
if RANK == MASTER_RANK:
raise ValueError(
"(Master) batch_size {} should divide dataset_size {}.".format(
batch_size, dataset_size))
else:
sys.exit(1)
# Compute number of batches to process
n_batches = dataset_size // batch_size
# Flags
save_volume = False
with_intensities = False
given_orientations = True
# Constants
photons_dtype = np.uint8
photons_max = np.iinfo(photons_dtype).max
# Load beam parameters
beam = ps.Beam(beam_file)
# Increase the beam fluence
if not np.isclose(beam_fluence_increase_factor, 1.0):
beam.set_photons_per_pulse(beam_fluence_increase_factor *
beam.get_photons_per_pulse())
# Load geometry of detector
det = ps.PnccdDetector(geom=geom_file, beam=beam)
# Get the shape of the diffraction pattern
diffraction_pattern_height = det.detector_pixel_num_x.item()
diffraction_pattern_width = det.detector_pixel_num_y.item()
# Define path to output HDF5 file
output_file = get_output_file_name(dataset_name, dataset_size,
diffraction_pattern_height,
diffraction_pattern_width)
cspi_synthetic_dataset_file = os.path.join(output_dir, output_file)
# Generate uniform orientations
if given_orientations and RANK == MASTER_RANK:
print("(Master) Generate {} uniform orientations".format(dataset_size))
orientations = ps.get_uniform_quat(dataset_size, True)
sys.stdout.flush()
# Create a particle object
if RANK == GPU_RANKS[0]:
# Load PDB
print("(GPU 0) Reading PDB file: {}".format(pdb_file))
particle = ps.Particle()
particle.read_pdb(pdb_file, ff='WK')
# Calculate diffraction volume
print("(GPU 0) Calculating diffraction volume")
experiment = ps.SPIExperiment(det, beam, particle)
else:
experiment = ps.SPIExperiment(det, beam, None)
sys.stdout.flush()
# Transfer diffraction volume to CPU memory
buffer = asnumpy(experiment.volumes[0])
# GPU rank broadcasts diffraction volume to other ranks
COMM.Bcast(buffer, root=1)
# This condition is necessary if the script is run on more than one machine (each machine having 1 GPU and 9 CPU)
if RANK in GPU_RANKS[1:]:
experiment.volumes[0] = xp.asarray(experiment.volumes[0])
if RANK == MASTER_RANK:
# Create output directory if it does not exist
if not os.path.exists(output_dir):
print("(Master) Creating output directory: {}".format(output_dir))
os.makedirs(output_dir)
# Create image directory if it does not exist
if not os.path.exists(img_dir):
print(
"(Master) Creating image output directory: {}".format(img_dir))
os.makedirs(img_dir)
print("(Master) Creating HDF5 file to store the datasets: {}".format(
cspi_synthetic_dataset_file))
f = h5.File(cspi_synthetic_dataset_file, "w")
f.create_dataset("pixel_position_reciprocal",
data=det.pixel_position_reciprocal)
f.create_dataset("pixel_index_map", data=det.pixel_index_map)
if given_orientations:
f.create_dataset("orientations", data=orientations)
f.create_dataset("photons", (dataset_size, 4, 512, 512), photons_dtype)
# Create a dataset to store the diffraction patterns
f.create_dataset("diffraction_patterns",
(dataset_size, diffraction_pattern_height,
diffraction_pattern_width),
dtype='f')
if save_volume:
f.create_dataset("volume", data=experiment.volumes[0])
if with_intensities:
f.create_dataset("intensities", (dataset_size, 4, 512, 512),
np.float32)
f.close()
sys.stdout.flush()
# Make sure file is created before others open it
COMM.barrier()
# Add the atomic coordinates of the particle to the HDF5 file
if RANK == GPU_RANKS[0]:
atomic_coordinates = particle.atom_pos
f = h5.File(cspi_synthetic_dataset_file, "a")
dset_atomic_coordinates = f.create_dataset("atomic_coordinates",
atomic_coordinates.shape,
dtype='f')
dset_atomic_coordinates[...] = atomic_coordinates
f.close()
# Make sure file is closed before others open it
COMM.barrier()
# Keep track of the number of images processed
n_images_processed = 0
if RANK == MASTER_RANK:
# Send batch numbers to non-Master ranks
for batch_n in tqdm(range(n_batches)):
# Receive query for batch number from a rank
i_rank = COMM.recv(source=MPI.ANY_SOURCE)
# Send batch number to that rank
COMM.send(batch_n, dest=i_rank)
# Send orientations as well
if given_orientations:
batch_start = batch_n * batch_size
batch_end = (batch_n + 1) * batch_size
COMM.send(orientations[batch_start:batch_end], dest=i_rank)
# Tell non-Master ranks to stop asking for more data since there are no more batches to process
for _ in range(N_RANKS - 1):
# Send one "None" to each rank as final flag
i_rank = COMM.recv(source=MPI.ANY_SOURCE)
COMM.send(None, dest=i_rank)
else:
# Get the HDF5 file
f = h5.File(cspi_synthetic_dataset_file, "r+")
# Get the dataset used to store the photons
h5_photons = f["photons"]
# Get the dataset used to store the diffraction patterns
h5_diffraction_patterns = f["diffraction_patterns"]
# Get the dataset used to store intensities
if with_intensities:
h5_intensities = f["intensities"]
while True:
# Ask for batch number from Master rank
COMM.send(RANK, dest=MASTER_RANK)
# Receive batch number from Master rank
batch_n = COMM.recv(source=MASTER_RANK)
# If batch number is final flag, stop
if batch_n is None:
break
# Receive orientations as well from Master rank
if given_orientations:
orientations = COMM.recv(source=MASTER_RANK)
experiment.set_orientations(orientations)
# Define a Numpy array to hold a batch of photons
np_photons = np.zeros((batch_size, 4, 512, 512), photons_dtype)
# Define a Numpy array to hold a batch of diffraction patterns
np_diffraction_patterns = np.zeros(
(batch_size, diffraction_pattern_height,
diffraction_pattern_width))
# Define a Numpy array to hold a batch of intensities
if with_intensities:
np_intensities = np.zeros((batch_size, 4, 512, 512),
np.float32)
# Define the batch start and end offsets
batch_start = batch_n * batch_size
batch_end = (batch_n + 1) * batch_size
# Generate batch of snapshots
for i in range(batch_size):
# Generate image stack
image_stack_tuple = experiment.generate_image_stack(
return_photons=True,
return_intensities=with_intensities,
always_tuple=True)
# Photons
photons = image_stack_tuple[0]
# # Raise exception if photon max exceeds max of uint8
# if photons.max() > photons_max:
# raise RuntimeError("Value of photons too large for type {}.".format(photons_dtype))
np_photons[i] = asnumpy(photons.astype(photons_dtype))
# Assemble the image stack into a 2D diffraction pattern
np_diffraction_pattern = experiment.det.assemble_image_stack(
image_stack_tuple)
# Add the assembled diffraction pattern to the batch
np_diffraction_patterns[i] = np_diffraction_pattern
# Save diffraction pattern as PNG file
data_index = batch_start + i
save_diffraction_pattern_as_image(data_index, img_dir,
np_diffraction_pattern)
# Intensities
if with_intensities:
np_intensities[i] = asnumpy(image_stack_tuple[1].astype(
np.float32))
# Update the number of images processed
n_images_processed += 1
# Add the batch of photons to the HDF5 file
h5_photons[batch_start:batch_end] = np_photons
# Add the batch of diffraction patterns to the HDF5 file
h5_diffraction_patterns[
batch_start:batch_end] = np_diffraction_patterns
if with_intensities:
h5_intensities[batch_start:batch_end] = np_intensities
# Close the HDF5 file
f.close()
sys.stdout.flush()
# Wait for ranks to finish
COMM.barrier()
def __init__(self, N_pixel, det_size, det_distance, beam=None):
"""
Initialize the detector
:param N_pixel: Number of pixels per dimension.
:param det_size: Length of detector sides (m).
:param det_distance: Sample-Detector distance (m).
"""
super(SimpleSquareDetector, self).__init__()
ar = xp.arange(N_pixel)
sep = float(det_size) / N_pixel
x = -det_size / 2 + sep / 2 + ar * sep
y = -det_size / 2 + sep / 2 + ar * sep
X, Y = xp.meshgrid(x, y, indexing='xy')
Xar, Yar = xp.meshgrid(ar, ar, indexing='xy')
self.panel_num = 1
self.panel_pixel_num_x = N_pixel
self.panel_pixel_num_y = N_pixel
self._shape = (1, N_pixel, N_pixel)
# Define all properties the detector should have
self._distance = det_distance
self.center_z = self._distance * xp.ones(self._shape, dtype=xp.float64)
p_center_x = xp.stack((X, ))
p_center_y = xp.stack((Y, ))
self.pixel_width = sep * xp.ones(self._shape, dtype=xp.float64)
self.pixel_height = sep * xp.ones(self._shape, dtype=xp.float64)
self.center_x = p_center_x
self.center_y = p_center_y
# construct the the pixel position array
self.pixel_position = xp.zeros(self._shape + (3, ))
self.pixel_position[..., 0] = self.center_x
self.pixel_position[..., 1] = self.center_y
self.pixel_position[..., 2] = self.center_z
# Pixel map
p_map_x = xp.stack((Xar, ))
p_map_y = xp.stack((Yar, ))
# [panel number, pixel num x, pixel num y]
self.pixel_index_map = xp.stack((p_map_x, p_map_y), axis=-1)
# Detector pixel number info
self.detector_pixel_num_x = asnumpy(
xp.max(self.pixel_index_map[..., 0]) + 1)
self.detector_pixel_num_y = asnumpy(
xp.max(self.pixel_index_map[..., 1]) + 1)
# Panel pixel number info
# number of pixels in each panel in x/y direction
self.panel_pixel_num_x = self.pixel_index_map.shape[1]
self.panel_pixel_num_y = self.pixel_index_map.shape[2]
# total number of pixels (px*py)
self.pixel_num_total = np.prod(self._shape)
# Calculate the pixel area
self.pixel_area = xp.multiply(self.pixel_height, self.pixel_width)
# Get reciprocal space configurations and corrections.
self.initialize_pixels_with_beam(beam=beam)
