/ Fractals / pythagoras_tree.py
pythagoras_tree.py
  1  from PyQt5.QtCore import *
  2  from PyQt5.QtGui import *
  3  from PyQt5.QtWidgets import *
  4  
  5  from vector import *
  6  
  7  ANGLE = 30
  8  
  9  
 10  class Square:
 11      def __init__(self, p1: Point, p2: Point):
 12          self.p1 = p1
 13          self.p2 = p2
 14          vector = (p2 - p1).to_vector()
 15          vector.angle -= 90
 16          self.p3 = p2 + vector.to_point()
 17          self.p4 = p1 + vector.to_point()
 18  
 19          self.triangle = None
 20  
 21      def __iter__(self):
 22          yield self.p1, self.p2, self.p3, self.p4
 23          if self.triangle:
 24              yield from self.triangle
 25  
 26      def iterate(self):
 27          if self.triangle is None:
 28              self.triangle = Triangle(self.p4, self.p3)
 29          else:
 30              self.triangle.iterate()
 31  
 32  
 33  class Triangle:
 34      def __init__(self, p1: Point, p2: Point):
 35          self.p1 = p1
 36          self.p2 = p2
 37          vector = (p2 - p1).to_vector()
 38          vector.angle -= ANGLE
 39          vector.distance *= math.cos(deg2rad(ANGLE))
 40          self.p3 = p1 + vector.to_point()
 41  
 42          self.square1 = self.square2 = None
 43  
 44      def __iter__(self):
 45          yield self.p1, self.p2, self.p3
 46          if self.square1 is not None:
 47              yield from self.square1
 48              yield from self.square2
 49  
 50      def iterate(self):
 51          if self.square1 is None:
 52              self.square1 = Square(self.p1, self.p3)
 53              self.square2 = Square(self.p3, self.p2)
 54          else:
 55              self.square1.iterate()
 56              self.square2.iterate()
 57  
 58  
 59  class PyTree(QWidget):
 60      def __init__(self):
 61          super().__init__()
 62  
 63          self.setWindowTitle("Pythagoras Tree - Iteration 0")
 64          self.setFixedSize(900, 600)
 65  
 66          self.tree = Square(
 67              Point(self.width() / 2 - 50, self.height() - 20),
 68              Point(self.width() / 2 + 50, self.height() - 20)
 69          )
 70          self.iterations = [list(self.tree)]
 71          self.current_iteration = 0
 72  
 73          self.show()
 74  
 75      def keyPressEvent(self, e: QKeyEvent):
 76          if e.key() == Qt.Key_Q:
 77              self.close()
 78          elif e.key() == Qt.Key_Right:
 79              self.current_iteration += 1
 80              if len(self.iterations) <= self.current_iteration:
 81                  self.tree.iterate()
 82                  self.iterations.append(list(self.tree))
 83              self.setWindowTitle(f"Pythagoras Tree - Iteration {self.current_iteration}")
 84              self.repaint()
 85          elif e.key() == Qt.Key_Left:
 86              if self.current_iteration > 0:
 87                  self.current_iteration -= 1
 88                  self.setWindowTitle(f"Pythagoras Tree - Iteration {self.current_iteration}")
 89                  self.repaint()
 90  
 91      def paintEvent(self, _: QPaintEvent):
 92          qp = QPainter(self)
 93          qp.setPen(Qt.white)
 94          qp.setBrush(Qt.white)
 95          qp.drawRect(self.rect())
 96  
 97          qp.setPen(Qt.black)
 98  
 99          for polygon in self.iterations[self.current_iteration]:
100              qp.drawPolygon(*(QPoint(*p) for p in polygon))
101  
102  
103  if __name__ == '__main__':
104      app = QApplication([])
105      tree = PyTree()
106      app.exec_()