Symmetric tree using `pyplot`

Straight Tree.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Straight Tree\n",
    "\n",
    "Plots a pretty symmetric simple model of a tree"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Declare a simple Dot structure that stores co-ordinates:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Dot:\n",
    "    def __init__(self, x=0, y=0):\n",
    "        self.x = x\n",
    "        self.y = y\n",
    "    \n",
    "    def __iter__(self):\n",
    "        return iter((self.x, self.y))\n",
    "    \n",
    "    def __repr__(self):\n",
    "        return f'<Dot {self.x:.2f} {self.y:.2f}>'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can describe a Vector and its methods for the Cartesian-Polar transformation and vice versa:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "\n",
    "class Vector:\n",
    "    @staticmethod    \n",
    "    def angle_to_dots(angle, length, start=None):\n",
    "        if not start:\n",
    "            start = Dot()\n",
    "        \n",
    "        x = start.x + math.sin(math.radians(angle)) * length\n",
    "        y = start.y + math.cos(math.radians(angle)) * length\n",
    "        \n",
    "        return start, Dot(x, y)\n",
    "    \n",
    "    @staticmethod\n",
    "    def dots_to_angle(start, end):\n",
    "        \n",
    "        a = math.degrees(math.atan((end.y - start.y) / (end.x - start.x)))\n",
    "        l = math.sqrt(abs((end.x - start.x) ** 2 + (end.y - start.y) ** 2))\n",
    "        \n",
    "        return start, a, l\n",
    "    \n",
    "    def __init__(self, start, end=None, angle=None, length=None):        \n",
    "        self.start = start\n",
    "        \n",
    "        if end:\n",
    "            self.end = end\n",
    "            _, self.angle, self.length = self.dots_to_angle(start, end)\n",
    "        else:            \n",
    "            self.angle = angle\n",
    "            self.length = length\n",
    "            _, self.end = self.angle_to_dots(angle, length, start)\n",
    "    \n",
    "    def __iter__(self):\n",
    "        return iter((self.start, self.end))\n",
    "    \n",
    "    def __repr__(self):\n",
    "        return f'<Line {self.start} {self.end} l={self.length:.2f} a={self.angle:.2f}>'\n",
    "    \n",
    "    def flat(self):\n",
    "        return tuple(map(tuple, zip(*self)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our configurable Tree class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Tree:\n",
    "    def __init__(self, parent=None, length=10, k=0.8, angle=0, c=30, nodes=4, depth=6):\n",
    "        self.parent = parent\n",
    "        \n",
    "        if self.parent:\n",
    "            self.start = self.parent.vector.end\n",
    "        else:\n",
    "            self.start = Dot(0,0)\n",
    "        \n",
    "        self.length = length\n",
    "        self.angle = angle\n",
    "        self.depth = depth\n",
    "        self.vector = Vector(self.start, angle=self.angle, length=self.length)\n",
    "        \n",
    "        self.nodes = []\n",
    "        \n",
    "        if depth > 0:\n",
    "            for child in range(nodes):\n",
    "                self.nodes.append(Tree(\n",
    "                    parent=self, \n",
    "                    length=k * length, \n",
    "                    k=k,\n",
    "                    angle=(angle - c / 2) + child * (c / (nodes - 1)),\n",
    "                    c=c,\n",
    "                    nodes=nodes,\n",
    "                    depth=depth - 1\n",
    "                    ))\n",
    "    \n",
    "    def __repr__(self):\n",
    "        return f'<Tree {self.vector} nodes={self.nodes}>'\n",
    "    \n",
    "    def flat(self):\n",
    "        flats = [(self.depth, self.vector.flat())]\n",
    "        for node in self.nodes:\n",
    "            node_flats = node.flat()\n",
    "            for node_flat in node_flats:\n",
    "                if isinstance(node_flat, list):\n",
    "                    flats.extend(node_flat)\n",
    "                else:\n",
    "                    flats.append(node_flat)\n",
    "        return flats\n",
    "                              "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Build a Tree and flatten branch vectors to the plain list:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "flats = Tree(length=30, k=0.8, depth=5, c=60, nodes=4).flat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "for weight, l in flats:\n",
    "    plt.plot(*l, linewidth=weight)\n",
    "\n",
    "\n",
    "plt.axis('equal')\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}