def __init__(self, dev, selfAb, lac, grid_concentration, p, n_element, sample_height_n, minibatch_size, sample_size_n, sample_size_cm, this_aN_dic, probe_energy, probe_cts, theta_st, theta_end, n_theta, this_theta_idx, n_det, P_minibatch, det_size_cm, det_from_sample_cm): """ Initialize the attributes of PPM. """ super(PPM, self).__init__() # inherit the __init__ from nn.Module. self.dev = dev self.selfAb = selfAb self.lac = lac.to(self.dev) self.grid_concentration = grid_concentration self.theta_ls = -tc.linspace(theta_st, theta_end, n_theta + 1)[:-1] self.this_theta_idx = this_theta_idx self.n_element = n_element self.sample_height_n = sample_height_n self.minibatch_size = minibatch_size self.sample_size_n = sample_size_n self.p = p # indicate which minibatch to calculate the gradient self.xp = self.init_xp() # initialize the values of the minibatch self.probe_energy = probe_energy self.this_aN_dic = this_aN_dic self.n_element = tc.as_tensor(len(self.this_aN_dic)).to(self.dev) self.element_ls = np.array(list(this_aN_dic.keys())) self.aN_ls = np.array(list(this_aN_dic.values())) self.probe_attCS_ls = tc.as_tensor( xlib_np.CS_Total(self.aN_ls, self.probe_energy).flatten()).to(self.dev) self.sample_size_cm = sample_size_cm self.fl_all_lines_dic = self.init_fl_all_lines_dic() self.n_lines = tc.as_tensor(self.fl_all_lines_dic["n_lines"]).to( self.dev) self.FL_line_attCS_ls = tc.as_tensor( xlib_np.CS_Total(self.aN_ls, self.fl_all_lines_dic["fl_energy"])).float().to( self.dev) self.detected_fl_unit_concentration = tc.as_tensor( self.fl_all_lines_dic["detected_fl_unit_concentration"]).float( ).to(self.dev) self.n_line_group_each_element = tc.IntTensor( self.fl_all_lines_dic["n_line_group_each_element"]).to(self.dev) self.dia_len_n = int((self.sample_height_n**2 + self.sample_size_n**2 + self.sample_size_n**2)**0.5) self.n_voxel_batch = self.minibatch_size * self.sample_size_n self.n_voxel = self.sample_height_n * self.sample_size_n**2 self.n_det = n_det self.P_minibatch = P_minibatch.to(dev) self.det_size_cm = det_size_cm self.det_from_sample_cm = det_from_sample_cm self.SA_theta = self.init_SA_theta() self.probe_cts = probe_cts self.probe_before_attenuation_flat = self.init_probe()
def create_phantom_tooth(size_x, energy, elem_tooth, elem_implant, pixel_size, isimplant): """ Create phantom, share is two flowers (with 3 and 6 petals) :param sx: size of phantom :param energy: energy, set in keV :param elem1: Number of the chemical element :param elem2: Number of the chemical element :param pixel_size: size of one pixel, set in microns :return: 2d array of phantom """ xrl_np.XRayInit() phantom = np.zeros((size_x, size_x)) sx_half = size_x / 2 sq = size_x / 14 #calculate mu density_tooth = xrl_np.ElementDensity(np.array([elem_tooth])) cross_section_tooth = xrl_np.CS_Total(np.array([elem_tooth]), np.array([energy])) mu_tooth = density_tooth * cross_section_tooth density_implant = xrl_np.ElementDensity(np.array([elem_implant])) cross_section_implant = xrl_np.CS_Total(np.array([elem_implant]), np.array([energy])) mu_implant = density_implant * cross_section_implant #buld mesh y, x = np.meshgrid(range(size_x), range(size_x)) xx = (x - sx_half).astype('float32') yy = (y - sx_half).astype('float32') r = np.sqrt(xx * xx + yy * yy) tetta = np.arctan2(yy, xx) #make teeth mask_tooth = r <= sq * (1 + np.cos(2 * tetta) + np.sin(2 * tetta)**2) mask_tooth += (xx * xx + yy * yy) <= (0.09 * size_x)**2 mask_tooth += np.roll(mask_tooth, size_x // 3, axis=0) + np.roll( mask_tooth, -size_x // 3, axis=0) # make 3 teeth phantom[mask_tooth] = mu_tooth[0] #make implant mask_implant = (xx / (0.11 * size_x))**2 + (yy / (0.07 * size_x))**2 < 1 mask_implant *= y <= sx_half mask_implant *= ((xx / (0.11 * size_x))**2 + (((yy - 0.025 * size_x) / (0.07 * size_x)))**2) > 1 if (isimplant): phantom[mask_implant] = mu_implant[0] phantom *= pixel_size print("for Ca:") print(mu_tooth) print("for Au") print(mu_implant) return phantom
def attenuation_3d(src_path, theta_st, theta_end, n_theta, sample_height_n, sample_size_n, sample_size_cm, this_aN_dic, probe_energy, dev): n_element = len(this_aN_dic) theta_ls = -tc.linspace(theta_st, theta_end, n_theta + 1)[:-1] grid_concentration = tc.tensor(np.load(src_path)).float() aN_ls = np.array(list(this_aN_dic.values())) probe_attCS_ls = tc.tensor( xlib_np.CS_Total(aN_ls, probe_energy).flatten()).float().to(dev) att_exponent_acc_map = tc.zeros( (len(theta_ls), sample_height_n, sample_size_n, sample_size_n + 1), device=dev) for i, theta in enumerate(theta_ls): theta = tc.tensor(theta, device=dev) concentration_map_rot = rotate(grid_concentration, theta, dev) for j in range(n_element): lac_single = concentration_map_rot[j] * probe_attCS_ls[j] lac_acc = tc.cumsum(lac_single, axis=2) lac_acc = tc.cat((tc.zeros( (sample_height_n, sample_size_n, 1), device=dev), lac_acc), dim=2) att_exponent_acc = lac_acc * (sample_size_cm / sample_size_n) att_exponent_acc_map[i, :, :, :] += att_exponent_acc attenuation_map_flat = tc.exp(-(att_exponent_acc_map[:, :, :, :-1])).view( n_theta, sample_height_n * sample_size_n * sample_size_n) transmission = tc.exp(-att_exponent_acc_map[:, :, :, -1]).view( n_theta, sample_height_n * sample_size_n) return attenuation_map_flat, transmission
def self_absorption_att_ratio_single_theta_3d(src_path, n_det, P, det_size_cm, det_from_sample_cm, det_ds_spacing_cm, sample_size_n, sample_size_cm, sample_height_n, this_aN_dic, probe_energy, dev, theta): fl_all_lines_dic = MakeFLlinesDictionary(this_aN_dic, probe_energy, sample_size_n.cpu().numpy(), sample_size_cm.cpu().numpy(), fl_line_groups=np.array( ["K", "L", "M"]), fl_K=fl_K, fl_L=fl_L, fl_M=fl_M, group_lines=True) n_voxel = sample_height_n * sample_size_n * sample_size_n dia_len_n = int( (sample_height_n**2 + sample_size_n**2 + sample_size_n**2)**0.5) n_lines = tc.as_tensor(fl_all_lines_dic["n_lines"]).to(dev) aN_ls = np.array(list(this_aN_dic.values())) grid_concentration = tc.from_numpy(np.load(src_path)).float().to(dev) n_element = len(this_aN_dic) # generate an arrary of total attenuation cross section with the dimension: (n_element, n_elemental_lines) # The component in the array represents the total attenuation cross section at some line energy in some element (with unitary concentration) FL_line_attCS_ls = tc.as_tensor( xlib_np.CS_Total(aN_ls, fl_all_lines_dic["fl_energy"])).float().to(dev) concentration_map_rot = rotate(grid_concentration, theta, dev).float() concentration_map_rot_flat = concentration_map_rot.view( n_element, n_voxel).float() # lac: linear attenuation coefficient = concentration * attenuation_cross_section, # dimension: n_element, n_lines, n_voxel(FL source), n_voxel) lac = concentration_map_rot_flat.view(n_element, 1, 1, n_voxel) * FL_line_attCS_ls.view( n_element, n_lines, 1, 1) lac = lac.expand(-1, -1, n_voxel, -1).float() att_exponent = tc.stack([ lac[:, :, P[m][0].to(dtype=tc.long), P[m][1].to(dtype=tc.long)] * P[m][2].view(1, 1, -1).repeat(n_element, n_lines, 1) for m in range(n_det) ]) ## summing over the attenation exponent contributed by all intersecting voxels, dim = (n_det, n_element, n_lines, n_voxel (FL source)) att_exponent_voxel_sum = tc.sum(att_exponent.view(n_det, n_element, n_lines, n_voxel, dia_len_n), axis=-1) ## calculate the attenuation caused by all elements and get an array of dim = (n_det, n_lines, n_voxel (FL source)), and then take the average over n_det FL ray paths ## Final dim = (n_lines, n_voxel (FL source)) representing the attenuation ratio of each fluorescence line emitting from each source voxel. SA_att = tc.mean(tc.exp(-tc.sum(att_exponent_voxel_sum, axis=1)), axis=0) return SA_att
def absorption(energy, element): """Returns total element absorption for given energy. Both energy and element could be 1d-arrays """ element = np.array(element) dens = xraylib.ElementDensity(element) cross = xraylib.CS_Total(element, energy) return dens.reshape(-1, 1) * cross
def attenuation(src_path, n_theta, theta_ls, sample_size, sample_size_l, element_ls, an_lib): """ Calculate the attenuation ratio of the incident beam before the beam travels to a certain voxel Assuming that the x-ray probe goes along the direction of axis=1 of the sample array Parameters ---------- theta_ls: ndarray The angles that the sample rotates from the initial angle in the experiment sample_size: int scalar sample size in number of pixles on one side, assuing a N x N-pixel sample sample_size_l: scalar sample size in mm element_ls: ndarray elements in the sample Returns: ndarray ------- dimension of the returned array is n_theta x sample_size x sample_size """ an_ls = np.array(list(an_lib.values())) probe_energy = np.array([20.0]) ## genrate the library of the total attenuation cross section for the involved elements at 20 keV cs_probe_ls = xlib_np.CS_Total(an_ls, probe_energy).flatten() cs_probe_lib = dict(zip(element_ls, cs_probe_ls)) att_acc_map = np.zeros((n_theta, sample_size, sample_size)) for i, theta in enumerate(theta_ls): for j, element in enumerate(element_ls): concentration_map_fname = os.path.join(src_path, element + '_map.tiff') concentration_map = dxchange.reader.read_tiff( concentration_map_fname) concentration_map_rot = sp_rotate(concentration_map, theta, reshape=False, order=1) lac_single = concentration_map_rot * cs_probe_lib[element] lac_acc = np.cumsum(lac_single, axis=1) lac_acc = np.insert(lac_acc, 0, np.zeros(sample_size), axis=1) lac_acc = np.delete(lac_acc, sample_size, axis=1) att_acc = lac_acc * (sample_size_l / sample_size) att_acc_map[i, :, :] += att_acc return np.exp(-att_acc_map)
def attenuation_3d(src_path, theta_st, theta_end, n_theta, sample_height_n, sample_size_n, sample_size_cm, this_aN_dic, probe_energy, dev): """ Parameters ---------- src_path : string the path of the elemental concentration grid theta_st: float The initial angle of the sample theta_end: float The final angle of the sample n_theta: integer The number of sample angles sample_height_n : integer The height of the sample along the rotational axis (in number of pixels) sample_size_n: int scalar sample size in number of pixles on the side along the probe propagation axis sample_size_cm: scalar sample size in cm on the side along the probe propagation axis this_aN_dic: dictionary a dictionary of items with key = element symbol (string), and value = atomic number e.g. this_aN_dic = {"C":6, "O": 8} probe_energy : ndarray This array is an array with only 1 element. The element is the keV energy of the incident beam. dev : string specify "cpu" or the cuda diveice (ex: cuda:0) Returns ------- attenuation_map_flat : torch tensor an array of attenuation ratio before the probe enters each voxel. dim 0: all angles of the sample dim 1: all voxels (flattened 3D array) transmission : TYPE DESCRIPTION. """ n_element = len(this_aN_dic) theta_ls = -tc.linspace(theta_st, theta_end, n_theta + 1)[:-1] grid_concentration = tc.tensor(np.load(src_path)).float().to(dev) aN_ls = np.array(list(this_aN_dic.values())) probe_attCS_ls = tc.tensor( xlib_np.CS_Total(aN_ls, probe_energy).flatten()).float().to(dev) att_exponent_acc_map = tc.zeros( (len(theta_ls), sample_height_n, sample_size_n, sample_size_n + 1), device=dev) for i, theta in enumerate(theta_ls): theta = tc.tensor(theta, device=dev) concentration_map_rot = rotate(grid_concentration, theta, dev) for j in range(n_element): lac_single = concentration_map_rot[j] * probe_attCS_ls[j] lac_acc = tc.cumsum(lac_single, axis=2) lac_acc = tc.cat((tc.zeros( (sample_height_n, sample_size_n, 1), device=dev), lac_acc), dim=2) att_exponent_acc = lac_acc * (sample_size_cm / sample_size_n) att_exponent_acc_map[i, :, :, :] += att_exponent_acc attenuation_map_flat = tc.exp(-(att_exponent_acc_map[:, :, :, :-1])).view( n_theta, sample_height_n * sample_size_n * sample_size_n).float().to(dev) transmission = tc.exp(-att_exponent_acc_map[:, :, :, -1]).view( n_theta, sample_height_n * sample_size_n).float().to(dev) return attenuation_map_flat, transmission
## Define sample size in number of pixles on one side, assuing a N x N-pixel sample sample_size = 20 ## Define sample size in cm on one side sample_size_l = 0.01 src_path = '../data/sample1' ## Define probe posision, the position is defined to pass through the center of the voxel prob_pos_ls = np.array([x for x in np.arange(sample_size)]) + 0.5 element_ls = np.array(["C", "O", "Si", "Ca", "Fe"]) an_lib = {"C": 6, "O": 8, "Si": 14, "Ca": 20, "Fe": 26} an_ls = np.array(list(an_lib.values())) ## genrate the library of the total attenuation cross section for the involved elements at 20 keV probe_energy = np.array([20.0]) cs_probe_ls = xlib_np.CS_Total(an_ls, probe_energy).flatten() cs_probe_lib = dict(zip(element_ls, cs_probe_ls)) aw_ls = xlib_np.AtomicWeight(an_ls) aw_lib = dict(zip(element_ls, aw_ls)) n_det = 5 det_energy_u = 20 n_det_energy_bins = 2000 det_energy_list = np.linspace(det_energy_u / n_det_energy_bins, det_energy_u, n_det_energy_bins) ## genrate the library of the total attenuation cross section for the involved elements from 0-20keV att_cs_ls = xlib_np.CS_Total(an_ls, det_energy_list) att_cs_lib = dict(zip(element_ls, att_cs_ls))
def attenuation_theta(src_path, theta_st, theta_end, n_theta, this_theta_idx, sample_size_n, sample_size_cm, this_aN_dic, probe_energy): """ Calculate the attenuation ratio of the incident beam before the beam travels to a certain voxel. Assuming that the x-ray probe goes along the direction of axis=1 of the sample array. Calculate the transmission ratio at the exit of the sample. Parameters ---------- src_path: string the path of the elemental map theta_st: float The initial angle of the sample, in degree theta_end: float The final angle of the sample, in degree sample_size_n: int scalar sample size in number of pixles on one side, assuing a square sample of N x N pixels sample_size_cm: scalar sample size in cm this_aN_dic: dictionary a dictionary of items with key = element symbol (string), and value = atomic number e.g. this_aN_dic = {"C":6, "O": 8} Returns: 2 ndarrays ------- The 1st array is the attenuation ratio of the incident beam before the beam travels to a certain voxel. Assuming that the x-ray probe goes along the direction of axis=1 of the sample array. dimension of the 1st returned array is (n_theta, sample_size_n, sample_size_n) [note: sample_size may not be the same as the input argument because of padding] The 2nd array is the transmission ratio of the incident beam at the exit of the sample. dimension of the 2st returned array is (n_theta, sample_size_n) [note: sample_size may not be the same as the input argument because of padding] """ element_ls = np.array(list(this_aN_dic.keys())) aN_ls = np.array(list(this_aN_dic.values())) probe_attCS_ls = xlib_np.CS_Total(aN_ls, probe_energy).flatten() probe_attCS_dic = dict(zip(element_ls, probe_attCS_ls)) theta_ls = -np.linspace(theta_st, theta_end, n_theta) theta = theta_ls[this_theta_idx] grid_concentration = np.load( os.path.join(src_path, 'grid_concentration.npy')) att_exponent_acc_map = np.zeros((sample_size_n, sample_size_n + 1)) for j, element in enumerate(element_ls): concentration_map = grid_concentration[j] concentration_map_rot = sp_rotate(concentration_map, theta, reshape=False, order=1) lac_single = concentration_map_rot * probe_attCS_dic[element] lac_acc = np.cumsum(lac_single, axis=1) lac_acc = np.insert(lac_acc, 0, np.zeros(concentration_map.shape[0]), axis=1) att_exponent_acc = lac_acc * (sample_size_cm / sample_size_n) att_exponent_acc_map += att_exponent_acc attenuation_map_theta_flat = np.exp( -(att_exponent_acc_map[:, :-1].reshape(sample_size_n * sample_size_n))) transmission_theta = np.exp(-att_exponent_acc_map[:, -1]) return attenuation_map_theta_flat, transmission_theta
def self_absorption_att_ratio_single_theta(src_path, n_det, det_size_cm, det_from_sample_cm, sample_size_n, sample_size_cm, this_aN_dic, probe_energy, theta): """ Parameters ---------- grid_concentration : TYPE DESCRIPTION. n_det : TYPE DESCRIPTION. det_size_cm : TYPE DESCRIPTION. det_from_sample_cm : TYPE DESCRIPTION. sample_size_n : TYPE DESCRIPTION. sample_size_cm : TYPE DESCRIPTION. this_aN_dic : TYPE DESCRIPTION. probe_energy : TYPE DESCRIPTION. theta : TYPE DESCRIPTION. padding : TYPE, optional DESCRIPTION. The default is True. Returns ------- SA_theta : ndarray dimension: (sample_size_n * sample_size_n, n_elemental_line) """ aN_ls = np.array(list(this_aN_dic.values())) element_ls = np.array(list(this_aN_dic.keys())) grid_concentration = np.load( os.path.join(src_path, 'grid_concentration.npy')) P = intersecting_length_fl_detectorlet(n_det, det_size_cm, det_from_sample_cm, sample_size_n, sample_size_cm) fl_all_lines_dic = MakeFLlinesDictionary(this_aN_dic, probe_energy, sample_size_n, sample_size_cm, fl_line_groups=np.array( ["K", "L", "M"]), fl_K=fl_K, fl_L=fl_L, fl_M=fl_M, group_lines=True) # generate an arrary of total attenuation cross section with the dimension: (n_element, n_elemental_lines) # The component in the array represents the total attenuation cross section at some line energy in some element (with unitary concentration) FL_line_attCS_ls = xlib_np.CS_Total(aN_ls, fl_all_lines_dic["fl_energy"]) SA_theta = np.zeros((sample_size_n * sample_size_n, len(fl_all_lines_dic["(element_name, Line)"]))) for j in tqdm(np.arange(sample_size_n * sample_size_n)): att_exponent_elemental_sum_temp = np.zeros( (len(element_ls), n_det, len(fl_all_lines_dic["(element_name, Line)"]))) for k, element in enumerate(element_ls): concentration_map = grid_concentration[k] concentration_map_rot = sp_rotate(concentration_map, theta, reshape=False, order=1) ## flattened concentration_map after rotation: (sample_size_n * sample_size_n) concentration_map_rot_flat = concentration_map_rot.flatten() ## llinear attenuation coefficient for each fl-line at each voxel: (sample_size * sample_size, len(fl_lines_all["(element_name, Line)"])) lac = np.array([ FL_line_attCS * concentration_map_rot_flat for FL_line_attCS in FL_line_attCS_ls[k] ]) lac = np.transpose(lac) ## att_exponent = [(intersecting_length_path1 * lac), (intersecting_length_path2 * lac), ..., (intersecting_length_path5 * lac)]: ## att_exponent (for each fl-line, at each_voxel, for each beam path): (n_det, sample_size * sample_size, len(fl_lines_all["(element_name, Line)"])) att_exponent = np.array( [P[m, j, :][:, np.newaxis] * lac for m in range(n_det)]) ## att_exponent summing over voxels (for each line, for each beam path): (n_det, n_elemental_line) att_exponent_voxel_sum = np.sum(att_exponent, axis=1) ## filling att_exponent_voxel_sum to att_exponent_elemental_sum for each element att_exponent_elemental_sum_temp[k, :, :] = att_exponent_voxel_sum ## summing over the attenation exponent contributed by each element att_exponent_elemental_sum = np.sum(att_exponent_elemental_sum_temp, axis=0) ## calculate the attenuation caused by all elements att = np.exp(-att_exponent_elemental_sum) ## calculate the attenuation averaged all paths att_path_ave = np.average(att, axis=0) SA_theta[j, :] = att_path_ave return SA_theta
def reconstruct_jXRFT_tomography( dev, selfAb, recon_idx, cont_from_check_point, use_saved_initial_guess, recon_path, f_initial_guess, f_recon_grid, grid_path, f_grid, data_path, f_XRF_data, f_XRT_data, this_aN_dic, ini_kind, f_recon_parameters, n_epoch, n_minibatch, minibatch_size, b, lr, init_const, fl_line_groups, fl_K, fl_L, fl_M, group_lines, theta_st, theta_end, n_theta, sample_size_n, sample_height_n, sample_size_cm, probe_energy, probe_cts, det_size_cm, det_from_sample_cm, det_ds_spacing_cm, P_folder, f_P, ): if not os.path.exists(recon_path): os.makedirs(recon_path) stdout_options = { 'output_folder': recon_path, 'save_stdout': True, 'print_terminal': False } loss_fn = nn.MSELoss() X_true = tc.from_numpy( np.load(os.path.join(grid_path, f_grid)).astype(np.float32)).to(dev) dia_len_n = int( (sample_height_n**2 + sample_size_n**2 + sample_size_n**2)**0.5) n_voxel_minibatch = minibatch_size * sample_size_n n_voxel = sample_height_n * sample_size_n**2 aN_ls = np.array(list(this_aN_dic.values())) fl_all_lines_dic = MakeFLlinesDictionary(this_aN_dic, probe_energy, sample_size_n.cpu().numpy(), sample_size_cm.cpu().numpy(), fl_line_groups, fl_K, fl_L, fl_M, group_lines) FL_line_attCS_ls = tc.as_tensor( xlib_np.CS_Total(aN_ls, fl_all_lines_dic["fl_energy"])).float().to(dev) n_lines = fl_all_lines_dic["n_lines"] minibatch_ls_0 = tc.arange(n_minibatch).to(dev) n_batch = (sample_height_n * sample_size_n) // (n_minibatch * minibatch_size) theta_ls = -tc.linspace(theta_st, theta_end, n_theta + 1)[:-1].to(dev) n_element = len(this_aN_dic) if not os.path.exists(P_folder): os.makedirs(P_folder) P_save_path = os.path.join(P_folder, f_P) longest_int_length, n_det, P = intersecting_length_fl_detectorlet_3d( det_size_cm, det_from_sample_cm, det_ds_spacing_cm, sample_size_n.cpu().numpy(), sample_size_cm.cpu().numpy(), sample_height_n.cpu().numpy(), P_save_path) P = tc.from_numpy(P).float() # P = P.view(n_det, 3, dia_len_n * sample_height_n * sample_size_n * sample_size_n) if cont_from_check_point == False: if use_saved_initial_guess: X = np.load(os.path.join(recon_path, f_initial_guess) + '.npy') X = tc.from_numpy(X).float().to(dev) ## Save the initial guess which will be used in reconstruction and will be updated to the current reconstructing result np.save(os.path.join(recon_path, f_recon_grid) + '.npy', X.cpu()) else: X = initialize_guess_3d(dev, ini_kind, grid_path, f_grid, recon_path, f_recon_grid, f_initial_guess, init_const) with open(os.path.join(recon_path, f_recon_parameters), "w") as recon_params: recon_params.write("starting_epoch = 0\n") recon_params.write("n_epoch = %d\n" % n_epoch) recon_params.write(str(this_aN_dic) + "\n") recon_params.write("n_minibatch = %d\n" % n_minibatch) recon_params.write("minibatch_size = %d\n" % minibatch_size) recon_params.write("b = %.9f\n" % b) recon_params.write("learning rate = %f\n" % lr) recon_params.write("theta_st = %.2f\n" % theta_st) recon_params.write("theta_end = %.2f\n" % theta_end) recon_params.write("n_theta = %d\n" % n_theta) recon_params.write("sample_size_n = %d\n" % sample_size_n) recon_params.write("sample_height_n = %d\n" % sample_height_n) recon_params.write("sample_size_cm = %.2f\n" % sample_size_cm) recon_params.write("probe_energy = %.2f\n" % probe_energy[0]) recon_params.write("probe_cts = %.2e\n" % probe_cts) recon_params.write("det_size_cm = %.2f\n" % det_size_cm) recon_params.write("det_from_sample_cm = %.2f\n" % det_from_sample_cm) recon_params.write("det_ds_spacing_cm = %.2f\n" % det_ds_spacing_cm) loss_minibatch = tc.zeros(n_minibatch * n_batch * n_theta * n_epoch, device=dev) mse_epoch = tc.zeros(n_epoch, len(this_aN_dic), device=dev) rand_idx = tc.randperm(n_theta) theta_ls = theta_ls[rand_idx] for epoch in tqdm(range(n_epoch)): t0_epoch = time.perf_counter() for idx, theta in enumerate(theta_ls): this_theta_idx = rand_idx[idx] # for each theta, load the current object and update the grid concentration that is used to calculate the absorption. # X dimension: [C, N, H, W] X = np.load(os.path.join(recon_path, f_recon_grid) + '.npy').astype(np.float32) X = tc.from_numpy(X).to(dev) ## Calculate lac using the current X. lac (linear attenuation coefficient) has the dimension of [n_element, n_lines, n_voxel_minibatch, n_voxel] if selfAb == True: X_ap_rot = rotate(X, theta, dev).view( n_element, sample_height_n * sample_size_n, sample_size_n) lac = X_ap_rot.view(n_element, 1, 1, n_voxel) * FL_line_attCS_ls.view( n_element, n_lines, 1, 1) lac = lac.expand(-1, -1, n_voxel_minibatch, -1).float() else: lac = tc.tensor(0.) ## load data y1: XRF data, y2: XRT data y1_true = tc.from_numpy( np.load( os.path.join(data_path, f_XRF_data) + '_{}'.format(this_theta_idx) + '.npy').astype( np.float32)).to(dev) y2_true = tc.from_numpy( np.load( os.path.join(data_path, f_XRT_data) + '_{}'.format(this_theta_idx) + '.npy').astype( np.float32)).to(dev) for m in range(n_batch): minibatch_ls = n_minibatch * m + minibatch_ls_0 P_this_batch = P[:, :, minibatch_ls[0] * dia_len_n * minibatch_size * sample_size_n:(minibatch_ls[0] + len(minibatch_ls)) * dia_len_n * minibatch_size * sample_size_n] P_this_batch = P_this_batch.view( n_det, 3, len(minibatch_ls), dia_len_n * minibatch_size * sample_size_n) P_this_batch = P_this_batch.permute(2, 0, 1, 3) for ip, p in enumerate(minibatch_ls): model = PPM(dev, selfAb, lac, X, p, n_element, sample_height_n, minibatch_size, sample_size_n, sample_size_cm, this_aN_dic, probe_energy, probe_cts, theta_st, theta_end, n_theta, this_theta_idx, n_det, P_this_batch[ip], det_size_cm, det_from_sample_cm).to(dev) tc.cuda.empty_cache() optimizer = tc.optim.Adam(model.parameters(), lr=lr) y1_hat, y2_hat = model() XRF_loss = loss_fn( y1_hat, y1_true[:, minibatch_size * p:minibatch_size * (p + 1)]) XRT_loss = loss_fn( y2_hat, y2_true[minibatch_size * p:minibatch_size * (p + 1)]) loss = XRF_loss + b * XRT_loss loss_minibatch[ (n_minibatch * n_batch * n_theta) * epoch + (n_minibatch * n_batch) * this_theta_idx + n_minibatch * m + ip] = float(loss) optimizer.zero_grad() loss.backward() optimizer.step() X[:, minibatch_size * p // sample_size_n:minibatch_size * (p + 1) // sample_size_n, :, :] = model.xp.detach() del model del P_this_batch X = tc.clamp(X, 0, float('inf')) np.save( os.path.join(recon_path, f_recon_grid) + '.npy', X.detach().cpu().numpy()) del lac tc.cuda.empty_cache() mse_epoch[epoch] = tc.mean(tc.square(X - X_true).view( X.shape[0], X.shape[1] * X.shape[2] * X.shape[3]), dim=1) per_epoch_time = time.perf_counter() - t0_epoch print_flush( val=per_epoch_time, output_file=f'per_epoch_time_mb_size_{minibatch_size}.csv', **stdout_options) tqdm._instances.clear() mse_epoch_tot = tc.mean(mse_epoch, dim=1) fig6 = plt.figure(figsize=(15, 5)) gs6 = gridspec.GridSpec(nrows=1, ncols=2, width_ratios=[1, 1]) fig6_ax1 = fig6.add_subplot(gs6[0, 0]) fig6_ax1.plot(loss_minibatch.detach().cpu().numpy()) fig6_ax1.set_xlabel('minibatch') fig6_ax1.set_ylabel('loss') fig6_ax1.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.5e')) fig6_ax2 = fig6.add_subplot(gs6[0, 1]) fig6_ax2.plot(mse_epoch_tot.detach().cpu().numpy()) fig6_ax2.set_xlabel('epoch') fig6_ax2.set_ylabel('mse of model') fig6_ax2.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.2f')) plt.savefig(os.path.join(recon_path, 'loss_and_tot_mse.pdf')) fig7 = plt.figure(figsize=(X.shape[0] * 6, 4)) gs7 = gridspec.GridSpec(nrows=1, ncols=X.shape[0], width_ratios=[1] * X.shape[0]) for i in range(X.shape[0]): fig7_ax1 = fig7.add_subplot(gs7[0, i]) fig7_ax1.plot(mse_epoch[:, i].detach().cpu().numpy()) fig7_ax1.set_xlabel('epoch') fig7_ax1.set_ylabel('mse of model (each element)') fig7_ax1.set_title(str(list(this_aN_dic.keys())[i])) fig7_ax1.yaxis.set_major_formatter( mtick.FormatStrFormatter('%.2f')) plt.savefig(os.path.join(recon_path, 'mse_model.pdf')) np.save(os.path.join(recon_path, 'loss_minibatch.npy'), loss_minibatch.detach().cpu().numpy()) np.save(os.path.join(recon_path, 'mse_model_elements.npy'), mse_epoch.detach().cpu().numpy()) np.save(os.path.join(recon_path, 'mse_model.npy'), mse_epoch_tot.detach().cpu().numpy()) dxchange.write_tiff(X.detach().cpu().numpy(), os.path.join(recon_path, f_recon_grid) + "_" + str(recon_idx), dtype='float32', overwrite=True) if cont_from_check_point == True: recon_idx += 1 loss_minibatch = tc.from_numpy( np.load(os.path.join(recon_path, 'loss_minibatch.npy')).astype(np.float32)) mse_epoch = tc.from_numpy( np.load(os.path.join(recon_path, 'mse_model_elements.npy')).astype(np.float32)) mse_epoch_tot = tc.from_numpy( np.load(os.path.join(recon_path, 'mse_model.npy')).astype(np.float32)) with open(os.path.join(recon_path, f_recon_parameters), "r") as recon_params: params_list = [] for line in recon_params.readlines(): params_list.append(line.rstrip("\n")) n_ending = len(params_list) with open(os.path.join(recon_path, f_recon_parameters), "a") as recon_params: n_start_last = n_ending - 18 previous_epoch = int( params_list[n_start_last][params_list[n_start_last].find("=") + 1:]) recon_params.write("\n") recon_params.write("###########################################\n") recon_params.write("starting_epoch = %d\n" % (previous_epoch + n_epoch)) recon_params.write("n_epoch = %d\n" % n_epoch) recon_params.write(str(this_aN_dic) + "\n") recon_params.write("n_minibatch = %d\n" % n_minibatch) recon_params.write("minibatch_size = %d\n" % minibatch_size) recon_params.write("b = %f\n" % b) recon_params.write("learning rate = %f\n" % lr) recon_params.write("theta_st = %.2f\n" % theta_st) recon_params.write("theta_end = %.2f\n" % theta_end) recon_params.write("n_theta = %d\n" % n_theta) recon_params.write("sample_size_n = %d\n" % sample_size_n) recon_params.write("sample_height_n = %d\n" % sample_height_n) recon_params.write("sample_size_cm = %.2f\n" % sample_size_cm) recon_params.write("probe_energy = %.2f\n" % probe_energy[0]) recon_params.write("probe_cts = %.2e\n" % probe_cts) recon_params.write("det_size_cm = %.2f\n" % det_size_cm) recon_params.write("det_from_sample_cm = %.2f\n" % det_from_sample_cm) recon_params.write("det_ds_spacing_cm = %.2f\n" % det_ds_spacing_cm) del recon_params del params_list loss_minibatch_cont = tc.zeros(n_minibatch * n_batch * n_theta * n_epoch, device=dev) mse_epoch_cont = tc.zeros(n_epoch, len(this_aN_dic), device=dev) rand_idx = tc.randperm(n_theta) theta_ls = theta_ls[rand_idx] for epoch in tqdm(range(n_epoch)): t0_epoch = time.perf_counter() for idx, theta in enumerate(theta_ls): this_theta_idx = rand_idx[idx] # for each theta, load the current object and update the grid concentration that is used to calculate the absorption. # X dimension: [C, N, H, W] X = np.load(os.path.join(recon_path, f_recon_grid) + '.npy').astype(np.float32) X = tc.from_numpy(X).to(dev) ## Calculate lac using the current X. lac (linear attenuation coefficient) has the dimension of [n_element, n_lines, n_voxel_minibatch, n_voxel] if selfAb == True: X_ap_rot = rotate(X, theta, dev).view( n_element, sample_height_n * sample_size_n, sample_size_n) lac = X_ap_rot.view(n_element, 1, 1, n_voxel) * FL_line_attCS_ls.view( n_element, n_lines, 1, 1) lac = lac.expand(-1, -1, n_voxel_minibatch, -1).float() else: lac = 0. ## load data y1: XRF data, y2: XRT data y1_true = tc.from_numpy( np.load( os.path.join(data_path, f_XRF_data) + '_{}'.format(this_theta_idx) + '.npy').astype( np.float32)).to(dev) y2_true = tc.from_numpy( np.load( os.path.join(data_path, f_XRT_data) + '_{}'.format(this_theta_idx) + '.npy').astype( np.float32)).to(dev) for m in range(n_batch): minibatch_ls = n_minibatch * m + minibatch_ls_0 P_this_batch = P[:, :, minibatch_ls[0] * dia_len_n * minibatch_size * sample_size_n:(minibatch_ls[0] + len(minibatch_ls)) * dia_len_n * minibatch_size * sample_size_n] P_this_batch = P_this_batch.view( n_det, 3, len(minibatch_ls), dia_len_n * minibatch_size * sample_size_n) P_this_batch = P_this_batch.permute(2, 0, 1, 3) for ip, p in enumerate(minibatch_ls): model = PPM(dev, selfAb, lac, X, p, n_element, sample_height_n, minibatch_size, sample_size_n, sample_size_cm, this_aN_dic, probe_energy, probe_cts, theta_st, theta_end, n_theta, this_theta_idx, n_det, P_this_batch[ip], det_size_cm, det_from_sample_cm).to(dev) tc.cuda.empty_cache() optimizer = tc.optim.Adam(model.parameters(), lr=lr) y1_hat, y2_hat = model() XRF_loss = loss_fn( y1_hat, y1_true[:, minibatch_size * p:minibatch_size * (p + 1)]) XRT_loss = loss_fn( y2_hat, y2_true[minibatch_size * p:minibatch_size * (p + 1)]) loss = XRF_loss + b * XRT_loss loss_minibatch_cont[ (n_minibatch * n_batch * n_theta) * epoch + (n_minibatch * n_batch) * this_theta_idx + n_minibatch * m + ip] = float(loss) optimizer.zero_grad() loss.backward() optimizer.step() X[:, minibatch_size * p // sample_size_n:minibatch_size * (p + 1) // sample_size_n, :, :] = model.xp.detach() del model del P_this_batch X = tc.clamp(X, 0, float('inf')) np.save( os.path.join(recon_path, f_recon_grid) + '.npy', X.detach().cpu().numpy()) del lac tc.cuda.empty_cache() mse_epoch_cont[epoch] = tc.mean(tc.square(X - X_true).view( X.shape[0], X.shape[1] * X.shape[2] * X.shape[3]), dim=1) per_epoch_time = time.perf_counter() - t0_epoch print_flush( val=per_epoch_time, output_file=f'per_epoch_time_mb_size_{minibatch_size}.csv', **stdout_options) tqdm._instances.clear() mse_epoch_tot_cont = tc.mean(mse_epoch_cont, dim=1) loss_minibatch = tc.cat((loss_minibatch, loss_minibatch_cont.cpu())) mse_epoch = tc.cat((mse_epoch, mse_epoch_cont.cpu())) mse_epoch_tot = tc.cat((mse_epoch_tot, mse_epoch_tot_cont.cpu())) fig6 = plt.figure(figsize=(15, 5)) gs6 = gridspec.GridSpec(nrows=1, ncols=2, width_ratios=[1, 1]) fig6_ax1 = fig6.add_subplot(gs6[0, 0]) fig6_ax1.plot(loss_minibatch.detach().cpu().numpy()) fig6_ax1.set_xlabel('minibatch') fig6_ax1.set_ylabel('loss') fig6_ax1.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.5e')) fig6_ax2 = fig6.add_subplot(gs6[0, 1]) fig6_ax2.plot(mse_epoch_tot.detach().cpu().numpy()) fig6_ax2.set_xlabel('epoch') fig6_ax2.set_ylabel('mse of model') fig6_ax2.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.2f')) plt.savefig(os.path.join(recon_path, 'loss_and_tot_mse.pdf')) fig7 = plt.figure(figsize=(X.shape[0] * 6, 4)) gs7 = gridspec.GridSpec(nrows=1, ncols=X.shape[0], width_ratios=[1] * X.shape[0]) for i in range(X.shape[0]): fig7_ax1 = fig7.add_subplot(gs7[0, i]) fig7_ax1.plot(mse_epoch_cont[:, i].detach().cpu().numpy()) fig7_ax1.set_xlabel('epoch') fig7_ax1.set_ylabel('mse of model (each element)') fig7_ax1.set_title(str(list(this_aN_dic.keys())[i])) fig7_ax1.yaxis.set_major_formatter( mtick.FormatStrFormatter('%.2f')) plt.savefig(os.path.join(recon_path, 'mse_model.pdf')) np.save(os.path.join(recon_path, 'loss_minibatch.npy'), loss_minibatch.detach().cpu().numpy()) np.save(os.path.join(recon_path, 'mse_model_elements.npy'), mse_epoch.detach().cpu().numpy()) np.save(os.path.join(recon_path, 'mse_model.npy'), mse_epoch_tot.detach().cpu().numpy()) dxchange.write_tiff(X.detach().cpu().numpy(), os.path.join(recon_path, f_recon_grid) + "_" + str(recon_idx), dtype='float32', overwrite=True)
def self_absorption_att_ratio_single_theta_3d(src_path, det_size_cm, det_from_sample_cm, det_ds_spacing_cm, sample_size_n, sample_size_cm, sample_height_n, this_aN_dic, probe_energy, dev, theta): fl_all_lines_dic = MakeFLlinesDictionary(this_aN_dic, probe_energy, sample_size_n.numpy(), sample_size_cm.numpy(), fl_line_groups=np.array( ["K", "L", "M"]), fl_K=fl_K, fl_L=fl_L, fl_M=fl_M, group_lines=True) P_all = intersecting_length_fl_detectorlet_3d( det_size_cm, det_from_sample_cm, det_ds_spacing_cm, sample_size_n, sample_size_cm, sample_height_n) n_det = P_all[0] P = tc.from_numpy(P_all[1]) n_lines = fl_all_lines_dic["n_lines"] aN_ls = np.array(list(this_aN_dic.values())) element_ls = np.array(list(this_aN_dic.keys())) grid_concentration = tc.from_numpy(np.load(src_path)).float() n_element = len(this_aN_dic) # generate an arrary of total attenuation cross section with the dimension: (n_element, n_elemental_lines) # The component in the array represents the total attenuation cross section at some line energy in some element (with unitary concentration) FL_line_attCS_ls = xlib_np.CS_Total(aN_ls, fl_all_lines_dic["fl_energy"]) concentration_map_rot = rotate(grid_concentration, theta, dev) concentration_map_rot_flat = concentration_map_rot.view( n_element, sample_height_n * sample_size_n * sample_size_n) SA_theta = tc.zeros( (n_lines, sample_height_n * sample_size_n * sample_size_n), device=dev) for j in tqdm(range(sample_height_n * sample_size_n * sample_size_n)): att_exponent_elemental_sum_temp = tc.zeros((n_element, n_det, n_lines), device=dev) for k in range(n_element): ## linear attenuation coefficient for each fl-line at each voxel: (sample_height_n * sample_size_n * sample_size_n, n_lines) lac = tc.stack([ FL_line_attCS * concentration_map_rot_flat[k] for FL_line_attCS in FL_line_attCS_ls[k] ], dim=1) # print(lac.shape) ## att_exponent = [(intersecting_length_path1 * lac), (intersecting_length_path2 * lac), ..., (intersecting_length_path5 * lac)]: ## att_exponent (for each fl-line, at each_voxel, for each beam path): (self.n_det, sample_size * sample_size, self.n_lines) att_exponent = tc.stack( [tc.unsqueeze(P[m, j, :], dim=1) * lac for m in range(n_det)]) # print(att_exponent.shape) ## att_exponent summing over voxels (for each line, for each beam path): (self.n_det, n_elemental_line) att_exponent_voxel_sum = tc.sum(att_exponent, axis=1) ## filling att_exponent_voxel_sum to att_exponent_elemental_sum for each element att_exponent_elemental_sum_temp[k, :, :] = att_exponent_voxel_sum ## summing over the attenation exponent contributed by each element att_exponent_elemental_sum = tc.sum(att_exponent_elemental_sum_temp, axis=0) ## calculate the attenuation caused by all elements att = tc.exp(-att_exponent_elemental_sum) ## calculate the attenuation averaged all paths att_path_ave = tc.mean(att, axis=0) SA_theta[:, j] = att_path_ave # SA_theta = np.array(SA_theta) return SA_theta
def attlen(material, E, density=None): E = np.atleast_1d(E) E = E.astype('double') Z = np.atleast_1d(getElementZ(material)) ma = xraylib_np.AtomicWeight(Z) return xraylib_np.CS_Total(Z, np.atleast_1d(E)) / ma * density