Python fractal tree using SVG

Question

I made a simple program to generate fractal tree using PATH object of SVG. Any suggestions?

import math
from random import randint
# const - upper limit for randint
s = 10
"""Simple fractal tree using SVG and recursion.
Usage:
 Create Root object bt Branch(x1=400, y1=800, x2=400, y2=600, color=60, size=35)
 x1, y1, x2, y2 - start points of root
 Generate Tree Tree(lenght=200, angle=-20, depth=9, x1=400, y1=600, size=35, color=60, outlist=resutlist)
 lenght - lenght of start branch
 angle - start angle of branch
 depth - number of tree level
 x1, y1 - start point of branch
"""
class Branch():
 """Class represents a single branch."""
 def __init__(self, x1, y1, x2, y2, color, size):
 """Assigning values."""
 self.x1 = x1
 self.x2 = x2
 self.y1 = y1
 self.y2 = y2
 self.color = color
 self.size = size
 def __str__(self):
 """Return path SVG object with points, color and stroke of branch."""
 return '<path d="M {x1} {y1} L {x2} {y2}" stroke="rgb(100,{c},0)" stroke-width="{w}"/>\n'.format(
 x1=self.x1,
 y1=self.y1,
 x2=self.x2,
 y2=self.y2,
 w=self.size,
 c=self.color
 )
 def __repr__(self):
 """Return text represent object."""
 return self.__str__()
class Tree():
 """
 Class represents Tree.
 Tree is composed of Branch object.
 """
 def __init__(self, lenght, angle, depth, x1, y1, size, color, outlist):
 """Main point of start generation."""
 self.branches = self.drawbranch(lenght, angle, depth, x1, y1, size, color, outlist)
 def drawbranch(self, lenght, angle, depth, x1, y1, size, color, outlist):
 """Recursive function for generate three Branch object per iteration."""
 # if depth > 0
 if depth:
 # X value of second point
 x2 = x1 + lenght * math.cos(math.radians(angle))
 # Y value of second point
 y2 = y1 + lenght * math.sin(math.radians(angle))
 # modify lenght of single branch
 lenght = float(2.0 / 3.0 * lenght)
 # modify size of single branch
 size = float(2.0 / 3.0 * size) + 1
 # modify color of single branch
 color += 6
 # X value of B point
 bx = x1
 # Y value of B point
 by = y2
 # X value of C point
 cx = -x2 + 2 * x1
 # Y value of C point
 cy = y2
 # Create A point
 b1 = Branch(x1, y1, x2, y2, color, size)
 # Add to list
 outlist.append(str(b1))
 # Call drawbranch function (recursion)
 self.drawbranch(lenght, angle + randint(-10, s), depth - 1, x2, y2, size, color, outlist)
 # Create B point
 b2 = Branch(x1, y1, bx, by, color, size)
 # Add to list
 outlist.append(str(b2))
 # Calculate new angle
 nangle = angle + randint(-1, 0) * randint(1, s)
 # Call drawbranch function (recursion)
 self.drawbranch(lenght, nangle, depth - 1, bx, by, size, color, outlist)
 # Create C point
 b3 = Branch(x1, y1, cx, cy, color, size)
 # Add to list
 outlist.append(str(b3))
 # Calculate new angle
 nangle = angle + randint(0, 1) * randint(1, s)
 # Call drawbranch function (recursion)
 self.drawbranch(lenght, nangle, depth - 1, cx, cy, size, color, outlist)
 # Return list of branches
 return outlist
 def write_svg(self, output='drzewko.svg'):
 """Function that write all branches to SVG file."""
 with open(output, 'w') as outfile:
 # Write SVG declaration
 outfile.write('<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 800 800" version="1.1">\n')
 # Map to str all branches and write it into file
 outfile.writelines(map(str, self.branches))
 # End of SVG file
 outfile.write('</svg>\n')
print "Start generating, please wait.."
# Create empty list
resutlist = []
# Create root of Tree and add to list
resutlist.append(Branch(x1=400, y1=800, x2=400, y2=600, color=60, size=35))
# Call Tree object
t = Tree(lenght=200, angle=-20, depth=9, x1=400, y1=600, size=35, color=60, outlist=resutlist)
# After generate tree save to file
t.write_svg()
print "Done, check SVG file"

Results:

enter image description here

• πάντα ῥεῖ: The generated tree doesn't look very natural. Needs some more randomness probably.

Answer 1

Interesting problem thanks for sharing.

I have refactored your code and it is included in it entirety below. The most significant change made was to introduce a Point class, to encapsulate (x, y) pairs, and to provide a bit of syntactic sugar to math on same.

Make Point as class:

So this code:

# X value of C point
cx = -x2 + 2 * x1
# Y value of C point
cy = y2

can become:

c = Point(-p2.x + 2 * p1.x, p2.y)

via a namedtuple that provides an immutable datastructure whose attributes can be accessed with things like p2.x.

Comments:

I removed many of your comments. In several of these cases, I renamed a variable to include the information that was present in the comment.

Use object attributes

Before the edits, the variable branches was being passed recursively, but it was eventually assigned as an attribute of the Tree class. So the new code just starts there, and branches is an attribute of the class from the beginning.

Type cast to float?

In general you do not need to cast floats to floats. So something like:

 branch_length = float(2.0 / 3.0 * length)

can just be:

 branch_length = 2.0 / 3.0 * length

since 2.0 is already a float.

Complete Code:

import math
from random import randint
from collections import namedtuple
# const - upper limit for randint
s = 10
class Point(namedtuple('Point', 'x y')):
 def __str__(self):
 return'{} {}'.format(self.x, self.y)
 def __add__(self, other):
 assert isinstance(other, Point)
 return Point(self.x + other.x, self.y + other.y)
 def __mul__(self, other):
 return Point(self.x * other, self.y * other)
 def __rmul__(self, other):
 return self.__mul__(other)
class Branch(namedtuple('Branch', 'p1 p2 color size')):
 def __str__(self):
 """Path SVG object with points, color and stroke of branch."""
 return ('<path d="M {p1} L {p2}" '
 'stroke="rgb(100,{c},0)" stroke-width="{w}"/>\n'.
 format(p1=self.p1, p2=self.p2, w=self.size, c=self.color))
 def __repr__(self):
 return self.__str__()
class Tree(object):
 def __init__(self, length, angle, depth, point, size, color, outlist):
 """Main point of start generation."""
 self.branches = outlist
 self.draw_branches(length, angle, depth, point, size, color)
 def draw_branches(self, length, angle, depth, p1, size, color):
 """ Recursively generate three Branch objects per iteration."""
 if depth <= 0:
 return
 p2 = p1 + length * Point(
 math.cos(math.radians(angle)),
 math.sin(math.radians(angle))
 )
 # set some new characteristics for the next level
 branch_length = 2.0 / 3.0 * length
 branch_size = 2.0 / 3.0 * size + 1
 color += 6
 # Calculate new angle and recurse
 self.branches.append(Branch(p1, p2, color, branch_size))
 nangle = angle + randint(-10, s)
 self.draw_branches(branch_length, nangle, depth - 1,
 p2, branch_size, color)
 # Calculate new angle and recurse
 b = Point(p1.x, p2.y)
 self.branches.append(Branch(p1, b, color, branch_size))
 nangle = angle + randint(-1, 0) * randint(1, s)
 self.draw_branches(branch_length, nangle, depth - 1,
 b, branch_size, color)
 # Calculate new angle and recurse
 c = Point(-p2.x + 2 * p1.x, p2.y)
 self.branches.append(Branch(p1, c, color, branch_size))
 nangle = angle + randint(0, 1) * randint(1, s)
 self.draw_branches(branch_length, nangle, depth - 1,
 c, branch_size, color)
 def write_svg(self, output='drzewko.svg'):
 with open(output, 'w') as outfile:
 outfile.write('<svg xmlns="http://www.w3.org/2000/svg" '
 'viewBox="0 0 800 800" version="1.1">\n')
 outfile.writelines(map(str, self.branches))
 outfile.write('</svg>\n')
print("Start generating, please wait..")
# a starting point
resultlist = [Branch(Point(400, 800), Point(400, 600), color=60, size=35)]
# Build and save the tree as an svg
t = Tree(length=200, angle=-20, depth=9, point=Point(400, 600),
 size=35, color=60, outlist=resultlist)
t.write_svg()
print("Done, check SVG file")