from numpy import * from myarray import * from springs import compute_edge_lengths def forward_euler( p_n, v_n, edges, edge_rest_lengths, constraints, dt, verbose = True ): ''' Given a list of points 'p' as an n-by-2 array, a list of velocities 'v' 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 constraints denoting p[i] doesn't change, takes one forward Euler step of size 'dt' and returns the p, v arrays corresponding to the new time step. NOTE: 'edges' must not have both (i,j) and (j,i) ''' from springs import F ## f = m*a ## We assume m = 1, so multiplying by 1/m has no effect. a_n = F( p_n, edges, edge_rest_lengths ).reshape( p_n.shape ) if verbose: print 'a_n:' print a_n ## Apply the explicit update rules. ## NOTE: With no damping, this oscillates a lot. One way to apply damping is simply to ## multiply velocity by a factor < 1. damping = 1. #damping = .99 v_n_p_1 = damping * v_n + dt * a_n p_n_p_1 = p_n + dt * v_n ## Apply the constraints. for i in constraints: p_n_p_1[i] = p_n[i] v_n_p_1[i] = 0. return p_n_p_1, v_n_p_1 def backward_euler( p_n, v_n, edges, edge_rest_lengths, constraints, dt, verbose = True ): ''' Given a list of points 'p' as an n-by-2 array, a list of velocities 'v' 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 constraints denoting p[i] doesn't change, takes one backward Euler step of size 'dt' and returns the p, v arrays corresponding to the new time step. NOTE: 'edges' must not have both (i,j) and (j,i) ''' from springs import F, J, constrain_system XSTEP_THRESHOLD = 1e-5 F_THRESHOLD = 1e-8 MAX_ITERATIONS = 100 N = prod( p_n.shape ) dim = p_n.shape[1] ## Combine positions (p) and velocities (v) into a single, large vector. pv_n = emptyf( 2 * N ) pv_n[ : N ] = p_n.flatten() pv_n[ N : ] = v_n.flatten() ## Apply constraints. constrain_rows = [] for i in constraints: ## constrain positions constrain_rows.extend( range( i*dim, (i+1) * dim ) ) ## constraint velocities constrain_rows.extend( range( N + i*dim, N + (i+1) * dim ) ) ## zero initial velocities pv_n[ N + i*dim : N + (i+1) * dim ] = 0. ## Store the current positions and velocities. pv_n0 = pv_n.copy() ## g() is a function of pv (positions and velocities) that is zero when ## dt * (the derivative at pv) = pv - pv_n0, ## where pv_n0 is our initial pv. def g( pv ): va = emptyf( 2*N ) va[ :N ] = pv[ N: ] va[ N: ] = F( pv[ :N ].reshape( p_n.shape ), edges, edge_rest_lengths ) return dt*va - pv + pv_n0 ## dg_dpv() computes the Jacobian of g with respect to pv. def dg_dpv( pv ): dg = zerosf( (2*N,2*N) ) dg[ :N, N: ] = identityf( N ) dg[ N:, :N ] = J( pv[ :N ].reshape( p_n.shape ), edges, edge_rest_lengths ) return dt * dg - identityf( 2*N ) ## Newton-Raphson iterations to find the solution, a satisfying pv. iteration = 0 while True: if verbose: print '-- iteration', iteration, '--' iteration += 1 g_n = g( pv_n ) dg_n = dg_dpv( pv_n ) mag2_g_n = sum( g_n ** 2 ) if verbose: print '| g( pv_n ) |^2:', mag2_g_n if mag2_g_n < F_THRESHOLD: break constrain_system( dg_n, g_n, constrain_rows ) # pv_n_p_1 = pv_n - dot( linalg.inv( dg_n ), g_n ) ## <=> pv_n_p_1 - pv_n = -linalg.inv( dg_n ) * g_n ## <=> pv_n - pv_n_p_1 = linalg.inv( dg_n ) * g_n ## <=> dg_n * ( pv_n - pv_n_p_1 ) = g_n p_negative_delta = linalg.solve( dg_n, g_n ) ## pv_n - ( pv_n - pv_n_p_1 ) = pv_n_p_1 pv_n_p_1 = pv_n - p_negative_delta diff2 = sum( ( pv_n_p_1 - pv_n ) ** 2 ) if verbose: print '| pv_n+1 - pv_n |^2:', diff2 pv_n = pv_n_p_1 if diff2 < XSTEP_THRESHOLD: break if iteration >= MAX_ITERATIONS: print 'Diverged.' return p_n.copy(), v_n.copy() break return pv_n[ :N ].reshape( p_n.shape ), pv_n[ N: ].reshape( v_n.shape ) 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, 2 ] print 'constraints:', constraints p_initial = p_undeformed.copy() p_initial[1] += array( (.5,0) ) print 'p initial:', p_initial v_initial = 0. * p_initial print 'v initial:', v_initial import sys explicit = 'explicit' in sys.argv[1:] if explicit: dt = 1 N = 1 stepper, stepper_name = forward_euler, 'forward euler' else: dt = 1 N = 1 stepper, stepper_name = backward_euler, 'backward euler' t = 0. p_step = p_initial.copy() v_step = v_initial.copy() for n in range(N): p_step, v_step = stepper( p_step, v_step, edges, edge_rest_lengths, constraints, dt ) t += dt print '%s step, t = %s' % (stepper_name, t) print 'p:' print p_step.round(4) print 'v:' print v_step.round(4) def main(): test1() if __name__ == '__main__': main()