From 51be44b090dd61b8efc917936bfcbd5fb7d4226b Mon Sep 17 00:00:00 2001 From: Kanupriya Pande Date: Wed, 18 May 2022 17:33:37 -0700 Subject: [PATCH] modules for rigid alignment --- utilities/alignment_functions.py | 485 ++++++++++++++++++++++++++++++ utilities/projection_operators.py | 122 ++++++++ utilities/ray_voxel_utilities.py | 345 +++++++++++++++++++++ utilities/voxel_utilities.py | 108 +++++++ 4 files changed, 1060 insertions(+) create mode 100644 utilities/alignment_functions.py create mode 100644 utilities/projection_operators.py create mode 100644 utilities/ray_voxel_utilities.py create mode 100644 utilities/voxel_utilities.py diff --git a/utilities/alignment_functions.py b/utilities/alignment_functions.py new file mode 100644 index 0000000..9b161dc --- /dev/null +++ b/utilities/alignment_functions.py @@ -0,0 +1,485 @@ +import numpy as np +from scipy import sparse, optimize +from scipy.optimize import line_search +from scipy.optimize.linesearch import line_search_armijo, line_search_wolfe1 + + +class AlignmentUtilities(object): + + def __init__(self, proj, proj_obj, geometry): + + self.proj = proj + self.proj_obj = proj_obj + self.proj_mask = proj > 0 + self.geometry = geometry + + def cost(self, rec, angles, translations): + + phi, alpha, beta = angles + xyz_shift = translations + this_proj, _ = self.proj_obj.projection_gradient(rec=rec, alpha=alpha, beta=beta, phi=phi, xyz_shift=xyz_shift, + cor_shift=self.geometry.cor_shift) + + residual = self.proj.ravel() - this_proj + + return residual + + def gradient(self, rec, angles, translations): + + phi, alpha, beta = angles + xyz_shift = translations + this_proj, this_grad = self.proj_obj.projection_gradient(rec=rec, alpha=alpha, beta=beta, phi=phi, + xyz_shift=xyz_shift, cor_shift=self.geometry.cor_shift) + + residual = self.proj.ravel() - this_proj + this_grad *= -1 + + return residual, this_grad + + +def gradient_descent(x, cost_function, gradient_function, args=(), options={}): + + n_itmax = options['maxiter'] if 'maxiter' in options else 100 + step_search = options['step_search'] if 'step_search' in options else 'armijo' + eps = options['eps'] if 'eps' in options else 1.e-6 + verbose = options['verbose'] if 'verbose' in options else False + + align_obj, rec, angles_in, xyz_in, scale_factor = args + + cost = np.zeros(n_itmax+1) + stop = 0 + it = 0 + + f = cost_function(x, align_obj, rec, angles_in, xyz_in, scale_factor=scale_factor, return_vector=False) + fp = gradient_function(x, align_obj, rec, angles_in, xyz_in, scale_factor=scale_factor, return_vector=False) + + cost[it] = f + ls_counter = 0 + alpha = 0.0 + while not stop and it < n_itmax: + if verbose: + print(it, f, alpha, fp, x) + + # linesearch with armijo + search_dir = -fp#/np.linalg.norm(fp) + if step_search == 'armijo': + alpha, _, f_new = line_search_armijo(cost_function, x, search_dir, fp, cost[it], alpha0=1.0, + args=(align_obj, rec, angles_in, xyz_in, None)) + elif step_search == 'wolfe': + if it > 1: + old_old_val = cost[it-1] + else: + old_old_val = None + #alpha,_,_, f_new, f_old, fp_new = line_search(cost_function, gradient_function, x, -fp, gfk=fp, old_fval=cost[it], + # old_old_fval=old_old_val, + # args=(align_obj, rec, angles_in, xyz_in, None)) + alpha, _, _, f_new, f_old, fp_new = line_search_wolfe1(cost_function, gradient_function, x, search_dir, + gfk=fp, amax=1.e-3, amin=1.e-12, + args=(align_obj, rec, angles_in, xyz_in, None, None)) + if alpha is None: + print('%s line search failed' %(step_search)) + ls_success = False + ls_counter += 1 + alpha = 1.0 + # starting at alpha=1.e-6, try successively dividing alpha by 10 to find step size + # and terminate GD here + while not ls_success and alpha > 1.e-15: + alpha = alpha/10 + x_new = x - alpha * fp + f_new = cost_function(x_new, align_obj, rec, angles_in, xyz_in, scale_factor, False) + if f_new < cost[it]: + ls_success = True + + if not ls_success or ls_counter >= 2: + # either linesearch or more than 2 successive brute linesearch + # return previous values + stop = 2 + it += 1 + if verbose: + print('either linesearch failed or two successive brute linesearch iterations') + else: + x = x - alpha * fp + it += 1 + f = cost_function(x, align_obj, rec, angles_in, xyz_in, scale_factor, False) + fp = gradient_function(x, align_obj, rec, angles_in, xyz_in, scale_factor, False) + + cost[it] = f + if np.abs(cost[it] - cost[it-1])/max(cost[it], cost[it-1], 1.0) <= eps: + stop = 1 + + return x, f, stop + + +def cost_xzpab(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + translations = np.array([parameters[0]+xyz_in[0], xyz_in[1], xyz_in[2]+parameters[1]]) + phi, alpha, beta = angles_in + parameters[2:] + angles = np.array([phi, alpha, beta]) + + cost = align_obj.cost(rec, angles, translations) + + if return_vector: + return cost + + return 0.5 * np.linalg.norm(cost) ** 2 + + +def gradient_xzpab(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + translations = np.array([parameters[0]+xyz_in[0], xyz_in[1], xyz_in[2]+parameters[1]]) + phi, alpha, beta = angles_in + parameters[2:] + angles = np.array([phi, alpha, beta]) + + residual, s = align_obj.gradient(rec, angles, translations) + + vary_parameter = np.array([True, False, True, True, True, True]) + s = s[vary_parameter, :] + + if scale_factor is None: + scale_factor = np.array([1.0, 1.0, 1.0, 1.0, 1.0]) + + s = s * scale_factor[:, np.newaxis] + + if return_vector: + return s.T + + grad = np.dot(s, residual) + + return grad + + +def cost_xzab(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + translations = np.array([xyz_in[0]+parameters[0], xyz_in[1], xyz_in[2]+parameters[1]]) + phi = angles_in[0] + alpha, beta = angles_in[1:] + parameters[2:] + angles = np.array([phi, alpha, beta]) + + cost = align_obj.cost(rec, angles, translations) + + if return_vector: + return cost + + return 0.5*np.linalg.norm(cost)**2 + + +def gradient_xzab(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + translations = np.array([xyz_in[0]+parameters[0], xyz_in[1], xyz_in[2]+parameters[1]]) + phi = angles_in[0] + alpha, beta = angles_in[1:] + parameters[2:] + angles = np.array([phi, alpha, beta]) + + residual, s = align_obj.gradient(rec, angles, translations) + + vary_parameter = np.array([True, False, True, False, True, True]) + s = s[vary_parameter] + + if scale_factor is None: + scale_factor = np.array([1.0, 1.0, 1.0, 1.0]) + + s = s*scale_factor[:, np.newaxis] + + if return_vector: + return s.T + + grad = np.dot(s, residual) + + return grad + + +def cost_xz(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + translations = np.array([xyz_in[0]+parameters[0], xyz_in[1], xyz_in[2]+parameters[1]]) + + cost = align_obj.cost(rec, angles_in, translations) + + if return_vector: + return cost + + return 0.5 * np.linalg.norm(cost) ** 2 + + +def gradient_xz(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + translations = np.array([xyz_in[0]+parameters[0], xyz_in[1], xyz_in[2]+parameters[1]]) + + residual, s = align_obj.gradient(rec, angles_in, translations) + + vary_parameter = np.array([True, False, True, False, False, False]) + s = s[vary_parameter] + + if scale_factor is None: + scale_factor = np.array([1.0, 1.0]) + + s = s * scale_factor[:, np.newaxis] + + if return_vector: + return s.T + + grad = np.dot(s, residual) + + return grad + + +def gradient_xz_fd(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + translations = np.array([xyz_in[0]+parameters[0], xyz_in[1], xyz_in[2]+parameters[1]]) + eps = np.array([1.e-4, 1.e-3, 1.e-4]) + grad = np.zeros(3,) + for i in range(3): + tt = np.zeros(3,) + tt[i] = 1.0 + trans_plus = translations + tt*eps + trans_minus = translations - tt*eps + cost_plus = align_obj.cost(rec, angles_in, trans_plus) + cost_minus = align_obj.cost(rec, angles_in, trans_minus) + cost_plus = 0.5*np.linalg.norm(cost_plus)**2 + cost_minus = 0.5*np.linalg.norm(cost_minus)**2 + grad[i] = (cost_plus - cost_minus)/(2*eps[i]) + + return np.array([grad[0], grad[2]]) + + +def cost_x(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + translations = np.array([xyz_in[0]+parameters[0], xyz_in[1], xyz_in[2]]) + + cost = align_obj.cost(rec, angles_in, translations) + + if return_vector: + return cost + + return 0.5 * np.linalg.norm(cost) ** 2 + + +def gradient_x(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + translations = np.array([xyz_in[0]+parameters[0], xyz_in[1], xyz_in[2]]) + + residual, s = align_obj.gradient(rec, angles_in, translations) + + vary_parameter = np.array([True, False, False, False, False, False]) + s = s[vary_parameter] + + if scale_factor is None: + scale_factor = np.array([1.0]) + + s = s * scale_factor[:, np.newaxis] + + if return_vector: + return s.T + + grad = np.dot(s, residual) + + return grad + + +def cost_z(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + translations = np.array([xyz_in[0], xyz_in[1], xyz_in[2]+parameters[0]]) + + cost = align_obj.cost(rec, angles_in, translations) + + if return_vector: + return cost + + return 0.5 * np.linalg.norm(cost) ** 2 + + +def gradient_z(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + translations = np.array([xyz_in[0], xyz_in[1], xyz_in[2]+parameters[0]]) + + residual, s = align_obj.gradient(rec, angles_in, translations) + + vary_parameter = np.array([False, False, True, False, False, False]) + s = s[vary_parameter] + + if scale_factor is None: + scale_factor = np.array([1.0]) + + s = s * scale_factor[:, np.newaxis] + + if return_vector: + return s.T + + grad = np.dot(s, residual) + + return grad + + +def cost_ab(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + phi = angles_in[0] + alpha, beta = angles_in[1:]+parameters + angles = np.array([phi, alpha, beta]) + cost = align_obj.cost(rec, angles, xyz_in) + + if return_vector: + return cost + + return 0.5 * np.linalg.norm(cost) ** 2 + + +def gradient_ab(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + phi = angles_in[0] + alpha, beta = angles_in[1:] + parameters + angles = np.array([phi, alpha, beta]) + residual, s = align_obj.gradient(rec, angles, xyz_in) + + vary_parameter = np.array([False, False, False, False, True, True]) + s = s[vary_parameter] + + if scale_factor is None: + scale_factor = np.array([1.0, 1.0]) + + s = s * scale_factor[:, np.newaxis] + + if return_vector: + return s.T + + grad = np.dot(s, residual) + + return grad + + +def cost_a(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + phi = angles_in[0] + alpha = angles_in[1] + parameters[0] + beta = angles_in[2] + angles = np.array([phi, alpha, beta]) + cost = align_obj.cost(rec, angles, xyz_in) + + if return_vector: + return cost + + return 0.5 * np.linalg.norm(cost) ** 2 + + +def gradient_a(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + phi = angles_in[0] + alpha = angles_in[1] + parameters[0] + beta = angles_in[2] + angles = np.array([phi, alpha, beta]) + residual, s = align_obj.gradient(rec, angles, xyz_in) + + vary_parameter = np.array([False, False, False, False, True, False]) + s = s[vary_parameter] + + if scale_factor is None: + scale_factor = np.array([1.0]) + + s = s * scale_factor[:, np.newaxis] + + if return_vector: + return s.T + + grad = np.dot(s, residual)#[vary_parameter] + + return grad + + +def cost_b(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + phi = angles_in[0] + alpha = angles_in[1] + beta = angles_in[2] + parameters[0] + angles = np.array([phi, alpha, beta]) + cost = align_obj.cost(rec, angles, xyz_in) + + if return_vector: + return cost + + return 0.5 * np.linalg.norm(cost) ** 2 + + +def gradient_b(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + phi = angles_in[0] + alpha = angles_in[1] + beta = angles_in[2] + parameters[0] + angles = np.array([phi, alpha, beta]) + residual, s = align_obj.gradient(rec, angles, xyz_in) + + vary_parameter = np.array([False, False, False, False, False, True]) + s = s[vary_parameter] + + if scale_factor is None: + scale_factor = np.array([1.0]) + + s = s * scale_factor[:, np.newaxis] + + if return_vector: + return s.T + + grad = np.dot(s, residual)#[vary_parameter] + + return grad + + +def gradient_ab_fd(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + phi = angles_in[0] + alpha, beta = angles_in[1:] + parameters + angles = np.array([phi, alpha, beta]) + residual, s = align_obj.gradient(rec, angles, xyz_in) + + eps = np.array([1.e-3, 1.e-4, 1.e-4]) + grad = np.zeros(3,) + for i in range(3): + tt = np.zeros(3,) + tt[i] = 1.0 + angles_plus = angles_in + tt*eps + angles_minus = angles_in - tt*eps + + cost_plus = align_obj.cost(rec, angles_plus, xyz_in) + cost_minus = align_obj.cost(rec, angles_minus, xyz_in) + cost_plus = 0.5*np.linalg.norm(cost_plus)**2 + cost_minus = 0.5*np.linalg.norm(cost_minus)**2 + grad[i] = (cost_plus - cost_minus)/(2*eps[i]) + + return grad[1:] + + +def cost_xzb(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + translations = np.array([xyz_in[0]+parameters[0], xyz_in[1], xyz_in[2]+parameters[1]]) + phi = angles_in[0] + alpha = angles_in[1] + beta = angles_in[2] + parameters[2] + angles = np.array([phi, alpha, beta]) + cost = align_obj.cost(rec, angles, translations) + + if return_vector: + return cost + + return 0.5 * np.linalg.norm(cost) ** 2 + + +def gradient_xzb(parameters, align_obj, rec, angles_in, xyz_in, scale_factor=None, return_vector=False): + + translations = np.array([xyz_in[0] + parameters[0], xyz_in[1], xyz_in[2] + parameters[1]]) + phi = angles_in[0] + alpha = angles_in[1] + beta = angles_in[2] + parameters[2] + angles = np.array([phi, alpha, beta]) + residual, s = align_obj.gradient(rec, angles, translations) + + vary_parameter = np.array([True, False, True, False, False, True]) + s = s[vary_parameter] + + if scale_factor is None: + scale_factor = np.array([1.0, 1.0, 1.0]) + + s = s * scale_factor[:, np.newaxis] + + if return_vector: + return s.T + + grad = np.dot(s, residual) + + return grad diff --git a/utilities/projection_operators.py b/utilities/projection_operators.py new file mode 100644 index 0000000..1579647 --- /dev/null +++ b/utilities/projection_operators.py @@ -0,0 +1,122 @@ +import numpy as np +from scipy import sparse +import sys +from copy import deepcopy +from utilities.ray_voxel_utilities import forward_sparse as ray_forward_sparse +from utilities.ray_voxel_utilities import forward_proj_grad as ray_forward_proj_grad +from utilities.voxel_utilities import forward_sparse as vox_forward_sparse +from utilities.voxel_utilities import forward_proj_grad as vox_forward_proj_grad + + +class ProjectionMatrix(object): + + def __init__(self, geometry, precision=np.float32): + + self.geometry = geometry + self.precision = precision + self.n_proj = None + self.angles = None + self.xyz_shift = None + self.voxel_mask = None + + def projection_matrix(self, alpha=None, beta=None, phi=None, xyz_shift=None, voxel_mask=None): + + if phi is None: + self.n_proj = self.geometry.n_proj + phi = np.linspace(0., np.pi, self.n_proj) + else: + self.n_proj = np.size(phi) + + if alpha is None: + alpha = np.zeros_like(phi) + + if beta is None: + beta = np.zeros_like(phi) + + # translational shift + if xyz_shift is None: + xyz_shift = np.zeros((self.n_proj, 3)) + + phi = np.squeeze(phi) + alpha = np.squeeze(alpha) + beta = np.squeeze(beta) + xyz_shift = np.squeeze(xyz_shift) + if self.n_proj == 1: + phi = np.array([phi]) + alpha = np.array([alpha]) + beta = np.array([beta]) + xyz_shift = np.array([xyz_shift]) + + self.angles = np.array([phi, alpha, beta]).T + self.xyz_shift = xyz_shift + self.voxel_mask = voxel_mask + + weights, detector_inds, data_inds = self._forward_ray() + + weights = np.concatenate(weights) + detector_inds = np.concatenate(detector_inds) + data_inds = np.concatenate(data_inds) + + if self.voxel_mask is not None: + voxel_mask = self.voxel_mask.ravel().astype(bool) + mask = voxel_mask[data_inds] + if np.sum(mask) == 0: + print('entire object is masked') + weights *= 0.0 + else: + data_inds = data_inds[mask] + #data_inds = _rank_order(data_inds)[0] + detector_inds = detector_inds[mask] + weights = weights[mask] + + # note that weights at duplicate (row, col) pair are summed + proj_mat = sparse.coo_matrix((weights, (detector_inds, data_inds)), + shape=(self.n_proj * self.geometry.n_det, self.geometry.n_vox)) + + return sparse.csr_matrix(proj_mat) + + def _forward_voxel(self): + + weights, detector_inds, data_inds = [], [], [] + for iproj in range(self.n_proj): + # rotate centers for this set of angles and translations + this_geo = deepcopy(self.geometry) + this_geo.cor_shift = self.geometry.cor_shift[iproj] + phi, alpha, beta = self.angles[iproj] + xyz_shift = self.xyz_shift[iproj] + dat_inds, inds, wts = vox_forward_sparse(this_geo, alpha, beta, phi, xyz_shift) + + weights.append(wts.astype(self.precision, copy=False)) + detector_inds.append((inds + iproj * self.geometry.n_det).astype(np.int32)) + data_inds.append(dat_inds) + + return weights, detector_inds, data_inds + + def _forward_ray(self): + + weights, detector_inds, data_inds = [], [], [] + for iproj in range(self.n_proj): + phi, alpha, beta = self.angles[iproj] + xyz_shift = self.xyz_shift[iproj] + this_geo = deepcopy(self.geometry) + this_geo.cor_shift = self.geometry.cor_shift[iproj] + + dat_inds, det_inds, wts = ray_forward_sparse(this_geo, alpha, beta, phi, xyz_shift) + + weights.append(wts.astype(self.precision, copy=False)) + data_inds.append(dat_inds.astype(np.int32)) + detector_inds.append(det_inds + iproj * self.geometry.n_det) + + return weights, detector_inds, data_inds + + def projection_gradient(self, rec, alpha, beta, phi, xyz_shift, cor_shift): + + this_geo = deepcopy(self.geometry) + this_geo.cor_shift = cor_shift + + proj_img, gradient = ray_forward_proj_grad(this_geo, alpha, beta, phi, xyz_shift, rec) + + proj_img = proj_img.astype(self.precision, copy=False) + gradient = gradient.astype(self.precision, copy=False) + + return proj_img.ravel(), gradient.reshape(6, -1) diff --git a/utilities/ray_voxel_utilities.py b/utilities/ray_voxel_utilities.py new file mode 100644 index 0000000..1ed851b --- /dev/null +++ b/utilities/ray_voxel_utilities.py @@ -0,0 +1,345 @@ +import numpy as np +from utilities.rotations import rot_x, rot_y, rot_z, der_rot_x, der_rot_y, der_rot_z +from src import ray_wt_grad + + +def transform_points(x, alpha, beta, phi, t): + + rot_pa = np.dot(rot_z(phi), rot_x(alpha)) + xp = np.dot(rot_y(beta), x) + t[:, np.newaxis] + xp = np.dot(rot_pa, xp) + + return xp + + +def derivative_ray_points(source_points, ray_vector, alpha, beta, phi, xyz_shift): + """ + Rx --> R_phi * R_alpha * (R_beta x + t) + p_j = Rs + j*r_hat where r_hat = (Rd-Rs)/|d-s| + :param source_points: (ub-transformed) source points for this ray + :param ray_vector: (un-transformed) ray-vector = detector_points - source_points + :param angles: rotation angles + :param xyz_shift: translations + :return: ndarray with derivatives (6, 3) + """ + R_p = rot_z(phi) + R_a = rot_x(alpha) + R_b = rot_y(beta) + dR_p = der_rot_z(phi) + dR_a = der_rot_x(alpha) + dR_b = der_rot_y(beta) + R_pa = np.dot(R_p, R_a) + R_ab = np.dot(R_a, R_b) + + der = np.zeros((9, 3, source_points.shape[1])) + # we have 9 components here since we separate the derivative wrt angles into a part that does not depend on + # step/ray_length and that which does + # example derivative wrt phi = dR_p R_a (R_b s + xyz_shift) + (step/ray_length) * dR_p R_a R_b (d-s) + rpa_dot_eye = np.dot(R_pa, np.eye(3)) + der[0] = rpa_dot_eye[:, 0][:, np.newaxis] # d/dt_x + der[1] = rpa_dot_eye[:, 1][:, np.newaxis] # d/dt_y + der[2] = rpa_dot_eye[:, 2][:, np.newaxis] # d/dt_z + + Rb_st = np.dot(R_b, source_points) + xyz_shift[:, np.newaxis] + der[3] = np.dot(dR_p, np.dot(R_a, Rb_st)) + der[4] = np.dot(R_p, np.dot(dR_a, Rb_st)) + der[5] = np.dot(R_pa, np.dot(dR_b, source_points)) + der[6] = np.dot(dR_p, np.dot(R_ab, ray_vector))[:, np.newaxis] + der[7] = np.dot(R_p, np.dot(dR_a, np.dot(R_b, ray_vector)))[:, np.newaxis] + der[8] = np.dot(R_pa, np.dot(dR_b, ray_vector))[:, np.newaxis] + return der + + +def forward_sparse(geometry, alpha, beta, phi, xyz_shift): + """ + Compute detector and data indices and trilinear interpolation weights for projection + at angle phi with jitter angles alpha and beta, and jitter translations xyz. The indices + and weights are used to make a sparse matrix for forward projection. + :param geometry: geometry object with detector and sample information + :param alpha: rotation about X-axis + :param beta: rotation about Y-axis + :param phi: tomographic rotation about Z-axis + :param xyz_shift: Translational jitter + :return: dat_ind, det_ind, wts + Create sparse matrix by sparse.coo_matrix((wts, (det_inds, dat_inds)), shape=(geo.n_det, geo.n_vox)) + Using sparse matrix create projection by sparse.csr_matrix.dot(pmat, 3d_obj.ravel()) + This function calls fortran function trilinear_ray_sparse for fast computation. + """ + N = geometry.vox_shape + n_rays = geometry.n_det + vox_ds = geometry.vox_ds + + geometry.source_centers[0, :] += geometry.cor_shift[0] + geometry.det_centers[0, :] += geometry.cor_shift[0] + p0 = transform_points(geometry.source_centers, alpha, beta, phi, xyz_shift) - geometry.vox_origin[:, np.newaxis] + p1 = transform_points(geometry.det_centers, alpha, beta, phi, xyz_shift) - geometry.vox_origin[:, np.newaxis] + + # apply center of rotation shift + #rot_phi_cor = np.dot(rot_z(phi), geometry.cor_shift) + ##rot_phi_cor[1] = 0.0 + #p0 += rot_phi_cor[:, np.newaxis] + #p1 += rot_phi_cor[:, np.newaxis] + + # now ray trace + step_size = geometry.step_size + r = p1 - p0 + r_length = np.linalg.norm(r, axis=0) + r_hat = r / r_length + n = int(r_length[0] / step_size) # how many points on the ray + r_points = np.zeros((3, n_rays, n)) + r_points[:, :, :] = p0[:, :, np.newaxis] + step = np.zeros((n_rays, n)) + for j in range(n): + r_points[:, :, j] += j * step_size * r_hat + step[:, j] = j * step_size / r_length[0] + + floor_points = np.array([np.floor(r_points[dim]) for dim in range(3)]).astype(np.int32) + + w_ceil = r_points - floor_points.astype(np.float64) + w_floor = 1. - w_ceil + + # fortran routine computes weights in double precision + dat_inds, det_inds, wts, n_inds = ray_wt_grad.trilinear_ray_sparse(np.asfortranarray(floor_points.astype(np.int32)), + np.asfortranarray(w_floor.astype(np.float64)), + N[0], N[1], N[2], n_rays, n) + + dat_inds = dat_inds[:n_inds] + det_inds = det_inds[:n_inds] + wts = wts[:n_inds] + + return dat_inds, det_inds, wts + + +def forward_proj_grad(geometry, alpha, beta, phi, xyz_shift, rec): + """ + Same as forward_sparse, except this function directly computes the forward projection + and the gradient of this projection wrt parameters (alpha, beta, phi, xyz). + :param geometry: + :param alpha: + :param beta: + :param phi: + :param xyz_shift: + :param rec: + :return: + """ + N = geometry.vox_shape + n_rays = geometry.n_det + vox_ds = geometry.vox_ds + + geometry.source_centers[0, :] += geometry.cor_shift[0] + geometry.det_centers[0, :] += geometry.cor_shift[0] + + p0 = transform_points(geometry.source_centers, alpha, beta, phi, xyz_shift) - geometry.vox_origin[:, np.newaxis] + p1 = transform_points(geometry.det_centers, alpha, beta, phi, xyz_shift) - geometry.vox_origin[:, np.newaxis] + + # apply center of rotation shift + #rot_phi_cor = np.dot(rot_z(phi), geometry.cor_shift) + #p0 += rot_phi_cor[:, np.newaxis] + #p1 += rot_phi_cor[:, np.newaxis] + + # now ray trace + step_size = geometry.step_size + r = p1 - p0 + r_length = np.linalg.norm(r, axis=0) + r_hat = r / r_length + n = int(r_length[0] / step_size) # how many points on the ray + r_points = np.zeros((3, n_rays, n)) + r_points[:, :, :] = p0[:, :, np.newaxis] + step = np.zeros((n_rays, n)) + for j in range(n): + r_points[:, :, j] += j * step_size * r_hat + step[:, j] = j * step_size / r_length[0] + + floor_points = np.array([np.floor(r_points[dim]) for dim in range(3)]).astype(np.int32) + + w_ceil = r_points - floor_points.astype(np.float64) + w_floor = 1. - w_ceil + + untransformed_ray = geometry.det_centers - geometry.source_centers + der = derivative_ray_points(geometry.source_centers, untransformed_ray[:, 0], alpha, beta, phi, xyz_shift) + + # fortran routine computes det_img and grad_det_img in double precision + # input 3d rec also passed in in double precision + # integers are int32 + det_img, grad_det_img = ray_wt_grad.trilinear_ray_interp(np.asfortranarray(floor_points), + np.asfortranarray(w_floor), + N[0], N[1], N[2], n_rays, n, + np.asfortranarray(rec.ravel().astype(np.float64)), + np.asfortranarray(step.astype(np.float64)), + np.asfortranarray(der.astype(np.float64))) + return det_img, grad_det_img + + +def ray_tracing_trilinear(vox_shape, p0, p1, vox_ds, step_size, precision=np.float32, return_der=False): + """ + March along parallel rays and compute trilinear interpolation weights at each point on ray. + Since rays are parallel, r_hat and r_length will be same for all. + :param vox_shape: shape of object pierced by the rays + :param p0: source ray-points + :param p1: detector points + :param vox_ds: downsampling of object wrt original voxel size + :param step_size: step size to increment points + :param precision: precision for interpolation weights + :param return_der: if weight derivatives should be returned (needed for parameter optimization) + :return wts: trilinear interpolations weights + :return det_inds: mapped detector indices + :return dat_inds: corners of voxel indices corresponding within which sampled point lies + """ + + N = vox_shape + n_rays = p0.shape[1] + r = p1 - p0 + r_length = np.linalg.norm(r, axis=0) + r_hat = r/r_length + n = int(r_length[0]/step_size) # how many points on the ray + r_points = np.zeros((3, n_rays, n)) + + r_points[:, :, :] = p0[:, :, np.newaxis] + for j in range(n): + r_points[:, :, j] += j * step_size * r_hat + + floor_points = np.array([np.floor(r_points[dim]) for dim in range(3)]).astype(np.int32) + ceil_points = floor_points + 1 + + w_ceil = r_points - floor_points.astype(np.float32) + w_floor = 1. - w_ceil + + in_0 = np.logical_and(floor_points[0] >= 0, ceil_points[0] < vox_ds[0] * N[0]) + in_1 = np.logical_and(floor_points[1] >= 0, ceil_points[1] < vox_ds[1] * N[1]) + in_2 = np.logical_and(floor_points[2] >= 0, ceil_points[2] < vox_ds[2] * N[2]) + in_volume = in_0 * in_1 * in_2 + + ray_ind = np.where(in_volume)[0] + pixel_ind = np.where(in_volume)[1] + floor_points = floor_points[:, ray_ind, pixel_ind] + ceil_points = ceil_points[:, ray_ind, pixel_ind] + norm_factor = np.linalg.norm(r_hat[:, ray_ind], axis=0) ** (1 / 3.) + w_floor = w_floor[:, ray_ind, pixel_ind] * norm_factor + w_ceil = w_ceil[:, ray_ind, pixel_ind] * norm_factor + + det_inds = np.hstack((ray_ind, ray_ind, ray_ind, ray_ind, ray_ind, ray_ind, ray_ind, ray_ind)) + + wts = np.hstack((w_floor[0] * w_floor[1] * w_floor[2], + w_floor[0] * w_floor[1] * w_ceil[2], + w_floor[0] * w_ceil[1] * w_floor[2], + w_floor[0] * w_ceil[1] * w_ceil[2], + w_ceil[0] * w_floor[1] * w_floor[2], + w_ceil[0] * w_floor[1] * w_ceil[2], + w_ceil[0] * w_ceil[1] * w_floor[2], + w_ceil[0] * w_ceil[1] * w_ceil[2])) + + # now divide by step to get data indices + floor_points = np.array([floor_points[dim] / vox_ds[dim] for dim in range(3)]).astype(np.int32) + ceil_points = np.array([ceil_points[dim] / vox_ds[dim] for dim in range(3)]).astype(np.int32) + + inds = np.hstack(((floor_points[0] * N[1] + floor_points[1]) * N[2] + floor_points[2], + (floor_points[0] * N[1] + floor_points[1]) * N[2] + ceil_points[2], + (floor_points[0] * N[1] + ceil_points[1]) * N[2] + floor_points[2], + (floor_points[0] * N[1] + ceil_points[1]) * N[2] + ceil_points[2], + (ceil_points[0] * N[1] + floor_points[1]) * N[2] + floor_points[2], + (ceil_points[0] * N[1] + floor_points[1]) * N[2] + ceil_points[2], + (ceil_points[0] * N[1] + ceil_points[1]) * N[2] + floor_points[2], + (ceil_points[0] * N[1] + ceil_points[1]) * N[2] + ceil_points[2])) + + der_wts = None + if return_der: + der_wts = np.hstack((np.array([-w_floor[1] * w_floor[2], -w_floor[0] * w_floor[2], -w_floor[0] * w_floor[1]]), + np.array([-w_floor[1] * w_ceil[2], -w_floor[0] * w_ceil[2], w_floor[0] * w_floor[1]]), + np.array([-w_ceil[1] * w_floor[2], w_floor[0] * w_floor[2], -w_floor[0] * w_ceil[1]]), + np.array([-w_ceil[1] * w_ceil[2], w_floor[0] * w_ceil[2], w_floor[0] * w_ceil[1]]), + np.array([w_floor[1] * w_floor[2], -w_ceil[0] * w_floor[2], -w_ceil[0] * w_floor[1]]), + np.array([w_floor[1] * w_ceil[2], -w_ceil[0] * w_ceil[2], w_ceil[0] * w_floor[1]]), + np.array([w_ceil[1] * w_floor[2], w_ceil[0] * w_floor[2], -w_ceil[0] * w_ceil[1]]), + np.array([w_ceil[1] * w_ceil[2], w_ceil[0] * w_ceil[2], w_ceil[0] * w_ceil[1]]))) + der_wts = der_wts.astype(precision, copy=False) + + return wts.astype(precision, copy=False), det_inds, inds, der_wts + + +def ray_weights_der(p0, p1, geometry, angles, xyz_shift, rec, return_der=True): + + N = geometry.vox_shape + n_rays = geometry.n_det + vox_ds = geometry.vox_ds + phi, alpha, beta = angles + + # now ray trace + step_size = geometry.step_size + r = p1 - p0 + r_length = np.linalg.norm(r, axis=0) + r_hat = r / r_length + n = int(r_length[0] / step_size) # how many points on the ray + r_points = np.zeros((3, n_rays, n)) + r_points[:, :, :] = p0[:, :, np.newaxis] + step = np.zeros((n_rays, n)) + for j in range(n): + r_points[:, :, j] += j * step_size * r_hat + step[:, j] = j*step_size/r_length[0] + + floor_points = np.array([np.floor(r_points[dim]) for dim in range(3)]).astype(np.int32) + ceil_points = floor_points + 1 + + w_ceil = r_points - floor_points.astype(np.float32) + w_floor = 1. - w_ceil + + in_0 = np.logical_and(floor_points[0] >= 0, ceil_points[0] < vox_ds[0] * N[0]) + in_1 = np.logical_and(floor_points[1] >= 0, ceil_points[1] < vox_ds[1] * N[1]) + in_2 = np.logical_and(floor_points[2] >= 0, ceil_points[2] < vox_ds[2] * N[2]) + in_volume = in_0 * in_1 * in_2 + + ray_ind = np.where(in_volume)[0] + pixel_ind = np.where(in_volume)[1] + step = step[in_volume] + floor_points = floor_points[:, ray_ind, pixel_ind] + ceil_points = ceil_points[:, ray_ind, pixel_ind] + w_floor = w_floor[:, ray_ind, pixel_ind] + w_ceil = w_ceil[:, ray_ind, pixel_ind] + + wts = np.hstack((w_floor[0] * w_floor[1] * w_floor[2], + w_floor[0] * w_floor[1] * w_ceil[2], + w_floor[0] * w_ceil[1] * w_floor[2], + w_floor[0] * w_ceil[1] * w_ceil[2], + w_ceil[0] * w_floor[1] * w_floor[2], + w_ceil[0] * w_floor[1] * w_ceil[2], + w_ceil[0] * w_ceil[1] * w_floor[2], + w_ceil[0] * w_ceil[1] * w_ceil[2])) + + inds = np.hstack(((floor_points[0] * N[1] + floor_points[1]) * N[2] + floor_points[2], + (floor_points[0] * N[1] + floor_points[1]) * N[2] + ceil_points[2], + (floor_points[0] * N[1] + ceil_points[1]) * N[2] + floor_points[2], + (floor_points[0] * N[1] + ceil_points[1]) * N[2] + ceil_points[2], + (ceil_points[0] * N[1] + floor_points[1]) * N[2] + floor_points[2], + (ceil_points[0] * N[1] + floor_points[1]) * N[2] + ceil_points[2], + (ceil_points[0] * N[1] + ceil_points[1]) * N[2] + floor_points[2], + (ceil_points[0] * N[1] + ceil_points[1]) * N[2] + ceil_points[2])) + + der_wts = np.hstack((np.array([-w_floor[1] * w_floor[2], -w_floor[0] * w_floor[2], -w_floor[0] * w_floor[1]]), + np.array([-w_floor[1] * w_ceil[2], -w_floor[0] * w_ceil[2], w_floor[0] * w_floor[1]]), + np.array([-w_ceil[1] * w_floor[2], w_floor[0] * w_floor[2], -w_floor[0] * w_ceil[1]]), + np.array([-w_ceil[1] * w_ceil[2], w_floor[0] * w_ceil[2], w_floor[0] * w_ceil[1]]), + np.array([w_floor[1] * w_floor[2], -w_ceil[0] * w_floor[2], -w_ceil[0] * w_floor[1]]), + np.array([w_floor[1] * w_ceil[2], -w_ceil[0] * w_ceil[2], w_ceil[0] * w_floor[1]]), + np.array([w_ceil[1] * w_floor[2], w_ceil[0] * w_floor[2], -w_ceil[0] * w_ceil[1]]), + np.array([w_ceil[1] * w_ceil[2], w_ceil[0] * w_ceil[2], w_ceil[0] * w_ceil[1]]))) + + inds = inds.reshape(8, -1) + wts = wts.reshape(8, -1) + der_wts = der_wts.reshape(3, 8, -1) + gradient = np.zeros((6, n_rays)) + proj = np.zeros((n_rays, )) + untransformed_ray = geometry.det_centers - geometry.source_centers + der = derivative_ray_points(geometry.source_centers, untransformed_ray[:, 0], alpha, beta, phi, xyz_shift) + for i, ri in enumerate(ray_ind): + # get derivatives of points wrt angles and translations + g = np.zeros((6, 3)) + g[:3, :] = der[:3, :, ri] + g[3, :] = der[3, :, ri] + step[i]*der[6, :, ri] + g[4, :] = der[4, :, ri] + step[i]*der[7, :, ri] + g[5, :] = der[5, :, ri] + step[i]*der[8, :, ri] + g0 = np.sum(der_wts[0, :, i]*rec.ravel()[inds[:, i]]) + g1 = np.sum(der_wts[1, :, i]*rec.ravel()[inds[:, i]]) + g2 = np.sum(der_wts[2, :, i]*rec.ravel()[inds[:, i]]) + gradient[:, ri] += (g0*g[:, 0] + g1*g[:, 1] + g2*g[:, 2]) + proj[ri] += np.sum(rec.ravel()[inds[:, i]]*wts[:, i]) + + return proj, gradient diff --git a/utilities/voxel_utilities.py b/utilities/voxel_utilities.py new file mode 100644 index 0000000..a05bc44 --- /dev/null +++ b/utilities/voxel_utilities.py @@ -0,0 +1,108 @@ +import numpy as np +from utilities.rotations import rot_x, rot_y, rot_z, der_rot_x, der_rot_y, der_rot_z +from src import vox_wt_grad + + +def rigid_transformation(x, alpha, beta, phi, xyz): + # compute rigid body transformation given alignment parameters + # taking into account center-of-rotation shift + # Ax --> R_b (R_a R_t(x) + s). + # Center-of-rotation shift c affects derivatives wrt alpha, beta and theta only. + + R_b = rot_y(beta) + R_a = rot_x(alpha) + R_t = rot_z(phi) + + rtx = np.dot(R_t, x) + ratx = np.dot(R_a, rtx) + ax = np.dot(R_b, ratx + xyz[:, np.newaxis]) + + return ax + + +def derivative_rigid(x, a, b, t, s): + # compute derivatives of rigid body transformation wrt alignment parameters + # taking into account center-of-rotation shift + # Ax --> R_b (R_a(R_t(x + c) - c) + t). + # Center-of-rotation shift c affects derivatives wrt alpha, beta and theta only. + + R_b = rot_y(b) + R_a = rot_x(a) + R_t = rot_z(t) + dR_b = der_rot_y(b) + dR_a = der_rot_x(a) + dR_t = der_rot_z(t) + + rtx = np.dot(R_t, x) + ratx = np.dot(R_a, rtx) + rba = np.dot(R_b, R_a) + # order of derivatives in matrix sx, sy, sz, t, a, b, cx + der = np.zeros((6, x.shape[0], x.shape[1])) + der[0] = np.dot(R_b, np.array([1.0, 0.0, 0.0]))[:, np.newaxis] # = R_b[:, 0] + der[1] = np.dot(R_b, np.array([0.0, 1.0, 0.0]))[:, np.newaxis] # = R_b[:, 1] + der[2] = np.dot(R_b, np.array([0.0, 0.0, 1.0]))[:, np.newaxis] # = R_b[:, 2] + der[3] = np.dot(rba, np.dot(dR_t, x)) + der[4] = np.dot(R_b, np.dot(dR_a, rtx)) + der[5] = np.dot(dR_b, ratx + s[:, np.newaxis]) + + return der + + +def forward_sparse(geometry, alpha, beta, phi, xyz_shift): + """ + Compute bilinear interpolation weights for rotated voxel centers projected + along the y-axis onto planar detector. This function prodces detector and + data indices, and interpolation weights that are used to make a sparse + forward projection matrix. + """ + + det_shape = geometry.det_shape + rot_centers = rigid_transformation(geometry.vox_centers, alpha, beta, phi, xyz_shift) + orig = geometry.vox_origin - geometry.cor_shift + dx = geometry.vox_ds + + floor_x = np.floor((rot_centers[0] - orig[0]) / dx[0]).astype(np.int32) + floor_z = np.floor((rot_centers[2] - orig[2]) / dx[2]).astype(np.int32) + alpha_x = ((rot_centers[0] - orig[0] - floor_x * dx[0]) / dx[0]).astype(np.float32) + alpha_z = ((rot_centers[2] - orig[2] - floor_z * dx[2]) / dx[2]).astype(np.float32) + + dat_inds, det_inds, wts, n_inds = vox_wt_grad.bilinear_sparse(geometry.n_vox, + np.asfortranarray(floor_x), + np.asfortranarray(floor_z), + np.asfortranarray(alpha_x), + np.asfortranarray(alpha_z), + det_shape[0], det_shape[1]) + dat_inds = dat_inds[:n_inds] + det_inds = det_inds[:n_inds] + wts = wts[:n_inds] + + return dat_inds, det_inds, wts + + +def forward_proj_grad(geometry, alpha, beta, phi, xyz_shift, rec): + """ + Same as forward_sparse, except this function directly computes the projection + and gradients of the projection wrt to parameters alpha, beta, phi, xyz. + """ + + det_shape = geometry.det_shape + rot_centers = rigid_transformation(geometry.vox_centers, alpha, beta, phi, xyz_shift) + orig = geometry.vox_origin - geometry.cor_shift + dx = geometry.vox_ds + + floor_x = np.floor((rot_centers[0] - orig[0]) / dx[0]).astype(np.int32) + floor_z = np.floor((rot_centers[2] - orig[2]) / dx[2]).astype(np.int32) + alpha_x = ((rot_centers[0] - orig[0] - floor_x * dx[0]) / dx[0]).astype(np.float32) + alpha_z = ((rot_centers[2] - orig[2] - floor_z * dx[2]) / dx[2]).astype(np.float32) + + der = derivative_rigid(geometry.vox_centers, alpha, beta, phi, xyz_shift) + der = der.astype(np.float32) + + det_img, gradient = vox_wt_grad.bilinear_vox_interp(geometry.n_vox, np.asfortranarray(floor_x), + np.asfortranarray(floor_z), + np.asfortranarray(alpha_x), + np.asfortranarray(alpha_z), + np.asfortranarray(rec.ravel()), + det_shape[0], det_shape[1], + np.asfortranarray(der)) + return det_img.ravel(), gradient.reshape(6, -1)