polar-sun-plot

polar_sun_plot.py

#!/usr/bin/env python

"""This plots a sun-like design with alternating curves
The intention is for the curves to alternate being over and under as in Celtic knots
"""

import matplotlib.pyplot as plt
from matplotlib import colors
import numpy as np
import argparse

class WeaveError(ValueError):
    pass

class WeaveCalculator(object):
    def __init__(self, points, curves, debug=False):
        self.points = points
        self.curves = curves
        self.debug = debug

    def debug_log(self, message):
        if self.debug:
            print(message)

    def calc_intersection(self, point, curve):
        return (self.curves*3-1 - point - curve) % self.curves

    def calc_ordering(self):
        in_front = {}
        at_back = {}
        order_segments = self.points * self.curves * 2
        for c in range(0, self.curves):
            n_int = c
            polarity_switch = None
            for tc in range(order_segments):
                existing_order = in_front.get(tc, []) + at_back.get(tc, [])
                iw = self.calc_intersection(tc, c)
                if (iw != c):
                    n_int = n_int + 1
                    polarity_adjusted_int = n_int + (polarity_switch or 0)
                    ordering, actual_ordering = bool(polarity_adjusted_int % 2), None
                    if c in existing_order or iw in existing_order:
                        if c not in existing_order or iw not in existing_order:
                            raise ValueError("Found partial ordering half way through")
                        actual_ordering = existing_order.index(c) < existing_order.index(iw)
                    if polarity_switch is None and actual_ordering is not None:
                        polarity_switch = 1 - int(bool(actual_ordering == ordering))
                        if polarity_switch:
                            self.debug_log(f"Color {c} needed a polarity switch for " \
                                           f"ordering {'<>'[ordering]}, actual ordering {'<>'[actual_ordering]}")
                        else:
                            self.debug_log(f"Color {c} didn't need a polarity switch for " \
                                           f"ordering {'<>'[ordering]}, actual ordering {'<>'[actual_ordering]}")
                        ordering = (not ordering) if polarity_switch else ordering
                    if actual_ordering is not None:
                        self.debug_log(f"priority({c:2},{tc:2}) is {'<>'[actual_ordering]} priority({iw:2},{tc:2}): " \
                                       f"existing order {', '.join(str(o) for o in existing_order)}")
                        if actual_ordering != ordering:
                            raise WeaveError("Actual ordering doesn't match expected")
                        continue
                    if ordering:
                        in_front.setdefault(tc, []).append(c)
                        at_back.setdefault(tc, []).append(iw)
                    else:
                        at_back.setdefault(tc, []).append(c)
                        in_front.setdefault(tc, []).append(iw)
                    self.debug_log(f"priority({c:2},{tc:2}) should be {'<>'[ordering]} priority({iw:2},{tc:2})")
        return {tc: in_front.get(tc, []) + at_back.get(tc, []) for tc in range(order_segments)}

can_order = {}
for p in range(2, 24):
    for c in range(2, 12):
        try:
            WeaveCalculator(p, c).calc_ordering()
            can_order[p, c] = True
        except WeaveError:
            can_order[p, c] = False

class SunCurve(WeaveCalculator):
    points = 4 # how many peaks each curve has
    curves = 3
    peak = 2
    trough = 1.1
    line_width = 2
    colors = []
    bg_colors = []
    bg_alphas = []
    accuracy = 1440 # points per peak
    weave = True
    debug = False

    def make_default_colors(self, alpha=1.0):
        hues = np.linspace(54, 36, self.curves)/360
        sats = np.linspace(93, 93, self.curves)/100
        values = np.linspace(100, 100, self.curves)/100
        hsv = np.array([hues, sats, values]).transpose()
        return [colors.to_rgba(c, 1.0) for c in colors.hsv_to_rgb(hsv)]

    def make_default_bg_colors(self):
        return self.make_default_colors()

    def __init__(self, args):
        # args can be set to match any named class attribute
        for arg in dir(args):
            if arg.startswith('_'):
                continue
            if hasattr(type(self), arg):
                arg_value = getattr(args, arg)
                setattr(self, arg, arg_value)
        if not self.colors:
            self.colors = self.make_default_colors()
        if not self.bg_colors:
            self.bg_colors = self.make_default_bg_colors()
        if not self.bg_alphas:
            self.bg_alphas = [0.5]

    @property
    def theta(self):
        if getattr(self, "_theta", None) is None:
            offset = np.pi/self.points/self.curves/2
            self._theta = np.arange(offset, 2*np.pi + offset, 2*np.pi/(self.accuracy*self.points))
        return self._theta

    # The peak function is graphed at https://www.geogebra.org/graphing/d7ezedpd
    # This is repeated semicircles. They are then calculated in polar co-ordinates

    def calc_curve(self, theta_points, c):
        mod_func = np.mod(theta_points*self.curves*self.points/(2*np.pi) + c, self.curves)
        return self.peak - np.sqrt(1 - (mod_func * 2/self.curves - 1)**2)*(self.peak - self.trough)

    def setup_plot(self):
        fig, ax = plt.subplots(subplot_kw={'projection': 'polar'})
        ax.set_rmax(self.peak)
        ax.set_rticks([])  # No radial ticks
        ax.set_axis_off()
        ax.grid(False)
        # ax.set_title(f"Sun with two alternating plots of {points} points", va='bottom')
        return ax

    def get_color(self, c):
        return self.colors[c % len(self.colors)]

    def get_bg_color(self, c):
        return self.bg_colors[c % len(self.bg_colors)]

    def get_bg_alpha(self, c):
        return self.bg_alphas[c % len(self.bg_alphas)]

    def plot_background(self, ax):
        # colouring in areas
        full_curves = np.array([self.calc_curve(self.theta, c) for c in range(self.curves)])
        full_curves.sort(axis=0)
        bg_curves = [0]
        for c in range(self.curves):
            bg_curve = full_curves[c]
            bg_curves.append(bg_curve)

        for c in range(self.curves):
            ax.fill_between(self.theta, bg_curves[c], bg_curves[c+1],
                            color=self.get_bg_color(c), alpha=self.get_bg_alpha(c), linewidth=0, zorder=-2)

    def plot_weave(self, ax):
        full_ordering = self.calc_ordering()
        for tc, tc_order in full_ordering.items():
            self.debug_log(f"ordering at {tc:2}: {', '.join(str(o) for o in tc_order)}")

        # these segments are separated to allow different orderings to create a weave
        theta_parts = np.split(self.theta, len(full_ordering))
        for c in range(0, self.curves):
            color = self.get_color(c)
            for tc, theta_c_segment in enumerate(theta_parts):
                tc_order = full_ordering[tc]
                rc = self.calc_curve(theta_c_segment, c)
                priority_c = tc_order.index(c) if c in tc_order else -1
                ax.plot(theta_c_segment, rc, color=color, linewidth=self.line_width, zorder=priority_c)
                if self.debug:
                    mp = int(len(theta_c_segment)/2)
                    iw = self.calc_intersection(tc, c)
                    debug_text = f"{tc%self.curves},{"/" if priority_c == -1 else priority_c},{"/" if iw == c else iw}"
                    text_rot = theta_c_segment[mp]*180/np.pi-90
                    ax.text(theta_c_segment[mp], self.trough-0.1*(c+1), debug_text,
                            color=c, ha='center', va='center', rotation=text_rot, rotation_mode='anchor')

    def plot_non_weave(self, ax):
        for c in range(self.curves):
            ax.plot(self.theta, self.calc_curve(self.theta, c), color=self.get_color(c), linewidth=self.line_width)

    def plot_foreground(self, ax):
        if self.weave:
            try:
                self.plot_weave(ax)
            except WeaveError as e:
                print(f"Can't calculate a weave for {self.points} points and {self.curves} curves: plotting unwoven")
                self.plot_non_weave(ax)
        else:
            self.plot_non_weave(ax)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--debug', action='store_true', default=False, help="Add debug text")
    parser.add_argument('-p', '--points', type=int, default=SunCurve.points, help="Number of points to each curve")
    parser.add_argument('-c', '--curves', type=int, default=SunCurve.curves, help="Number of curves")
    parser.add_argument('--peak', type=float, default=SunCurve.peak, help="Radius of peaks")
    parser.add_argument('--trough', type=float, default=SunCurve.trough, help="Radius of troughs")
    parser.add_argument('--accuracy', type=int, default=SunCurve.accuracy, help="Points between two peaks")
    parser.add_argument('-l', '--line-width', type=float, default=SunCurve.line_width, help="Width of lines")
    parser.add_argument('--no-weave', action='store_false', dest='weave', default=SunCurve.weave,
                        help="Don't try to weave curves")
    parser.add_argument('-C', '--color', action='append', dest='colors', default=[],
                        help="Custom color to use for foreground plot")
    parser.add_argument('-B', '--bg-color', action='append', dest='bg_colors', default=[],
                        help="Custom color to use for background plot")
    parser.add_argument('-A', '--bg-alpha', action='append', dest='bg_alphas', default=[], type=float,
                        help="Custom alpha to use for background plot")
    parser.add_argument('filename', nargs='?', help="Filename to save to instead of plotting to screen")
    args = parser.parse_args()
    sun_curve = SunCurve(args)
    ax = sun_curve.setup_plot()
    sun_curve.plot_background(ax)
    sun_curve.plot_foreground(ax)
    if args.filename:
        plt.savefig("polar_sun_plot.svg")
    else:
        plt.show()

requirements.txt

matplotlib
numpy