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
