from numpy import * from myarray import * ## This is the spring constant k = 1. ## This is our threshold for whether a number is approximately zero. ZERO_THRESHOLD = 1e-8 def F_ij( p_i, p_j, r_ij ): ''' Returns the force of the spring from 'p_i' to 'p_j' with rest length 'r_ij' acting on 'p_i'. Note that F_ij( p_i, p_j, r_ij ) equals -F_ij( p_j, p_i, r_ij ). ''' p_ij = p_i - p_j len_p_ij = sqrt( sum( p_ij ** 2 ) ) if abs( len_p_ij ) < ZERO_THRESHOLD: result = 0. * p_ij else: result = -k * ( len_p_ij - r_ij ) / len_p_ij * p_ij return result def F( p, edges, edges_rest_lengths ): ''' Returns a vector containing the force at every point. Note that the input 'p' is assumed to be a number-of-points by dimension array, where dimension is 2 for our example. The result is flattened into a single vector of size number-of-points times dimension. ''' dim = p.shape[1] Fp = zerosf( prod( p.shape ) ) ## Loop over every edge and its corresponding rest length. ## (zip() simply combines two lists into one so we can loop over them together.) for (i,j), r_ij in zip( edges, edges_rest_lengths ): assert i != j Fij = F_ij( p[i], p[j], r_ij ) Fp[ i*dim : (i+1) * dim ] += Fij ## 'edges' contains edges uniquely, so 'edges' will contain (i,j) but not (j,i). ## This means that we must add -Fij to row j as well as Fij to row i. Fp[ j*dim : (j+1) * dim ] += -Fij return Fp def dF_ij_d_p_i( p_i, p_j, r_ij ): ''' Returns the derivative with respect to 'p_i' of the force of the spring from 'p_i' to 'p_j' with rest length 'r_ij' acting on 'p_i'. Our dimension is 2, so this is a 2x2 quantity. Note that dF_ij_d_p_i( p_i, p_j, r_ij ) equals dF_ij_d_p_i( p_j, p_i, r_ij ). ''' dim = p_i.shape[0] p_ij = p_i - p_j len_p_ij = sqrt( sum( p_ij ** 2 ) ) if abs( len_p_ij ) < ZERO_THRESHOLD: result = -k * identity( dim ) else: result = -k * identity( dim ) - k * r_ij / len_p_ij**3 * outer( p_ij, p_ij ) + k * r_ij / len_p_ij * identity( dim ) return result def J( p, edges, edges_rest_lengths ): ''' Returns a matrix containing the derivative of the force at every point with respect to each point. Note that the input 'p' is assumed to be a number-of-points by dimension array, where dimension is 2 for our example. The result is flattened is a square matrix (of type numpy.array), of size NxN, where N = number-of-points times dimension. ''' dim = p.shape[1] Jp = zerosf( ( prod( p.shape ), prod( p.shape ) ) ) ## Loop over every edge and its corresponding rest length. ## (zip() simply combines two lists into one so we can loop over them together.) for (i,j), r_ij in zip( edges, edges_rest_lengths ): assert i != j dF = dF_ij_d_p_i( p[i], p[j], r_ij ) assert ( ( Jp[ i*dim : (i+1) * dim, j*dim : (j+1) * dim ] - zeros( ( dim, dim ) ) ) ** 2 ).sum().sum() == 0. Jp[ i*dim : (i+1) * dim, j*dim : (j+1) * dim ] = -dF Jp[ i*dim : (i+1) * dim, i*dim : (i+1) * dim ] += dF ## 'edges' contains edges uniquely, so 'edges' will contain (i,j) but not (j,i). ## This means that we must add dF to the right places in column j as well. Jp[ j*dim : (j+1) * dim, i*dim : (i+1) * dim ] = -dF Jp[ j*dim : (j+1) * dim, j*dim : (j+1) * dim ] += dF return Jp def constrain_system( A, rhs, rows ): ''' This function modifies its input parameters, a system matrix 'A' and right-hand-side vector 'rhs', such that for every index i in 'rows', the row i of A is set to row i of the identity matrix and rhs[i] is set to zero. ''' for i in rows: A[ i, : ] = zeros( A.shape[1] ) ## We can also zero the column, which keeps the matrix symmetric, because ## we are zeroing the corresponding entries in the right-hand-side (x*0 = 0). A[ :, i ] = zeros( A.shape[0] ) A[ i, i ] = 1 rhs[i] = 0 return A, rhs def static_solution( p, edges, edges_rest_lengths, constraints, verbose = True ): ''' Given a list of points 'p' as an n-by-2 array, a list of (i,j) pairs 'edges' denoting an edge between points p[i] and p[j], a list of rest lengths (one for each edge in 'edges'), and a list of position constraints (i, position) denoting p[i] = position, uses Newton's method to solve for the positions where the forces are all zero. NOTE: 'edges' must not have both (i,j) and (j,i) ''' XSTEP_THRESHOLD = 1e-5 F_THRESHOLD = 1e-8 MAX_ITERATIONS = 100 p_n = p.copy().flatten() dim = p.shape[1] constrain_rows = [] for i, p_val in constraints: p_n[ i*dim : (i+1) * dim ] = p_val constrain_rows.extend( range( i*dim, (i+1) * dim ) ) iteration = 0 while True: if verbose: print '-- iteration', iteration, '--' iteration += 1 Jp_n = J( p_n.reshape( p.shape ), edges, edges_rest_lengths ) Fp_n = F( p_n.reshape( p.shape ), edges, edges_rest_lengths ) mag2_Fp_n = sum( Fp_n ** 2 ) if verbose: print '| F( p_n ) |^2:', mag2_Fp_n if mag2_Fp_n < F_THRESHOLD: break constrain_system( Jp_n, Fp_n, constrain_rows ) # p_n_p_1 = p_n - dot( linalg.inv( Jp_n ), Fp_n ) ## <=> p_n_p_1 - p_n = -linalg.inv( Jp_n ) * Fp_n ## <=> p_n - p_n_p_1 = linalg.inv( Jp_n ) * Fp_n ## <=> Jp_n * ( p_n - p_n_p_1 ) = Fp_n p_negative_delta = linalg.solve( Jp_n, Fp_n ) ## p_n - ( p_n - p_n_p_1 ) = p_n_p_1 p_n_p_1 = p_n - p_negative_delta diff2 = sum( ( p_n_p_1 - p_n ) ** 2 ) if verbose: print '| p_n+1 - p_n |^2:', diff2 p_n = p_n_p_1 if diff2 < XSTEP_THRESHOLD: break if iteration >= MAX_ITERATIONS: print 'Diverged.' return p.copy() break return p_n.reshape( p.shape ) def compute_edge_lengths( p, edges ): ''' Given a list of (i,j) pairs 'edges' denoting an edge between points p[i] and p[j], returns a list of rest lengths, one for each edge in 'edges'. NOTE: 'edges' must not have both (i,j) and (j,i) ''' ## Check for duplicate edges, which are forbidden. edges = tuple( map( tuple, edges ) ) from sets import ImmutableSet as Set assert len( Set( map( Set, edges ) ) ) == len( edges ) result = [] for i,j in edges: len_p_ij = sqrt( sum( (p[i] - p[j]) ** 2 ) ) result.append( len_p_ij ) return result def test1(): print '===== test1() =====' #p_undeformed = arrayf( [[0,0], [1,1]] ) #p_undeformed = arrayf( [[0,0], [1,0], [1,1]] ) #p_undeformed = arrayf( [[0,0], [1,0], [0,1], [1,1]] ) p_undeformed = arrayf( [[0,0], [1,0], [2,0]] ) print 'p.shape:', p_undeformed.shape print 'p undeformed:', p_undeformed #edges = [ (0,1) ] edges = [ (0,1), (1,2) ] #edges = [ (0,1), (0,2), (1,3), (2,3) ] print 'edges:', edges ## Multiply p_undeformed by 0 to force 0 rest length springs #edge_rest_lengths = compute_edge_lengths( 0 * p_undeformed, edges ) edge_rest_lengths = compute_edge_lengths( p_undeformed, edges ) print 'edge_rest_lengths:', edge_rest_lengths #constraints = [ ( 0, p_undeformed[0] ) ] #constraints = [ ( 0, p_undeformed[0] ), ( 3, p_undeformed[3] ) ] constraints = [ ( 0, p_undeformed[0] ), ( 2, p_undeformed[2] ) ] print 'constraints:', constraints p_initial = p_undeformed.copy() p_initial[1] += array( (.5,0) ) print 'p initial:', p_initial p_solution = static_solution( p_initial, edges, edge_rest_lengths, constraints ) print 'static solution:', p_solution def main(): test1() if __name__ == '__main__': main()