Is there any way to create a number of equal class objects in a for loop in Python?

user_15592052 0
Viewed 62

I started learning Python about a month ago because I need it for a project at University, so I apologize in advance for any conceptual mistakes or erroneous assumptions I may make in my explanation.

Before explaining my problem, I checked a couple of questions like the following, but I don't think the answers are what I need:

Python create objects in for loop

The program simulates the trajectories of a number of fluid particles. The way I do this is by creating a 'line' and a 'dot' object that will store the values of the x- and y-coordinates of the particles and then animating them with matplotlib:

import numpy as np
from matplotlib import pyplot as plt
from matplotlib.animation import FuncAnimation
import math

# initializing the figure in 
# which the graph will be plotted
fig = plt.figure() 
   
# setting limits for both the x-axis and y-axis
axis = plt.axes(xlim =(0, 20), 
                ylim =(-2, 50))

# initializing a line variable for each
# trajectory to be plotted
line1, = axis.plot([], [], lw = 1)
line2, = axis.plot([], [], lw = 1)

redDot, = plt.plot([1], [10], 'ro', markersize = 2)

# data which the line will 
# contain (x, y). Lines are
# initially empty objects
def init(): 
    line1.set_data([], [])
    line2.set_data([], [])
    return line1, line2
   
# initializing empty values
# for x and y coordinates
xdata1, ydata1 = [], [] 
xdata2, ydata2 = [], []

# animation function 
def animate(i): 
    # t is a parameter which varies
    # with the frame number
    t = 0.1 * i 
       
    # x, y values to be plotted 
    x = math.exp(t) 
    y = 1 + (2/3)*t**(3/2) 
    
    # appending values to the previously 
    # empty x and y data holders 
    xdata1.append(x) 
    ydata1.append(y) 
    line1.set_data(xdata1, ydata1) 
    
    x = math.exp(t) 
    y = 10 + (2/3)*t**(3/2) 
    
    # appending values to the previously 
    # empty x and y data holders 
    xdata2.append(x) 
    ydata2.append(y) 
    line2.set_data(xdata2, ydata2)
    
    redDot.set_data(x, y)
    
    return line1, line2, redDot

anim = FuncAnimation(fig, animate, init_func = init,
                     frames = 200, interval = 20, blit = True)

The code above creates a line where it says line1, = axis.plot([], [], lw = 1). Immediately after that, it creates a point with redDot, = plt.plot([1], [10], 'ro', markersize = 2), which will follow the line. The problem is I intend to simulate 10 particles at the same time, and I believe there must be a cleaner, more efficient way to create all these objects than having to repeat those two lines of code 10 times. What's more, there are two lists that act as data holders for each particle, and the animate function needs to go through all of them too to assign the calculated values. It would make for a very lengthy code, and it would be downright impractical if I wanted to animate a larger number of particles.

Is there any way to do that with a loop? Something like:

for i in list(range(number_of_particles)):

    #number_of_particles lines and dots are created here

I'm wondering if I could store these objets in a list, such as:

my_list = []

for i in list(range(number_of_particles)):

    my_list.append(axis.plot([], [], lw = 1))

Then I would only have to call the items of said list in another for loop.

Again, I am just a newbie and I really appreciate any input from anyone who might want to help. Thanks!

matplotlib object for-loop python
1 Answers

I think your own suggestions work more or less, which I have just adapted here below.

  • To create a list of Matplotlib.lines.Line2D instances
number_of_particles = 10
lines = [axis.plot([], [], lw=1)[0] for i in range(number_of_particles)]
  • Storing the data for each line: Create a list to store the x and y data for each line, you can do the following, which simply creates a list of tuples, where each tuple contains a list for the x-coordinates and a list for the y-coordinates. This is similar to the spirit of your xdata1, ydata1 = [], [], but just a list of such tuples.
lines_data = [([], []) for i in range(number_of_particles)]
  • Updating each of the lines in the animate function. Here the built-in enumerate function gets you both the index of the ith iteration so you can use it in the loop, and the item in the iterable list. I suppose to have more flexibility in the functions you could perhaps add the constants to lists and then index them with i if needed.
for i, line in enumerate(lines):
    x = math.exp(t)
    y = i + (2/3)*t**(3/2)
    lines_data[i][0].append(x)
    lines_data[i][1].append(y)
    line.set_data(lines_data[i][0], lines_data[i][1])
  • Smart way to return a list. When returning use a * to have the list unpacked. Have just added the lines list to the existing return statement.
return line1, line2, redDot, *lines

Finally, I was not fully sure about the comment

What's more, there are two lists that act as data holders for each particle, and the animate function needs to go through all of them too to assign the calculated values.

Maybe you could elaborate on that, but it seems that the new computations are just appended to the end, which I suppose shouldn't have the worst impact.

user_82714440
Related