Plotting multiple functions in Sage from a for loop

edit

Hello, @EconJohn!

What you can do is to create a Graphics object (let's call it p) before entering the loop, then superimpose the plots of each k1 over the previous one and store the result in p, and then show the plot outside of the loop. Since yourk1s are defined as symbolic equalities, you will need to use the rhs() (right-hand side of the equality) method inside the plot() command.

Before writing the final version of the code, I would like to make a few observations: You don't need to define v0 and beta as symbolic variables, since you later assign to beta the value 0.6, which overwrites the effect of the var() command, converting it into a numerical variable; similarly, you assign v0 = log(k), which overwrites the definition made by var() (in this case, v0 IS automatically made a symbolic variable, since it is the logarithm of the symbolic variable k.) For a similar reason, it is unnecessary to define vk with var(). Finally, you define but never use the vk1 variable.

Here is the code:

#Step1: Name variables
k,k1 = var('k,k1')

#Step2: Initialize your values
beta = 0.6
v0 = log(k)
cols = ['red', 'darkgreen', 'blue', 'magenta', 'yellowgreen']
p = Graphics()

#Step 3: The Loop to Obtain Policy function
for n in range(5):
    vk = log(k-k1) + beta*v0(k=k1)
    FOC = vk.diff(k1)
    k1star = solve(FOC==0, k1)
    print(n, k1star)
    v0 = (vk).subs(k1=k1star[0].rhs())
    p += plot(k1star[0].rhs(), color=cols[n], legend_label=str(k1star[0]))
p.show()

I have taken the liberty of defining a list cols with some different colors in order to differentiate the plots. The instruction p = Graphics() defines p as an empty plot. Inside the loop, I have added the instruction

p += plot(k1star[0].rhs(), color=cols[n], legend_label=str(k1star[0]))

Remember that superimposing plots is equivalent to "adding" the plots, so this command generates a plot for every iteration and adds it to p. I have used the $n$-th color in the cols list. I have also added a label, indicating which color and formula correspond to which plot. The last line, p.show() finally shows the plot.

I hope this helps! Please, let me know if you need some additional help.