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)