Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

TensorMap slow import #352 #375

Merged
merged 4 commits into from
Jan 14, 2025
Merged

TensorMap slow import #352 #375

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ecd9b68
test
Jan 13, 2025
9a741ec
Adding cache=True to speed up tensorMap imports
Jan 13, 2025
c335f41
Removing _cImageD11 files and adding cache=True in point_by_point.py
Jan 14, 2025
fdc8596
With pbp modifications...
Jan 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 24 additions & 24 deletions ImageD11/sinograms/point_by_point.py
View file
Open in desktop
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
Open in desktop
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