{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Arrow Demo\n\nThree ways of drawing arrows to encode arrow \"strength\" (e.g., transition\nprobabilities in a Markov model) using arrow length, width, or alpha (opacity).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import itertools\n\nimport matplotlib.pyplot as plt\nimport numpy as np\n\n\ndef make_arrow_graph(ax, data, size=4, display='length', shape='right',\n                     max_arrow_width=0.03, arrow_sep=0.02, alpha=0.5,\n                     normalize_data=False, ec=None, labelcolor=None,\n                     **kwargs):\n    \"\"\"\n    Makes an arrow plot.\n\n    Parameters\n    ----------\n    ax\n        The Axes where the graph is drawn.\n    data\n        Dict with probabilities for the bases and pair transitions.\n    size\n        Size of the plot, in inches.\n    display : {'length', 'width', 'alpha'}\n        The arrow property to change.\n    shape : {'full', 'left', 'right'}\n        For full or half arrows.\n    max_arrow_width : float\n        Maximum width of an arrow, in data coordinates.\n    arrow_sep : float\n        Separation between arrows in a pair, in data coordinates.\n    alpha : float\n        Maximum opacity of arrows.\n    **kwargs\n        `.FancyArrow` properties, e.g. *linewidth* or *edgecolor*.\n    \"\"\"\n\n    ax.set(xlim=(-0.25, 1.25), ylim=(-0.25, 1.25), xticks=[], yticks=[],\n           title=f'flux encoded as arrow {display}')\n    max_text_size = size * 12\n    min_text_size = size\n    label_text_size = size * 4\n\n    bases = 'ATGC'\n    coords = {\n        'A': np.array([0, 1]),\n        'T': np.array([1, 1]),\n        'G': np.array([0, 0]),\n        'C': np.array([1, 0]),\n    }\n    colors = {'A': 'r', 'T': 'k', 'G': 'g', 'C': 'b'}\n\n    for base in bases:\n        fontsize = np.clip(max_text_size * data[base]**(1/2),\n                           min_text_size, max_text_size)\n        ax.text(*coords[base], f'${base}_3$',\n                color=colors[base], size=fontsize,\n                horizontalalignment='center', verticalalignment='center',\n                weight='bold')\n\n    arrow_h_offset = 0.25  # data coordinates, empirically determined\n    max_arrow_length = 1 - 2 * arrow_h_offset\n    max_head_width = 2.5 * max_arrow_width\n    max_head_length = 2 * max_arrow_width\n    sf = 0.6  # max arrow size represents this in data coords\n\n    if normalize_data:\n        # find maximum value for rates, i.e. where keys are 2 chars long\n        max_val = max((v for k, v in data.items() if len(k) == 2), default=0)\n        # divide rates by max val, multiply by arrow scale factor\n        for k, v in data.items():\n            data[k] = v / max_val * sf\n\n    # iterate over strings 'AT', 'TA', 'AG', 'GA', etc.\n    for pair in map(''.join, itertools.permutations(bases, 2)):\n        # set the length of the arrow\n        if display == 'length':\n            length = (max_head_length\n                      + data[pair] / sf * (max_arrow_length - max_head_length))\n        else:\n            length = max_arrow_length\n        # set the transparency of the arrow\n        if display == 'alpha':\n            alpha = min(data[pair] / sf, alpha)\n        # set the width of the arrow\n        if display == 'width':\n            scale = data[pair] / sf\n            width = max_arrow_width * scale\n            head_width = max_head_width * scale\n            head_length = max_head_length * scale\n        else:\n            width = max_arrow_width\n            head_width = max_head_width\n            head_length = max_head_length\n\n        fc = colors[pair[0]]\n\n        cp0 = coords[pair[0]]\n        cp1 = coords[pair[1]]\n        # unit vector in arrow direction\n        delta = cos, sin = (cp1 - cp0) / np.hypot(*(cp1 - cp0))\n        x_pos, y_pos = (\n            (cp0 + cp1) / 2  # midpoint\n            - delta * length / 2  # half the arrow length\n            + np.array([-sin, cos]) * arrow_sep  # shift outwards by arrow_sep\n        )\n        ax.arrow(\n            x_pos, y_pos, cos * length, sin * length,\n            fc=fc, ec=ec or fc, alpha=alpha, width=width,\n            head_width=head_width, head_length=head_length, shape=shape,\n            length_includes_head=True,\n            **kwargs\n        )\n\n        # figure out coordinates for text:\n        # if drawing relative to base: x and y are same as for arrow\n        # dx and dy are one arrow width left and up\n        orig_positions = {\n            'base': [3 * max_arrow_width, 3 * max_arrow_width],\n            'center': [length / 2, 3 * max_arrow_width],\n            'tip': [length - 3 * max_arrow_width, 3 * max_arrow_width],\n        }\n        # for diagonal arrows, put the label at the arrow base\n        # for vertical or horizontal arrows, center the label\n        where = 'base' if (cp0 != cp1).all() else 'center'\n        # rotate based on direction of arrow (cos, sin)\n        M = [[cos, -sin], [sin, cos]]\n        x, y = np.dot(M, orig_positions[where]) + [x_pos, y_pos]\n        label = r'$r_{_{\\mathrm{%s}}}$' % (pair,)\n        ax.text(x, y, label, size=label_text_size, ha='center', va='center',\n                color=labelcolor or fc)\n\n\nif __name__ == '__main__':\n    data = {  # test data\n        'A': 0.4, 'T': 0.3, 'G': 0.6, 'C': 0.2,\n        'AT': 0.4, 'AC': 0.3, 'AG': 0.2,\n        'TA': 0.2, 'TC': 0.3, 'TG': 0.4,\n        'CT': 0.2, 'CG': 0.3, 'CA': 0.2,\n        'GA': 0.1, 'GT': 0.4, 'GC': 0.1,\n    }\n\n    size = 4\n    fig = plt.figure(figsize=(3 * size, size), layout=\"constrained\")\n    axs = fig.subplot_mosaic([[\"length\", \"width\", \"alpha\"]])\n\n    for display, ax in axs.items():\n        make_arrow_graph(\n            ax, data, display=display, linewidth=0.001, edgecolor=None,\n            normalize_data=True, size=size)\n\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.13.2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}