plotter_experiments/lib/plot.py

#!/usr/bin/env python

import cairo
import numpy as np


class Paper:
    def __init__(self, width, height, margins=None):
        self.width = width
        self.height = height
        self.margins = {
            'top': 0,
            'right': 0,
            'bottom': 0,
            'left': 0
        }
        self.set_margins(margins)
        

    def set_margins(self, margins):
        if isinstance(margins, int) or isinstance(margins, float):
            self.margins['top'] = margins
            self.margins['right'] = margins
            self.margins['bottom'] = margins
            self.margins['left'] = margins
        elif isinstance(margins, list) and (isinstance(margins[0], int)
                                            or isinstance(margins[0], float)):
            if len(margins) == 2:
                self.margins['top'] = margins[0]
                self.margins['bottom'] = margins[0]
                self.margins['left'] = margins[1]
                self.margins['right'] = margins[1]
            elif len(margins) == 3:
                self.margins['top'] = margins[0]
                self.margins['left'] = margins[1]
                self.margins['right'] = margins[1]
                self.margins['bottom'] = margins[2]
            elif len(margins) == 4:
                self.margins['top'] = margins[0]
                self.margins['right'] = margins[1]
                self.margins['bottom'] = margins[2]
                self.margins['left'] = margins[3]
        else:
            self.margins = margins

    @property
    def content_width(self):
        return self.width - self.margins['left'] - self.margins['right']

    @property
    def content_height(self):
        return self.height - self.margins['top'] - self.margins['bottom']

    def size(self, honour_margins=True):
        if honour_margins:
            return np.array([self.content_width, self.content_height])
        else:
            return np.array([self.width, self.height])

    def is_within(self, point, honour_margins=True):
        return ((self.left(honour_margins) < point[0])
                and (point[0] < self.right(honour_margins))
                and (self.top(honour_margins) < point[1])
                and (point[1] < self.bottom(honour_margins)))

    def top_left(self, honour_margins=True):
        if honour_margins:
            return np.array([self.margins['left'],
                             self.height - self.margins['top']])
        else:
            return np.array([0, self.height])

    def bottom_left(self, honour_margins=True):
        if honour_margins:
            return np.array([self.margins['left'],
                             self.margins['bottom']])
        else:
            return np.array([0, 0])

    def top_right(self, honour_margins=True):
        if honour_margins:
            return np.array([self.width - self.margins['right'],
                             self.height - self.margins['top']])
        else:
            return np.array([self.width, self.height])

    def bottom_right(self, honour_margins=True):
        if honour_margins:
            return np.array([self.width - self.margins['right'],
                             self.margins['bottom']])
        else:
            return np.array([self.width, 0])

    def centre(self, honour_margins=True):
        return (self.bottom_left(honour_margins)
                + self.top_right(honour_margins)) / 2

    def top(self, honour_margins=True):
        if honour_margins:
            return self.margins['top']
        else:
            return 0

    def left(self, honour_margins=True):
        if honour_margins:
            return self.margins['left']
        else:
            return 0

    def right(self, honour_margins=True):
        if honour_margins:
            return self.width - self.margins['right']
        else:
            return self.width

    def bottom(self, honour_margins=True):
        if honour_margins:
            return self.height - self.margins['bottom']
        else:
            return self.height


A4_PORTRAIT = Paper(210, 297, 20)
A6_PORTRAIT = Paper(105, 148, 8)



class Plotter:

    def __init__(self, paper, line_width=0, colour=[0, 0, 0, 1]):
        self.paper = paper
        self.pen_down = False
        self.position = np.array([0, 0])

        self.layers = []
        self.add_layer(colour, line_width)

    def add_layer(self, colour=[0, 0, 0, 1], line_width=0.5, name='',
                  switch=True):
        self.layers.append({
            'line_width': line_width,
            'colour': colour,
            'name': name
        })
        if switch:
            self.switch_layer(len(self.layers) - 1)
        return len(self.layers) - 1

    def switch_layer(self, layer):
        self.current_layer = layer

    def move_to(self, point):
        self.pen_down = False
        self.position = point

    def line_to(self, point):
        self.pen_down = True
        self.position = point

    def draw_frame(self):
        self.move_to(self.paper.bottom_left())
        self.line_to(self.paper.bottom_right())
        self.line_to(self.paper.top_right())
        self.line_to(self.paper.top_left())
        self.line_to(self.paper.bottom_left())


class SVGPlotter(Plotter):

    def __init__(self, file_name, paper, line_width=0.5, colour=[0, 0, 0, 1]):
        dpi = 72
        self.file_name = file_name
        self.surface = cairo.SVGSurface(file_name,
                                        paper.width / 25.4 * dpi,
                                        paper.height / 25.4 * dpi)
        self.context = cairo.Context(self.surface)
        self.context.scale(dpi / 25.4, dpi / 25.4)
        self.context.set_line_cap(cairo.LINE_CAP_ROUND)

        super().__init__(paper, line_width, colour)

    def move_to(self, point):
        if self.pen_down:
            self.context.stroke()
        self.context.move_to(point[0], point[1])
        super().move_to(point)

    def line_to(self, point):
        self.context.line_to(point[0], point[1])
        super().line_to(point)

    def finalise(self):
        if self.pen_down:
            self.context.stroke()
            self.pen_down = False
        self.surface.finish()

    def switch_layer(self, layer):
        if self.pen_down:
            self.context.stroke()

        super().switch_layer(layer)
        self.context.set_source_rgba(
            *self.layers[self.current_layer]['colour'])
        self.context.set_line_width(
            self.layers[self.current_layer]['line_width'])

        return len(self.layers) - 1


class HPGLPlotter(Plotter):

    def __init__(self, paper, file_name_pattern):
        self.file_name_pattern = file_name_pattern
        self.resolution = (4000 / 81.5, 4000 / 80.5)
        self.file = None
        super().__init__(paper)

    def _layer_file_name(self, index):
        return self.file_name_pattern.format(
            index=index,
            name=self.layers[index]['name'],
            colour=self.layers[index]['colour']
        )

    def _mm_to_plotter(self, mm):
        p = (mm[0] * self.resolution[0],
             (self.paper.height - mm[1]) * self.resolution[1])
        return f'{p[0]:.0f},{p[1]:.0f}'

    def move_to(self, point):
        if self.pen_down:
            self.file.write(';PU')
        else:
            self.file.write(',')
        self.file.write(self._mm_to_plotter(point))
        return super().move_to(point)

    def line_to(self, point):
        if not self.pen_down:
            self.file.write(';PD')
        else:
            self.file.write(',')
        self.file.write(self._mm_to_plotter(point))
        return super().line_to(point)

    def add_layer(self, colour='black', line_width=0.5, name='', switch=True):
        new_index = super().add_layer(colour=colour, line_width=line_width,
                                      name=name, switch=False)
        layer_file = open(self._layer_file_name(new_index), 'wt')

        layer_file.write('IN;PU0,0')

        if switch:
            if self.file:
                self.file.close()
            self.file = layer_file
            self.current_layer = len(self.layers) - 1
        else:
            layer_file.close()

        return len(self.layers) - 1

    def switch_layer(self, layer):
        if self.pen_down:
            self.file.write(';PU')
            self.pen_down = False
        self.file.close()
        super().switch_layer(layer)
        self.file = open(self._layer_file_name(self.current_layer), 'at')

    def finalise(self):
        self.file.close()
        for i in range(len(self.layers)):
            with open(self._layer_file_name(i), 'at') as file:
                file.write(';PU0,0;IN')


class MultiPlotter(Plotter):
    def __init__(self):
        self.plotters = []

    def move_to(self, point):
        for plotter in self.plotters:
            plotter.move_to(point)

    def line_to(self, point):
        for plotter in self.plotters:
            plotter.line_to(point)

    def add_layer(self, *args, **kwargs):
        for plotter in self.plotters:
            plotter.add_layer(*args, **kwargs)

    def switch_layers(self, *args, **kwargs):
        for plotter in self.plotters:
            plotter.switch_layers(*args, **kwargs)

    def finalise(self):
        for plotter in self.plotters:
            plotter.finalise()

    def register_plotter(self, plotter):
        self.plotters.append(plotter)

    @property
    def paper(self):
        return self.plotters[0].paper