From 1576efd9c3b29b782a0d19453a4691ff20d4f055 Mon Sep 17 00:00:00 2001 From: Thomas Wemyss Date: Thu, 10 Dec 2020 12:42:21 +0000 Subject: [PATCH 1/3] Update variable names and add comments --- week09/improvement/trees.py | 44 ++++++++++++++++++++++++++----------- 1 file changed, 31 insertions(+), 13 deletions(-) diff --git a/week09/improvement/trees.py b/week09/improvement/trees.py index dfa03f3..be421e1 100644 --- a/week09/improvement/trees.py +++ b/week09/improvement/trees.py @@ -1,17 +1,35 @@ """This code produces a tree-like plot.""" -from math import sin, cos +from math import sin, cos, pi from matplotlib import pyplot as plt -s=1 -d=[[0,1,0]] -plt.plot([0,0],[0,1]) -for i in range(5): - n=[] - for j in range(len(d)): - n.append([d[j][0]+s*sin(d[j][2]-0.1), d[j][1]+s*cos(d[j][2]-0.1), d[j][2]-0.1]) - n.append([d[j][0]+s*sin(d[j][2]+0.1), d[j][1]+s*cos(d[j][2]+0.1), d[j][2]+0.1]) - plt.plot([d[j][0], n[-2][0]],[d[j][1], n[-2][1]]) - plt.plot([d[j][0], n[-1][0]],[d[j][1], n[-1][1]]) - d=n - s*=0.6 + +# This is the length of the first branches of the tree, and is scaled +# down for the following branches +base_length = 1 + +# Length of the initial vertical branch +trunk_length = 1 + +# How many times the tree should split +num_branches = 5 + +# First element is the x coordinate of the second element in the tree +# Second element is the y-height of the second element in the tree +# Angle of the second element in the tree (clockwise, relative to vertical) +tree_branch = [[0, trunk_length, 0]] + +# Plot the initial vertical "trunk" +plt.plot([0,0],[0,trunk_length]) + +for _ in range(num_branches): + tree = [] + for j in range(len(tree_branch)): + tree.append([tree_branch[j][0]+base_length*sin(tree_branch[j][2]-0.1), tree_branch[j][1]+base_length*cos(tree_branch[j][2]-0.1), tree_branch[j][2]-0.1]) + tree.append([tree_branch[j][0]+base_length*sin(tree_branch[j][2]+0.1), tree_branch[j][1]+base_length*cos(tree_branch[j][2]+0.1), tree_branch[j][2]+0.1]) + plt.plot([tree_branch[j][0], tree[-2][0]],[tree_branch[j][1], tree[-2][1]]) + plt.plot([tree_branch[j][0], tree[-1][0]],[tree_branch[j][1], tree[-1][1]]) + tree_branch = tree + base_length *= 0.6 + + plt.savefig('tree.png') \ No newline at end of file From c61b69909e9c1bbd43ec41417a8b0d1b9f29c85a Mon Sep 17 00:00:00 2001 From: Thomas Wemyss Date: Thu, 10 Dec 2020 12:49:38 +0000 Subject: [PATCH 2/3] Move branches into a class --- week09/improvement/trees.py | 61 ++++++++++++++++++++++++------------- 1 file changed, 40 insertions(+), 21 deletions(-) diff --git a/week09/improvement/trees.py b/week09/improvement/trees.py index be421e1..e2adad7 100644 --- a/week09/improvement/trees.py +++ b/week09/improvement/trees.py @@ -1,35 +1,54 @@ """This code produces a tree-like plot.""" -from math import sin, cos, pi +from math import sin, cos from matplotlib import pyplot as plt -# This is the length of the first branches of the tree, and is scaled -# down for the following branches -base_length = 1 +# Make a class to hold each branch +class Branch(): + def __init__(self, x_coord, y_coord, base_angle): + self.x = x_coord + self.y = y_coord + self.base_angle = base_angle -# Length of the initial vertical branch +# Set the length of the first branch (from 0,0 to 0,trunk_length) trunk_length = 1 -# How many times the tree should split -num_branches = 5 +# Set up a list holding the end of the first branch/trunk +# First element is the x coordinate of the centre of the tree +# The second element is the length of the intial trunk (the height where the branches start) +# Third element is angle of the branch (radians, relative to vertical) +tree=[Branch(0, trunk_length, 0)] + +# This scales the successive branches of the tree down in length +scale_factor = 0.6 -# First element is the x coordinate of the second element in the tree -# Second element is the y-height of the second element in the tree -# Angle of the second element in the tree (clockwise, relative to vertical) -tree_branch = [[0, trunk_length, 0]] +# This is the number of times that the tree branches after the initial trunk +num_branches = 5 -# Plot the initial vertical "trunk" +# Add a plot for the "trunk" branch plt.plot([0,0],[0,trunk_length]) +# Make the successive branches for _ in range(num_branches): - tree = [] - for j in range(len(tree_branch)): - tree.append([tree_branch[j][0]+base_length*sin(tree_branch[j][2]-0.1), tree_branch[j][1]+base_length*cos(tree_branch[j][2]-0.1), tree_branch[j][2]-0.1]) - tree.append([tree_branch[j][0]+base_length*sin(tree_branch[j][2]+0.1), tree_branch[j][1]+base_length*cos(tree_branch[j][2]+0.1), tree_branch[j][2]+0.1]) - plt.plot([tree_branch[j][0], tree[-2][0]],[tree_branch[j][1], tree[-2][1]]) - plt.plot([tree_branch[j][0], tree[-1][0]],[tree_branch[j][1], tree[-1][1]]) - tree_branch = tree - base_length *= 0.6 - + # Store all the branches from the current "end" here + branches = [] + # Loop over each branch in the tree and add the child nodes + for j in range(len(tree)): + # For two sub branches, we go at -0.1 radians and +0.1 radians from the parent/lower branch + for subbranch_angle in [-0.1, 0.1]: + # Create the new branch and store + new_node = Branch( + tree[j].x + trunk_length * sin(tree[j].base_angle + subbranch_angle), + tree[j].y + trunk_length * cos(tree[j].base_angle + subbranch_angle), + tree[j].base_angle + subbranch_angle + ) + branches.append(new_node) + # Plot the branches we just created + plt.plot([tree[j].x, branches[-2].x], [tree[j].y, branches[-2].y]) + plt.plot([tree[j].x, branches[-1].x], [tree[j].y, branches[-1].y]) + # Store the branches for the next iteration + tree=branches + # Reduce the trunk length for the next branches + trunk_length *= scale_factor plt.savefig('tree.png') \ No newline at end of file From 629c132576c954e0621266c733e34420b4463d9a Mon Sep 17 00:00:00 2001 From: Thomas Wemyss Date: Thu, 10 Dec 2020 12:56:03 +0000 Subject: [PATCH 3/3] Use argparse to load arguments --- week09/improvement/trees.py | 94 ++++++++++++++++++++----------------- 1 file changed, 52 insertions(+), 42 deletions(-) diff --git a/week09/improvement/trees.py b/week09/improvement/trees.py index e2adad7..6fab445 100644 --- a/week09/improvement/trees.py +++ b/week09/improvement/trees.py @@ -1,8 +1,10 @@ """This code produces a tree-like plot.""" +import argparse from math import sin, cos from matplotlib import pyplot as plt + # Make a class to hold each branch class Branch(): def __init__(self, x_coord, y_coord, base_angle): @@ -10,45 +12,53 @@ def __init__(self, x_coord, y_coord, base_angle): self.y = y_coord self.base_angle = base_angle -# Set the length of the first branch (from 0,0 to 0,trunk_length) -trunk_length = 1 - -# Set up a list holding the end of the first branch/trunk -# First element is the x coordinate of the centre of the tree -# The second element is the length of the intial trunk (the height where the branches start) -# Third element is angle of the branch (radians, relative to vertical) -tree=[Branch(0, trunk_length, 0)] - -# This scales the successive branches of the tree down in length -scale_factor = 0.6 - -# This is the number of times that the tree branches after the initial trunk -num_branches = 5 - -# Add a plot for the "trunk" branch -plt.plot([0,0],[0,trunk_length]) - -# Make the successive branches -for _ in range(num_branches): - # Store all the branches from the current "end" here - branches = [] - # Loop over each branch in the tree and add the child nodes - for j in range(len(tree)): - # For two sub branches, we go at -0.1 radians and +0.1 radians from the parent/lower branch - for subbranch_angle in [-0.1, 0.1]: - # Create the new branch and store - new_node = Branch( - tree[j].x + trunk_length * sin(tree[j].base_angle + subbranch_angle), - tree[j].y + trunk_length * cos(tree[j].base_angle + subbranch_angle), - tree[j].base_angle + subbranch_angle - ) - branches.append(new_node) - # Plot the branches we just created - plt.plot([tree[j].x, branches[-2].x], [tree[j].y, branches[-2].y]) - plt.plot([tree[j].x, branches[-1].x], [tree[j].y, branches[-1].y]) - # Store the branches for the next iteration - tree=branches - # Reduce the trunk length for the next branches - trunk_length *= scale_factor - -plt.savefig('tree.png') \ No newline at end of file +def make_tree(filename, num_levels = 5): + # Set the length of the first branch (from 0,0 to 0,trunk_length) + trunk_length = 1 + + # Set up a list holding the end of the first branch/trunk + # First element is the x coordinate of the centre of the tree + # The second element is the length of the intial trunk (the height where the branches start) + # Third element is angle of the branch (radians, relative to vertical) + tree=[Branch(0, trunk_length, 0)] + + # This scales the successive branches of the tree down in length + scale_factor = 0.6 + + # This is the number of times that the tree branches after the initial trunk + num_branches = num_levels + + # Add a plot for the "trunk" branch + plt.plot([0,0],[0,trunk_length]) + + # Make the successive branches + for _ in range(num_branches): + # Store all the branches from the current "end" here + branches = [] + # Loop over each branch in the tree and add the child nodes + for j in range(len(tree)): + # For two sub branches, we go at -0.1 radians and +0.1 radians from the parent/lower branch + for subbranch_angle in [-0.1, 0.1]: + # Create the new branch and store + new_branch = Branch( + tree[j].x + trunk_length * sin(tree[j].base_angle + subbranch_angle), + tree[j].y + trunk_length * cos(tree[j].base_angle + subbranch_angle), + tree[j].base_angle + subbranch_angle + ) + branches.append(new_branch) + # Plot the branch that we just created + plt.plot([tree[j].x, new_branch.x], [tree[j].y, new_branch.y]) + # Store the branches for the next iteration + tree=branches + # Reduce the trunk length for the next branches + trunk_length *= scale_factor + # Save the figure + plt.savefig(filename) + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("filename") + parser.add_argument("--num_branches") + args = parser.parse_args() + make_tree(args.filename) + \ No newline at end of file