Author: tlinnet Date: Wed Jun 25 02:14:38 2014 New Revision: 24295 URL: http://svn.gna.org/viewcvs/relax?rev=24295&view=rev Log: Speeded up the code of NS r1rho 3site. This was essential done to numpy einsum, and doing the dot operations in multiple dimensions. It was though necessary to realize, that to do the proper dot product operations, the outer two axis if M0 should be swapped, by rolling the outer axis one back. Task #7807 (https://gna.org/task/index.php?7807): Speed-up of dispersion models for Clustered analysis. Modified: branches/disp_spin_speed/lib/dispersion/ns_r1rho_3site.py branches/disp_spin_speed/target_functions/relax_disp.py Modified: branches/disp_spin_speed/lib/dispersion/ns_r1rho_3site.py URL: http://svn.gna.org/viewcvs/relax/branches/disp_spin_speed/lib/dispersion/ns_r1rho_3site.py?rev=24295&r1=24294&r2=24295&view=diff ============================================================================== --- branches/disp_spin_speed/lib/dispersion/ns_r1rho_3site.py (original) +++ branches/disp_spin_speed/lib/dispersion/ns_r1rho_3site.py Wed Jun 25 02:14:38 2014 @@ -56,11 +56,10 @@ """ # Python module imports. -from math import atan2 -from numpy import array, cos, dot, einsum, float64, log, multiply, sin, sum +from numpy import array, einsum, float64, isfinite, log, min, multiply, sin, sum +from numpy.ma import fix_invalid, masked_less # relax module imports. -from lib.float import isNaN from lib.dispersion.matrix_exponential import matrix_exponential_rank_NE_NS_NM_NO_ND_x_x # Repetitive calculations (to speed up calculations). @@ -332,16 +331,17 @@ return matrix -def ns_r1rho_3site(M0=None, r1rho_prime=None, omega=None, offset=None, r1=0.0, pA=None, pB=None, dw_AB=None, dw_BC=None, kex_AB=None, kex_BC=None, kex_AC=None, spin_lock_fields=None, relax_time=None, inv_relax_time=None, back_calc=None, num_points=None): +def ns_r1rho_3site(M0=None, M0_T=None, r1rho_prime=None, omega=None, offset=None, r1=0.0, pA=None, pB=None, dw_AB=None, dw_BC=None, kex_AB=None, kex_BC=None, kex_AC=None, spin_lock_fields=None, relax_time=None, inv_relax_time=None, back_calc=None, num_points=None): """The 3-site numerical solution to the Bloch-McConnell equation for R1rho data. This function calculates and stores the R1rho values. @keyword M0: This is a vector that contains the initial magnetizations corresponding to the A and B state transverse magnetizations. - @type M0: numpy float64, rank-1, 7D array + @type M0: numpy float array of rank [NE][NS][NM][NO][ND][9][1] + @keyword M0_T: This is a vector that contains the initial magnetizations corresponding to the A and B state transverse magnetizations, where the outer two axis has been swapped for efficient dot operations. @keyword r1rho_prime: The R1rho_prime parameter value (R1rho with no exchange). - @type r1rho_prime: numpy float array of rank [NS][NM][NO][ND] + @type r1rho_prime: numpy float array of rank [NE][NS][NM][NO][ND][1][9] @keyword omega: The chemical shift for the spin in rad/s. @type omega: numpy float array of rank [NS][NM][NO][ND] @keyword offset: The spin-lock offsets for the data. @@ -399,31 +399,20 @@ # Magnetization evolution. Rexpo_M0_mat = einsum('...ij,...jk', Rexpo_mat, M0) - # Loop over spins. - for si in range(NS): - # Loop over the spectrometer frequencies. - for mi in range(NM): - # Loop over offsets: - for oi in range(NO): - # Extract parameters from array. - num_points_i = num_points[0, si, mi, oi] - - # Loop over the time points, back calculating the R2eff values. - for j in range(num_points_i): - # Extract the preformed matrix that rotate the magnetization previous to spin-lock into the weff frame. - M0_i= M0[0, si, mi, oi, j, :, 0] - - # This matrix is a propagator that will evolve the magnetization with the matrix R. - Rexpo_i = Rexpo_mat[0, si, mi, oi, j] - - # Extract from the pre-formed Magnetization evolution matrix. - Rexpo_M0_mat_i = Rexpo_M0_mat[0, si, mi, oi, j, :, 0] - - # Magnetization evolution. - MA = dot(M0_i, Rexpo_M0_mat_i) - - # The next lines calculate the R1rho using a two-point approximation, i.e. assuming that the decay is mono-exponential. - if MA <= 0.0 or isNaN(MA): - back_calc[0, si, mi, oi, j] = 1e99 - else: - back_calc[0, si, mi, oi, j]= -inv_relax_time[0, si, mi, oi, j] * log(MA) + # Magnetization evolution, which include all dimensions. + MA_mat = einsum('...ij,...jk', M0_T, Rexpo_M0_mat)[:, :, :, :, :, 0, 0] + + # Insert safe checks. + if min(MA_mat) < 0.0: + mask_min_MA_mat = masked_less(MA_mat, 0.0) + # Fill with high values. + MA_mat[mask_min_MA_mat.mask] = 1e100 + + # Do back calculation. + back_calc[:] = -inv_relax_time * log(MA_mat) + + # Catch errors, taking a sum over array is the fastest way to check for + # +/- inf (infinity) and nan (not a number). + if not isfinite(sum(back_calc)): + # Replaces nan, inf, etc. with fill value. + fix_invalid(back_calc, copy=False, fill_value=1e100) Modified: branches/disp_spin_speed/target_functions/relax_disp.py URL: http://svn.gna.org/viewcvs/relax/branches/disp_spin_speed/target_functions/relax_disp.py?rev=24295&r1=24294&r2=24295&view=diff ============================================================================== --- branches/disp_spin_speed/target_functions/relax_disp.py (original) +++ branches/disp_spin_speed/target_functions/relax_disp.py Wed Jun 25 02:14:38 2014 @@ -432,6 +432,8 @@ # The A state initial Z magnetisation. M0_cos = multiply.outer( cos(theta_mat), M0_2 ) self.M0 = M0_sin + M0_cos + # Transpose M0, to prepare for dot operation. Roll the last axis one back, corresponds to a transpose for the outer two axis. + self.M0_T = rollaxis(self.M0, 6, 5) # Set up the model. if model == MODEL_NOREX: @@ -772,7 +774,7 @@ self.r20_struct[:] = multiply.outer( r1rho_prime.reshape(self.NE, self.NS, self.NM), self.no_nd_ones ) # Back calculate the R2eff values for each experiment type. - ns_r1rho_3site(M0=self.M0, r1rho_prime=self.r20_struct, omega=self.chemical_shifts, offset=self.offset, r1=self.r1, pA=pA, pB=pB, dw_AB=self.dw_AB_struct, dw_BC=self.dw_BC_struct, kex_AB=kex_AB, kex_BC=kex_BC, kex_AC=kex_AC, spin_lock_fields=self.spin_lock_omega1, relax_time=self.relax_times, inv_relax_time=self.inv_relax_times, back_calc=self.back_calc, num_points=self.num_disp_points) + ns_r1rho_3site(M0=self.M0, M0_T=self.M0_T, r1rho_prime=self.r20_struct, omega=self.chemical_shifts, offset=self.offset, r1=self.r1, pA=pA, pB=pB, dw_AB=self.dw_AB_struct, dw_BC=self.dw_BC_struct, kex_AB=kex_AB, kex_BC=kex_BC, kex_AC=kex_AC, spin_lock_fields=self.spin_lock_omega1, relax_time=self.relax_times, inv_relax_time=self.inv_relax_times, back_calc=self.back_calc, num_points=self.num_disp_points) # Clean the data for all values, which is left over at the end of arrays. self.back_calc = self.back_calc*self.disp_struct