def plot_concentric_shells_spherical_coords(image_by_slices, base_folder, idx_slices=None): """Plot the concentric shells for a given three-dimensional volumetric shape. The volumetric shape is the three-dimensional diffraction intensity, as calculated by :py:mod:`ai4materials.descriptors.diffraction3d.Diffraction3D`. Parameters: image_by_slices: np.ndarray, shape [n_slices, theta_bins_fine, phi_bins_fine] Three-dimensional array containing each concentric shell obtained in spherical coordinate, as calculated by :py:mod:`ai4materials.descriptors.diffraction3d.Diffraction3D` ``n_slices``, ``theta_bins_fine``, ``phi_bins_fine`` are given by the interpolation and the region of the space considered. In our case, ``n_slices=52``, ``theta_bins_fine=256``, ``phi_bins_fine=512``, as defined in :py:mod:`ai4materials.descriptors.diffraction3d.Diffraction3D` in ``phi_bins_fine`` and ``theta_bins_fine``. base_folder: str Folder to save the figures generated. The figures are saved in a subfolder folder ``shells_png`` of ``base_folder``. idx_slices: list of int, optional (default=None) List of integers defining which concentric shells to plot. If `None`, all concentric shells - in spherical coordinates - are plotted. .. codeauthor:: Angelo Ziletti <*****@*****.**> """ if idx_slices is None: idx_slices = range(image_by_slices.shape[0]) # create folder for saving files shells_images_folder = os.path.join(base_folder, 'shells_png') if not os.path.exists(shells_images_folder): os.makedirs(shells_images_folder) filename_png_list = [] for idx_slice in idx_slices: filename_png = os.path.join( shells_images_folder, 'desc_sph_coords_slice' + str(idx_slice) + '.png') filename_png_list.append(filename_png) logger.debug("Slide idx: {}".format(idx_slice)) logger.debug("Image max: {}".format(image_by_slices[idx_slice].max())) coeffs = SHExpandDH(image_by_slices[idx_slice], sampling=2) coeffs_filtered = coeffs.copy() imgs = [ MakeGridDH(coeffs_filtered[:, :, :], sampling=2), MakeGridDH(coeffs_filtered[:, :16, :], sampling=2), MakeGridDH(coeffs_filtered[:, :32, :], sampling=2), MakeGridDH(coeffs_filtered[:, :64, :], sampling=2) ] fig, axes = plt.subplots(nrows=2, ncols=2) for idx_ax, ax in enumerate(axes.flat): im = ax.imshow(imgs[idx_ax], interpolation='none') fig.subplots_adjust(right=0.8) cbar_ax = fig.add_axes([0.85, 0.15, 0.05, 0.7]) fig.colorbar(im, cax=cbar_ax) plt.savefig(filename_png, dpi=100, format="png")
def calculate(self, structure, min_nb_atoms=20, plot_3d=False, plot_slices=False, plot_slices_sph_coords=False, save_diff_intensity=True, **kwargs): """Calculate the descriptor for the given ASE structure. Parameters: structure: `ase.Atoms` object Atomic structure. min_nb_atoms: int, optional (default=20) If the structure contains less than ``min_nb_atoms``, the descriptor is not calculated and an array with zeros is return as descriptor. This is because the descriptor is expected to be no longer meaningful for such a small amount of atoms present in the chosen structure. """ if len(structure) > min_nb_atoms - 1: atoms = scale_structure( structure, scaling_type=self.atoms_scaling, atoms_scaling_cutoffs=self.atoms_scaling_cutoffs, extrinsic_scale_factor=self.extrinsic_scale_factor) # Source src = condor.Source(**self.param_source) # Detector # solid_angle_correction are meaningless for 3d diffraction det = condor.Detector(solid_angle_correction=False, **self.param_detector) # Atoms atomic_numbers = map(lambda el: el.number, atoms) atomic_numbers = [ atomic_number + 5 for atomic_number in atomic_numbers ] # atomic_numbers = [82 for atomic_number in atomic_numbers] # convert Angstrom to m (CONDOR uses meters) atomic_positions = map( lambda pos: [pos.x * 1E-10, pos.y * 1E-10, pos.z * 1E-10], atoms) par = condor.ParticleAtoms(atomic_numbers=atomic_numbers, atomic_positions=atomic_positions) s = "particle_atoms" condor_exp = condor.Experiment(src, {s: par}, det) res = condor_exp.propagate3d() # retrieve some physical quantities that might be useful for users intensity = res["entry_1"]["data_1"]["data"] fourier_space = res["entry_1"]["data_1"]["data_fourier"] phases = np.angle(fourier_space) % (2 * np.pi) # 3D diffraction calculation real_space = np.fft.fftshift( np.fft.ifftn( np.fft.fftshift(res["entry_1"]["data_1"]["data_fourier"]))) window = get_window(self.window, self.n_px) tot_density = window * real_space.real center_of_mass = ndimage.measurements.center_of_mass(tot_density) logger.debug("Tot density data dimensions: {}".format( tot_density.shape)) logger.debug( "Center of mass of total density: {}".format(center_of_mass)) # take the fourier transform of structure in real_space fft_coeff = fftpack.fftn(tot_density, shape=(self.nx_fft, self.ny_fft, self.nz_fft)) # now shift the quadrants around so that low spatial frequencies are in # the center of the 2D fourier transformed image. fft_coeff_shifted = fftpack.fftshift(fft_coeff) # calculate a 3D power spectrum power_spect = np.abs(fft_coeff_shifted)**2 if self.use_mask: xc = (self.nx_fft - 1.0) / 2.0 yc = (self.ny_fft - 1.0) / 2.0 zc = (self.nz_fft - 1.0) / 2.0 # spherical mask a, b, c = xc, yc, zc x, y, z = np.ogrid[-a:self.nx_fft - a, -b:self.ny_fft - b, -c:self.nz_fft - c] mask_int = x * x + y * y + z * z <= self.mask_r_min * self.mask_r_min mask_out = x * x + y * y + z * z >= self.mask_r_max * self.mask_r_max for i in range(self.nx_fft): for j in range(self.ny_fft): for k in range(self.nz_fft): if mask_int[i, j, k]: power_spect[i, j, k] = 0.0 if mask_out[i, j, k]: power_spect[i, j, k] = 0.0 # cut the spectrum and keep only the relevant part for crystal-structure recognition of # hexagonal closed packed (spacegroup=194) # simple cubic (spacegroup=221) # face centered cubic (spacegroup=225) # diamond (spacegroup=227) # body centered cubic (spacegroup=229) # this interval (20:108) might need to be varied if other classes are added power_spect_cut = power_spect[20:108, 20:108, 20:108] # zoom by two times using spline interpolation power_spect = ndimage.zoom(power_spect_cut, (2, 2, 2)) if save_diff_intensity: np.save( '/home/ziletti/Documents/calc_nomadml/rot_inv_3d/power_spect.npy', power_spect) # power_spect.shape = 176, 176, 176 if plot_3d: plot_3d_volume(power_spect) vox = np.copy(power_spect) logger.debug("nan in data: {}".format( np.count_nonzero(~np.isnan(vox)))) # optimized # these specifications are valid for a power_spect = power_spect[20:108, 20:108, 20:108] # and a magnification of 2 xyz_indices_r = get_slice_volume_indices( vox, min_r=32.0, dr=1.0, max_r=83., phi_bins=self.phi_bins, theta_bins=self.theta_bins) # slow - only for benchmarking the fast implementation below (shells_to_sph, interp_theta_phi_surfaces) # (vox_by_slices, theta_phi_by_slices) = _slice_3d_volume_slow(vox) # convert 3d shells (vox_by_slices, theta_phi_by_slices) = get_shells_from_indices( xyz_indices_r, vox) if plot_slices: plot_concentric_shells( vox_by_slices, base_folder=self.configs['io']['main_folder'], idx_slices=None, create_animation=False) image_by_slices = interp_theta_phi_surfaces( theta_phi_by_slices, theta_bins=self.theta_bins_fine, phi_bins=self.phi_bins_fine) if plot_slices_sph_coords: plot_concentric_shells_spherical_coords( image_by_slices, base_folder=self.configs['io']['main_folder'], idx_slices=None) coeffs_list = [] nl_list = [] ls_list = [] for idx_slice in range(image_by_slices.shape[0]): logger.debug("img #{} max: {}".format( idx_slice, image_by_slices[idx_slice].max())) # set to zero the spherical harmonics coefficients above self.sph_l_cutoff coeffs = SHExpandDH(image_by_slices[idx_slice], sampling=2) coeffs_filtered = coeffs.copy() coeffs_filtered[:, self.sph_l_cutoff:, :] = 0. coeffs = coeffs_filtered.copy() nl = coeffs.shape[0] ls = np.arange(nl) coeffs_list.append(coeffs) nl_list.append(nl) ls_list.append(ls) coeffs = np.asarray(coeffs_list).reshape(image_by_slices.shape[0], coeffs.shape[0], coeffs.shape[1], coeffs.shape[2]) sh_coeffs_list = [] for idx_slice in range(coeffs.shape[0]): sh_coeffs = SHCoeffs.from_array(coeffs[idx_slice]) sh_coeffs_list.append(sh_coeffs) sh_spectrum_list = [] for sh_coeff in sh_coeffs_list: sh_spectrum = sh_coeff.spectrum(convention='l2norm') sh_spectrum_list.append(sh_spectrum) sh_spectra = np.asarray(sh_spectrum_list).reshape( coeffs.shape[0], -1) # cut the spherical harmonics expansion to sph_l_cutoff order logger.debug( 'Spherical harmonics spectra maximum before normalization: {}'. format(sh_spectra.max())) sh_spectra = sh_spectra[:, :self.sph_l_cutoff] sh_spectra = (sh_spectra - sh_spectra.min()) / (sh_spectra.max() - sh_spectra.min()) # add results in ASE structure info descriptor_data = dict(descriptor_name=self.name, descriptor_info=str(self), diffraction_3d_sh_spectrum=sh_spectra) else: # return array with zeros for structures with less than min_nb_atoms sh_spectra = np.zeros((52, int(self.sph_l_cutoff))) descriptor_data = dict(descriptor_name=self.name, descriptor_info=str(self), diffraction_3d_sh_spectrum=sh_spectra) structure.info['descriptor'] = descriptor_data return structure