Search code examples
pythonarraysfunctionfor-loopcalling-convention

Calling and evolving a function inside of a for loop


My main function calculateTrajectories is defined below:

def calculateTrajectories(masses, positions, velocities, T, dt):

    #create lists for where we want to know positions, velocities at some time and convert them to numpy arrays
    current_positions = []
    current_velocities = []

    #call updateParticles function to get new positions, velocities at each step
        #loop starts at 0, ends at T, has step value of dt
    for i in range(0, T, dt):
        #show all the time steps in the total time range
        steps = np.array(i)

        #call updateParticles function
        Positions, Velocities = updateParticles(masses, positions, velocities, dt)

        #assign the position and velocity results to their respective lists to get turned into arrays
        current_positions.append(Positions)
        current_velocities.append(Velocities)

        #convert lists into numpy arrays
        new_positions = np.array(current_positions)
        new_velocities = np.array(current_velocities)

        #to make sure the calculation is working
        print(f"computing positions, velocities for time step {i}")

    return steps, new_positions, new_velocities

The updateParticles function uses a leapfrog integrator scheme that evolves the particles over a given time range. It takes the arguments (masses, positions, velocities, dt) where dt is the time step value. The T argument in the calculateTrajectories function is the total time that the calculation will have to go through.

I'm trying to feed this input below into the above function:

#from 1000 days into seconds
T4 = 86400000
#step value
dt4_test = 864000
#in kg
masses4 = [1.989e30, 5.972e24]
#converted to meters
positions4 = [(-448794, 0.0, 0.0),(1.4959742e11, 0.0, 0.0)]
#in m/s
velocities4 = [(0.0, -8.94e02, 0.0),(0.0, 2.98e4, 0.0)]

calculation4 = calculateTrajectories(masses4, positions4, velocities4, T4, dt4_test)
print(calculation4)

I get the correct print statement telling the step value and I also get arrays for both the new_positions and new_velocities. However, the position and velocity arrays do not change, they contain the same values for each step.

How do I edit my calculateTrajectories function so that it does the calculation using the most previous position and velocity in order to get the next one and the next and so on until the required time has finished? Or in other words, how can I call my updateParticles function inside the for loop so that it always uses the new_positions and new_velocities from the last step of the loop?

Thanks for any support with this!


Solution

  • updateParticles is continuously called with the original positions and velocities with this line:

    Positions, Velocities = updateParticles(masses, positions, velocities, dt)
    

    Need to use latest positions and velocities as in:

    Positions, Velocities = updateParticles(masses, current_positions[-1], current_velocities[-1], dt)
    

    New calculateTrajectories function

    def calculateTrajectories(masses, positions, velocities, T, dt):
    
        #create lists for where we want to know positions, velocities at some time and convert them to numpy arrays
        current_positions = [positions]
        current_velocities = [velocities]
    
        #call updateParticles function to get new positions, velocities at each step
            #loop starts at 0, ends at T, has step value of dt
        for i in range(0, T, dt):
            #show all the time steps in the total time range
            steps = np.array(i) #--this could be moved outside the for loop
    
            # call updateParticles function--pass in the latest position and velocities (i.e. last index of each array)
            Positions, Velocities = updateParticles(masses, current_positions[-1], current_velocities[-1], dt)
    
            # Update position and velocities
            current_positions.append(Positions)
            current_velocities.append(Velocities)
    
            #to make sure the calculation is working
            print(f"computing positions, velocities for time step {i}")
    
        #assign the position and velocity results to their respective lists to get turned into arrays
        # only need to convert to bumpy arrays once so moved outside for loop
        new_positions = np.array(current_positions)
        new_velocities = np.array(current_velocities)
    
        return steps, new_positions, new_velocities