from __future__ import division from scitbx.array_family import flex from scitbx import matrix from scitbx.linalg import eigensystem import operator import math # Exception class NoConvergenceException(RuntimeError): pass # Conversion between rotation representations def quaternion_to_rotmat(q): ( l, m, n, s ) = q return matrix.sqr( [ l*l - m*m - n*n + s*s, 2.0 * ( l*m - n*s ), 2.0 * ( l*n + m*s ), 2.0 * ( l*m + n*s ), -l*l + m*m -n*n + s*s, 2.0 * ( m*n - l*s ), 2.0 * ( l*n - m*s ), 2.0 * ( m*n + l*s ), -l*l - m*m + n*n + s*s, ] ) def quaternion_to_quatmat(q): ( l, m, n, s ) = q.elems return matrix.sqr( [ s, -n, m, l, n, s, -l, m, -m, l, s, n, -l, -m, -n, s, ] ) i_bar = matrix.sqr( [ -1, 0, 0, 0, 0, -1, 0, 0, 0, 0, -1, 0, 0, 0, 0, 1 ] ) # Functions to calculate rotational superposition def best_rotation(moving, reference, weights): return top_quaternion( p = p_matrix_from( a_sites = moving, b_sites = reference, weights = weights ) ) def top_quaternion(p): eigen = eigensystem.real_symmetric( p.as_flex_double_matrix() ) return eigen.vectors()[0:4] def p_matrix_from(a_sites, b_sites, weights): # Diamond, Acta Cryst A44, 211-216 (1988) assert len( a_sites ) == len( b_sites ) and len( a_sites ) == len( weights ) ( xa, ya, za ) = component_arrays( a_sites ) ( xb, yb, zb ) = component_arrays( b_sites ) m = matrix.sqr( [ flex.sum( weights * xa * xb ), flex.sum( weights * xa * yb ), flex.sum( weights * xa * zb ), flex.sum( weights * ya * xb ), flex.sum( weights * ya * yb ), flex.sum( weights * ya * zb ), flex.sum( weights * za * xb ), flex.sum( weights * za * yb ), flex.sum( weights * za * zb ), ] ) q = m + m.transpose() - 2 * matrix.identity(3) * m.trace() # Application of Sum(j, k)[ epsilon( i, j, k ) * m( j, k ) ] # where i,j,k = [ 0, 1, 2 ] and epsilon is the Levi-Civitta symbol v = ( m[5] - m[7], m[6] - m[2], m[1] - m[3] ) return matrix.sqr( [ q[0], q[1], q[2], v[0], q[3], q[4], q[5], v[1], q[6], q[7], q[8], v[2], v[0], v[1], v[2], 0, ] ) # Various vector operations def component_arrays(vec3_array): ( x, y, z ) = zip( *vec3_array ) return ( flex.double( x ), flex.double( y ), flex.double( z ) ) # Non-member functions def iterate(superposition, convergence, tolerance=1.0E-8, iterations = 100): assert 0 < convergence residual = superposition.residual() count = 0 tolerance *= superposition.set_count * superposition.site_count while True: superposition.full_iteration() new_residual = superposition.residual() diff = new_residual - residual if tolerance < diff: raise NoConvergenceException, ( "Divergence in refinement", count, new_residual, ) if abs( diff ) < convergence: break residual = new_residual count += 1 if iterations < count: raise NoConvergenceException, ( "Iteration limit exceeded", count, new_residual, ) return ( count, new_residual ) def weighted_error_between(left, right, weights): differences = left - right return flex.mean_weighted( differences.dot( differences ), weights ) def transform_sites(sites, rotation, translation): return rotation.elems * sites + translation class Matrix(object): """ A matrix """ def __init__(self, n, m): self.n = n self.m = m self.store = [ [ None ] * m for i in range( n ) ] def set(self, left, right, value): assert left < self.n and right < self.m self.store[ left ][ right ] = value def get(self, left, right): assert left < self.n and right < self.m return self.store[ left ][ right ] def rows(self): return self.store def off_diagonal_elements_in_rows(self): for ( index, row ) in enumerate( self.rows() ): yield row[ : index ] + row[ index + 1 : ] class SymmetricMatrix(object): """ A matrix that is symmetric and only half has to be calculated / stored """ def __init__(self, dimension): assert 0 <= dimension self.d = dimension self.store = [ None ] * ( ( ( self.d + 1 ) * self.d ) // 2 ) def set(self, left, right, value): self.store[ self.transform( left = left, right = right ) ] = value def get(self, left, right): return self.store[ self.transform( left = left, right = right ) ] def column(self, index): start1 = self.transform( left = index, right = 0 ) slice1 = self.store[ start1 : start1 + index + 1 ] slice2 = [ self.get( i, index ) for i in range( index + 1, self.d ) ] return slice1 + slice2 def transform(self, left, right): if left < right: ( left, right ) = ( right, left ) return ( ( left + 1 ) * left ) // 2 + right class ErrorOnRotation(object): """ Superposition of two coordinate sets """ def __init__(self, e, p): self.e = e self.p = p def with_rotated_reference(self, q): rho = quaternion_to_quatmat( q = q ) return self.__class__( e = self.e, p = rho * self.p * rho.transpose() ) def for_quaternion(self, q): square = self.e - 2 * ( q.transpose() * self.p * q )[0] if square < 0: print "Negative square: %s" % square return ( square if 0 <= square else 0 ) def best_rotation(self): return matrix.col( top_quaternion( p = self.p ) ) @classmethod def from_sites(cls, left, right, weights): assert len( left ) == len( right ) == len( weights ) diffs = left - right e = flex.sum( diffs.dot( diffs ) * weights ) p = p_matrix_from( a_sites = left, b_sites = right, weights = weights ) return cls( e = e, p = p ) class MultipleSuperposition(object): """ Base class for multiple superpositions """ def __init__(self, site_sets): self.set_count = len( site_sets ) assert 2 <= self.set_count self.site_count = len( site_sets[0] ) assert all( len( s ) == self.site_count for s in site_sets[1:] ) self.site_sets = site_sets def rotations(self): return [ matrix.sqr( quaternion_to_rotmat( q ) ) for q in self.qs ] def transformations(self): rotations = self.rotations() return ( rotations, [ -r * matrix.col( t ) for ( r, t ) in zip( rotations, self.centroids ) ] ) def transformed_sites(self): ( rotations, translations ) = self.transformations() return [ transform_sites( sites = sites, rotation = rot, translation = tra ) for ( sites, rot, tra ) in zip( self.site_sets, rotations, translations ) ] def rmsds(self, weights): transformed = self.transformed_sites() norm = 1.0 / ( self.set_count - 1 ) return [ norm * sum( [ weighted_rmsd_between( left = l, right = r, weights = weights ) for r in transformed ] ) for l in transformed ] class DiamondAlgorithm(MultipleSuperposition): """ Multiple superposition using Diamond's algorithm """ def __init__(self, site_sets, weights): super( DiamondAlgorithm, self ).__init__( site_sets = site_sets ) # Calculate initial rotations self.prepare_new_weights( weights = weights ) self.calculate_pairwise_error_matrix() column = [ r[0] for r in self.pw_matrix.rows() ] self.qs = ( [ matrix.col( [ 0, 0, 0, 1 ] ) ] + [ eor.best_rotation() for eor in column[1:] ] ) self.totals = [ None ] * self.set_count self.calculate_total_error_vector() # Commands def set_new_weights(self): self.calculate_pairwise_error_matrix() self.calculate_total_error_vector() def prepare_new_weights(self, weights): assert len( weights ) == self.site_count if flex.sum( weights ) <= 0: raise RuntimeError, "Zero weights" self.new_weight_data = weights def full_iteration(self): self.qs = [ matrix.col( top_quaternion( p = t.p ) ) for t in self.totals ] self.calculate_total_error_vector() # Queries def rmsd_between(self, left, right): if left == right: return 0.0 eor = ( self.pw_matrix.get( left = left, right = right ) .with_rotated_reference( q = self.qs[ right ] ) ) square = eor.for_quaternion( q = self.qs[ left ] ) / self.norm return math.sqrt( square if 0 <= square else 0 ) def weighted_rmsds(self): norm = 1.0 / ( self.set_count - 1 ) return [ math.sqrt( r / self.norm * norm ) for r in self.residuals() ] def unweighted_rmsds(self): transformed = self.transformed_sites() weights = flex.double( [ 1 ] * self.site_count ) norm = 1.0 / ( self.set_count - 1 ) return [ math.sqrt( norm * sum( [ weighted_error_between( left = l, right = r, weights = weights ) for r in transformed ] ) ) for l in transformed ] def rmsd(self): return math.sqrt( self.residual() / ( self.set_count * ( self.set_count - 1 ) * self.norm ) ) def residual(self): return sum( self.residuals() ) def residuals(self): return [ t.for_quaternion( q = q ) for ( q, t ) in zip( self.qs, self.totals ) ] def average_structure(self): return self.average_structure_from_transformed_sites( transformeds = self.transformed_sites() ) def distance_squares_from_average_structure(self): transformed = self.transformed_sites() average = self.average_structure_from_transformed_sites( transformeds = self.transformed_sites() ) differences = [ t - average for t in transformed ] return 1.0 / ( self.set_count - 1 ) * reduce( operator.add, [ diff.dot( diff ) for diff in differences ] ) # Internal methods def average_structure_from_transformed_sites(self, transformeds): return 1.0 / self.set_count * reduce( operator.add, transformeds ) def calculate_pairwise_error_matrix(self): self.norm = flex.sum( self.new_weight_data ) assert 0 < self.norm # Calculate centroids self.centroids = flex.vec3_double( [ s.mean_weighted( self.new_weight_data ) for s in self.site_sets ] ) origin_sets = [ c - o for ( c, o ) in zip( self.site_sets, self.centroids ) ] # Calculate starting pairwise squared error matrix self.pw_matrix = Matrix( n = self.set_count, m = self.set_count ) for i in range( self.set_count ): for j in range( i + 1, self.set_count ): pw = ErrorOnRotation.from_sites( left = origin_sets[ i ], right = origin_sets[ j ], weights = self.new_weight_data ) self.pw_matrix.set( left = i, right = j, value = pw ) self.pw_matrix.set( left = j, right = i, value = ErrorOnRotation( e = pw.e, p = i_bar * pw.p * i_bar ) ) def calculate_total_error_vector(self): for ( index, row ) in enumerate( self.pw_matrix.rows() ): transformed_off_diagonal = [ eor.with_rotated_reference( q = q ) for ( i, ( eor, q ) ) in enumerate( zip( row, self.qs ) ) if i != index ] sum_e = sum( [ eor.e for eor in transformed_off_diagonal ] ) sum_p = reduce( operator.add, [ eor.p for eor in transformed_off_diagonal ] ) self.totals[ index ] = ErrorOnRotation( e = sum_e, p = sum_p ) class WangSnoeyinkAlgorithm(MultipleSuperposition): """ Multiple structure alignment with gaps using the algorithm from Wang & Snoeyink """ def __init__(self, site_sets, selections, weights): super( WangSnoeyinkAlgorithm, self ).__init__( site_sets = site_sets ) assert len( selections ) == self.set_count assert all( len( s ) == self.site_count for s in selections ) self.selections = selections self.norms = reduce( operator.add, [ s.as_double() for s in self.selections ] ) if not all( 2 <= v for v in self.norms ): raise RuntimeError, "Single positions in supplied site sets" self.qs = [ matrix.col( [ 0, 0, 0, 1 ] ) ] * self.set_count self.origins = [ ( 0, 0, 0 ) ] * self.set_count self.prepare_new_weights( weights = weights ) self.set_new_weights() def set_new_weights(self): ( atomic_weights, sums ) = self.next_weight_data self.weights = atomic_weights self.sums = sums self.centroids = [ s.mean_weighted( w ) for ( s, w ) in zip( self.site_sets, self.weights ) ] self.calculate_superposition_errors() def prepare_new_weights(self, weights): assert len( weights ) == self.site_count atomic_weights = [ weights.deep_copy().set_selected( ~s, 0 ) for s in self.selections ] sums = [ flex.sum( w ) for w in atomic_weights ] if not all( 0 < s for s in sums ): raise RuntimeError, "Zero weights" self.next_weight_data = ( atomic_weights, sums ) def calculate_superposition_errors(self): average = self.average_structure() self.origins = [ average.mean_weighted( w ) for w in self.weights ] self.errors = [ ErrorOnRotation.from_sites( left = sites - centroid, right = average - origin, weights = weights ) for ( sites, centroid, origin, weights ) in zip( self.site_sets, self.centroids, self.origins, self.weights ) ] def full_iteration(self): self.qs = [ eor.best_rotation() for eor in self.errors ] self.calculate_superposition_errors() def rmsd_between(self, left, right): ( rotations, translations ) = self.transformations() lsites = transform_sites( sites = self.site_sets[ left ], rotation = rotations[ left ], translation = translations[ left ] ) rsites = transform_sites( sites = self.site_sets[ right ], rotation = rotations[ right ], translation = translations[ right ] ) return math.sqrt( weighted_error_between( left = lsites, right = rsites, weights = self.weights[ left ] * self.weights[ right ] ) ) def weighted_rmsds(self): return self.rmsds( weights = self.weights ) def unweighted_rmsds(self): return self.rmsds( weights = [ s.as_double() for s in self.selections ] ) def rmsd(self): return math.sqrt( 2.0 / ( self.set_count - 1 ) * self.residual() ) def residual(self): return sum( [ eor.for_quaternion( q = q ) / s for ( q, eor, s ) in zip( self.qs, self.errors, self.sums ) ] ) def transformations(self): ( rotations, translations ) = super( WangSnoeyinkAlgorithm, self ).transformations() return ( rotations, [ t + matrix.col( o ) for ( t, o ) in zip( translations, self.origins ) ], ) def average_structure(self): return self.average_structure_from_transformed_sites( transformeds = self.transformed_sites() ) def distance_squares_from_average_structure(self): transformeds = self.transformed_sites() average = self.average_structure_from_transformed_sites( transformeds = transformeds ) differences = [ s - average for s in transformeds ] return ( reduce( operator.add, [ s.as_double() * d.dot( d ) for ( d, s ) in zip( differences, self.selections ) ] ) / ( self.norms - 1.0 ) ) def rmsds(self, weights): transformed = self.transformed_sites() norm = 1.0 / ( self.set_count - 1 ) return [ math.sqrt( norm * sum( [ weighted_error_between( left = l, right = r, weights = wl * wr ) for ( r, wr ) in zip( transformed, weights ) ] ) ) for ( l, wl ) in zip( transformed, weights ) ] # Internal methods def average_structure_from_transformed_sites(self, transformeds): return ( reduce( operator.add, [ s.set_selected( ~sel, ( 0, 0, 0 ) ) for ( s, sel ) in zip( transformeds, self.selections ) ] ) / self.norms ) # Weighting class UnitWeightScheme(object): """ Unit weights """ def for_difference_squares(self, squares): return flex.double( [ 1.0 ] * len( squares ) ) def __str__(self): return "Unit weighting scheme" class RobustResistantWeightScheme(object): """ Robust-resistant weights """ def __init__(self, critical_value_square): self.critical_value_square = float( critical_value_square ) def for_difference_squares(self, squares): sqr_w = 1.0 - squares / self.critical_value_square sqr_w.set_selected( sqr_w < 0.0, 0.0 ) return sqr_w * sqr_w def __str__(self): return "Robust-resistant weighting scheme (critical value: %s)" % ( self.critical_value_square )