Skip to content

Commit

Permalink
Merge pull request #375 from CiolJegou/master
Browse files Browse the repository at this point in the history 
 TensorMap slow import #352
  • Loading branch information
@jadball
jadball authored Jan 14, 2025
2 parents 5ee0fc4 + fdc8596 commit 6ed46f3
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 36 deletions.
48 changes: 24 additions & 24 deletions ImageD11/sinograms/point_by_point.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@

# GOTO - find somewhere!
# Everything written in Numba could move to C?
@numba.njit(boundscheck=True)
@numba.njit(boundscheck=True, cache = True)
def hkluniq(ubi, gx, gy, gz, eta, m, tol, hmax):
"""count uniq peaks - move to cImageD11 in the future"""
# index from -hmax to hmax
Expand Down Expand Up @@ -82,7 +82,7 @@ def hkluniq(ubi, gx, gy, gz, eta, m, tol, hmax):

# stuff we need to compute g-vectors
# from ImageD11.transform
@numba.njit
@numba.njit(cache = True)
def detector_rotation_matrix(tilt_x, tilt_y, tilt_z):
r1 = np.array([[np.cos(tilt_z), -np.sin(tilt_z), 0.0], # note this is r.h.
[np.sin(tilt_z), np.cos(tilt_z), 0.0],
Expand All @@ -97,7 +97,7 @@ def detector_rotation_matrix(tilt_x, tilt_y, tilt_z):
return r2r1


@numba.njit
@numba.njit(cache = True)
def compute_xyz_lab(sc, fc,
y_center=0., y_size=0., tilt_y=0.,
z_center=0., z_size=0., tilt_z=0.,
Expand Down Expand Up @@ -132,7 +132,7 @@ def compute_xyz_lab(sc, fc,
return rotvec


@numba.njit
@numba.njit(cache=True)
def compute_tth_eta_from_xyz(peaks_xyz,
t_x=0.0, t_y=0.0, t_z=0.0,
wedge=0.0, # Wedge == theta on 4circ
Expand All @@ -147,7 +147,7 @@ def compute_tth_eta_from_xyz(peaks_xyz,
return tth, eta


@numba.njit
@numba.njit(cache=True)
def compute_tth_eta(sc, fc,
y_center=0., y_size=0., tilt_y=0.,
z_center=0., z_size=0., tilt_z=0.,
Expand Down Expand Up @@ -175,7 +175,7 @@ def compute_tth_eta(sc, fc,
return tth, eta


@numba.njit
@numba.njit(cache=True)
def compute_k_vectors(tth, eta, wvln):
"""
generate k vectors - scattering vectors in laboratory frame
Expand All @@ -195,7 +195,7 @@ def compute_k_vectors(tth, eta, wvln):
return k


@numba.njit
@numba.njit(cache=True)
def compute_g_from_k(k, omega, wedge=0, chi=0):
"""
Compute g-vectors with cached k-vectors
Expand Down Expand Up @@ -226,7 +226,7 @@ def compute_g_from_k(k, omega, wedge=0, chi=0):
return g


@numba.njit
@numba.njit(cache=True)
def compute_g_vectors(tth,
eta,
omega,
Expand All @@ -243,7 +243,7 @@ def compute_g_vectors(tth,
return compute_g_from_k(k, omega, wedge, chi)


@numba.njit
@numba.njit(cache=True)
def count_unique_peaks(hkl, etasign, dtyi):
N = hkl.shape[1] # Number of entries
indices = np.zeros(N, dtype=np.int64)
Expand All @@ -270,7 +270,7 @@ def count_unique_peaks(hkl, etasign, dtyi):
return indices


@numba.njit
@numba.njit(cache=True)
def count_unique_peaks_no_dtyi(hkl, etasign):
N = hkl.shape[1] # Number of entries
indices = np.zeros(N, dtype=np.int64)
Expand All @@ -297,7 +297,7 @@ def count_unique_peaks_no_dtyi(hkl, etasign):
return indices


@numba.njit
@numba.njit(cache=True)
def merge(hkl, etasign, dtyi, sum_intensity, sc, fc, omega, dty, xpos_refined, eta):
"""
merge peaks with the same (h, k, l, etasign, dtyi)
Expand All @@ -318,7 +318,7 @@ def merge(hkl, etasign, dtyi, sum_intensity, sc, fc, omega, dty, xpos_refined, e
return merged_sum_intensity, merged_sc, merged_fc, merged_omega, merged_dty, merged_xpos_refined, merged_eta


@numba.njit
@numba.njit(cache=True)
def get_voxel_idx(y0, xi0, yi0, sinomega, cosomega, dty, ystep):
"""
get peaks at xi0, yi0
Expand All @@ -331,7 +331,7 @@ def get_voxel_idx(y0, xi0, yi0, sinomega, cosomega, dty, ystep):
return idx, ydist


@numba.njit
@numba.njit(cache=True)
def compute_gve(sc, fc, omega, xpos,
distance, y_center, y_size, tilt_y, z_center, z_size, tilt_z, tilt_x,
o11, o12, o21, o22,
Expand Down Expand Up @@ -367,7 +367,7 @@ def compute_gve(sc, fc, omega, xpos,
return gve


@numba.njit
@numba.njit(cache=True)
def weighted_lstsq_ubi_fit(ydist, gve, hkl):
# run the weighted fit
# a.T @ gve = h => gve.T @ a = h.T => a = np.linalg.pinv(gve.T) @ h.T, same for b and c
Expand All @@ -382,7 +382,7 @@ def weighted_lstsq_ubi_fit(ydist, gve, hkl):
return w, ubifit, residuals, rank, sing_vals


@numba.njit
@numba.njit(cache=True)
def gve_norm(gve):
norms = np.zeros(gve.shape[1])
for i in range(gve.shape[1]):
Expand All @@ -392,7 +392,7 @@ def gve_norm(gve):
return norms


@numba.njit
@numba.njit(cache=True)
def divide_where(arr1, arr2, out, wherearr):
"""
Do arr1/arr2.
Expand All @@ -402,7 +402,7 @@ def divide_where(arr1, arr2, out, wherearr):
return np.where(wherearr != 0, div, out)


@numba.njit
@numba.njit(cache=True)
def ubi_to_unitcell(ubi):
# fast numba version, can't use guvec version from tensor_map.py here unfortunately
mt = np.dot(ubi, ubi.T)
Expand All @@ -414,7 +414,7 @@ def ubi_to_unitcell(ubi):
return np.array([a, b, c, al, be, ga])


@numba.njit
@numba.njit(cache=True)
def ubi_and_ucell_to_u(ubi, ucell):
# compute B
a, b, c = ucell[:3]
Expand Down Expand Up @@ -453,7 +453,7 @@ def ubi_and_ucell_to_u(ubi, ucell):
return u


@numba.njit
@numba.njit(cache=True)
def nb_choose_best(i, j, u, n, NY, ubiar,
minpeaks=6):
# map of the unique scores
Expand Down Expand Up @@ -486,7 +486,7 @@ def nb_choose_best(i, j, u, n, NY, ubiar,
return uniq, npk, ubi, best_rows


@numba.njit
@numba.njit(cache=True)
def nb_inv(mats, imats):
for i in range(mats.shape[0]):
for j in range(mats.shape[1]):
Expand All @@ -501,7 +501,7 @@ def nb_inv(mats, imats):
imats[i, j] = 42.


@numba.njit
@numba.njit(cache=True)
def nb_inv_3d(mats, imats):
for i in range(mats.shape[0]):
for j in range(mats.shape[1]):
Expand Down Expand Up @@ -1267,7 +1267,7 @@ def prepare_refine_bash(pbp_object, id11_code_path, output_filename):
return bash_script_path


@numba.njit
@numba.njit(cache=True)
def unique_with_counts(a):
# https://github.com/numba/numba/pull/2959/commits/6657709adcf128d97eaaf8371d9106d48e2360ba
b = np.sort(a.ravel())
Expand All @@ -1282,7 +1282,7 @@ def unique_with_counts(a):
return np.array(unique), np.array(counts)


@numba.njit(parallel=True)
@numba.njit(parallel=True,cache=True)
def compute_origins(singlemap, sample_mask,
gve, sinomega, cosomega, omega, dty,
sx_grid, sy_grid,
Expand Down Expand Up @@ -1385,7 +1385,7 @@ def compute_origins(singlemap, sample_mask,
return lx_modified


@numba.njit(parallel=True)
@numba.njit(parallel=True, cache=True)
def refine_map(refine_points, all_pbpmap_ubis, ri_col, rj_col, sx_grid, sy_grid, mask, # refinement stuff
sc, fc, eta, sum_intensity, sinomega, cosomega, omega, dty, dtyi, xpos, # icolf columns
ystep, y0,
Expand Down
24 changes: 12 additions & 12 deletions ImageD11/sinograms/tensor_map.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
# e.g these can be called and broadcasted for a UBI map of (1,200,500,3,3)

# take in a nxn float64 array, return a nxn float64 array
@numba.guvectorize([(numba.float64[:, :], numba.float64[:, :])], '(n,n)->(n,n)', nopython=True)
@numba.guvectorize([(numba.float64[:, :], numba.float64[:, :])], '(n,n)->(n,n)', nopython=True, cache = True)
def fast_invert(mat, res):
"""
Fast matrix invert.
Expand All @@ -42,7 +42,7 @@ def fast_invert(mat, res):


# take in a nxn float64 array, return a nxn float64 array
@numba.guvectorize([(numba.float64[:, :], numba.float64[:, :])], '(n,n)->(n,n)', nopython=True)
@numba.guvectorize([(numba.float64[:, :], numba.float64[:, :])], '(n,n)->(n,n)', nopython=True, cache = True)
def ubi_to_mt(ubi, res):
"""
Compute metric tensor from UBI matrix, the same way as in ``ImageD11.grain``.
Expand All @@ -64,7 +64,7 @@ def ubi_to_mt(ubi, res):


# take in a nxn float64 array and a dummy variable p-length float64 vector (required by numba), return a p-length float64 vector
@numba.guvectorize([(numba.float64[:, :], numba.float64[:], numba.float64[:])], '(n,n),(p)->(p)', nopython=True)
@numba.guvectorize([(numba.float64[:, :], numba.float64[:], numba.float64[:])], '(n,n),(p)->(p)', nopython=True, cache = True)
def mt_to_unitcell(mt, dum, res):
"""
Get unitcell parameters array ``[a, b, c, alpha, beta, gamma]`` from the metric tensor.
Expand Down Expand Up @@ -103,7 +103,7 @@ def mt_to_unitcell(mt, dum, res):


# take in a p-length float64 vector and a dummy variable nxn float64 array (required by numba), return a nxn float64 array
@numba.guvectorize([(numba.float64[:], numba.float64[:, :], numba.float64[:, :])], '(p),(n,n)->(n,n)', nopython=True)
@numba.guvectorize([(numba.float64[:], numba.float64[:, :], numba.float64[:, :])], '(p),(n,n)->(n,n)', nopython=True, cache = True)
def unitcell_to_b(unitcell, dum, res):
"""
Get B matrix from unitcell, the same way as in ``ImageD11.grain`` (via ``ImageD11.unitcell``).
Expand Down Expand Up @@ -160,7 +160,7 @@ def unitcell_to_b(unitcell, dum, res):

# take in a nxn float64 array, return a nxn float64 array
@numba.guvectorize([(numba.float64[:, :], numba.float64[:, :], numba.float64[:, :])], '(n,n),(n,n)->(n,n)',
nopython=True)
nopython=True, cache = True)
def ubi_and_b_to_u(ubi, b, res):
"""
Get U matrix from UBI and B matrices, the same way as in ``ImageD11.grain``.
Expand All @@ -187,7 +187,7 @@ def ubi_and_b_to_u(ubi, b, res):
res[...] = np.dot(b, ubi).T


@numba.guvectorize([(numba.float64[:, :], numba.float64[:], numba.float64[:, :])], '(n,n),(p)->(n,n)', nopython=True)
@numba.guvectorize([(numba.float64[:, :], numba.float64[:], numba.float64[:, :])], '(n,n),(p)->(n,n)', nopython=True, cache = True)
def ubi_and_unitcell_to_eps_sample(ubi, dzero_cell, res):
"""
Get Biot strain tensor (3x3) in sample reference frame from UBI array and unitcell array.
Expand Down Expand Up @@ -250,7 +250,7 @@ def ubi_and_unitcell_to_eps_sample(ubi, dzero_cell, res):
res[...] = em


@numba.guvectorize([(numba.float64[:, :], numba.float64[:], numba.float64[:, :])], '(n,n),(p)->(n,n)', nopython=True)
@numba.guvectorize([(numba.float64[:, :], numba.float64[:], numba.float64[:, :])], '(n,n),(p)->(n,n)', nopython=True, cache = True)
def ubi_and_unitcell_to_eps_crystal(ubi, dzero_cell, res):
"""
Get Biot strain tensor (3x3) in crystal reference frame from UBI and unitcell.
Expand Down Expand Up @@ -314,7 +314,7 @@ def ubi_and_unitcell_to_eps_crystal(ubi, dzero_cell, res):


@numba.guvectorize([(numba.float64[:, :], numba.float64[:, :], numba.float64[:, :])],
'(n,n),(n,n)->(n,n)', nopython=True)
'(n,n),(n,n)->(n,n)', nopython=True, cache = True)
def tensor_crystal_to_sample(tensor_crystal, U, res):
"""
Rotate tensor from crystal to sample reference frame.
Expand All @@ -340,7 +340,7 @@ def tensor_crystal_to_sample(tensor_crystal, U, res):


@numba.guvectorize([(numba.float64[:, :], numba.float64[:, :], numba.float64[:, :])],
'(n,n),(n,n)->(n,n)', nopython=True)
'(n,n),(n,n)->(n,n)', nopython=True, cache = True)
def tensor_sample_to_crystal(tensor_sample, U, res):
"""
Rotate tensor from sample to crystal reference frame.
Expand All @@ -366,7 +366,7 @@ def tensor_sample_to_crystal(tensor_sample, U, res):


@numba.guvectorize([(numba.float64[:, :], numba.float64[:, :], numba.float64[:, :], numba.float64[:], numba.float64[:, :])],
'(n,n),(k,k),(n,n),()->(n,n)', nopython=True)
'(n,n),(k,k),(n,n),()->(n,n)', nopython=True, cache = True)
def strain_crystal_to_stress_crystal(strain_crystal, stiffness_tensor, B0, phase_mask, res):
"""
Convert Biot strain tensor (ImageD11 coordinate system) to Biot stress tensor (ImageD11 coordinate system).
Expand Down Expand Up @@ -453,7 +453,7 @@ def strain_crystal_to_stress_crystal(strain_crystal, stiffness_tensor, B0, phase
res[...] = stress_crystal


@numba.guvectorize([(numba.float64[:, :], numba.float64[:])], '(n,n)->()', nopython=True)
@numba.guvectorize([(numba.float64[:, :], numba.float64[:])], '(n,n)->()', nopython=True, cache = True)
def sig_to_vm(sig, res):
"""Get von-Mises stress scalar from stress tensor"""
sig11 = sig[0, 0]
Expand Down Expand Up @@ -489,7 +489,7 @@ def _arctan2(y, x):


# take in a nxn float64 array and a dummy variable p-length float64 vector (required by numba), return a p-length float64 vector
@numba.guvectorize([(numba.float64[:, :], numba.float64[:], numba.float64[:])], '(n,n),(p)->(p)', nopython=True)
@numba.guvectorize([(numba.float64[:, :], numba.float64[:], numba.float64[:])], '(n,n),(p)->(p)', nopython=True, cache = True)
def u_to_euler(u, dum, res):
"""Get Euler angles (radians) from U matrix, the same way as in xfab.tools
To use (no need to pre-populate res):
Expand Down

0 comments on commit 6ed46f3

Please sign in to comment.