model_plotter.py

"""
I'ts >not< this: https://ds055uzetaobb.cloudfront.net/brioche/uploads/lRC3MIbnUy-germany_map_2.png?width=1200 ,
i.e. it is not the classical map graph.
"""

import matplotlib.pyplot as plt
import numpy as np


def get_polygon_centroid(points, 
                         n_samples_x = 100):
    """Get the centroid of a polygon given by an ordered list of points."""

    x_min = points[:,0].min()
    x_max = points[:,0].max()
    y_min = points[:,1].min()
    y_max = points[:,1].max()

    n_samples_y = int(np.ceil(n_samples_x*(y_max-y_min)/(x_max-x_min)))

    X = np.linspace(x_min, x_max, n_samples_x)
    Y = np.linspace(y_min, y_max, n_samples_y)

    X,Y = np.meshgrid(X,Y)
    XY  = np.stack([X,Y], axis=2)

    intersects = np.zeros(XY.shape[:2], dtype = int)

    for i in range(len(points)):
        start = points[i]
        end   = points[(i+1) % len(points)]

        v_along  = end-start
        v_orth   = np.array([v_along[1],-v_along[0]])

        start_as_orig = XY-start[None,None,:]
        coord_orth    = (start_as_orig*v_orth[None,None,:] ).sum(axis=2)

        in_y_range = (XY[:,:,1] < max(start[1],end[1])) & \
                     (XY[:,:,1] > min(start[1],end[1])) 

        if end[1] > start[1]:
            on_good_side = (coord_orth > 0)
        else:
            on_good_side = (coord_orth < 0)

        intersects += on_good_side & in_y_range

    in_polygon = (intersects % 2 == 1)  

    return XY[in_polygon].mean(axis=0)


class ModelPlotter(object):
    """Class for plotting the basic structure of Ribasim models."""

    def __init__(self, **kwargs):

        self.fig, self.ax = plt.subplots(**kwargs)

        self.basin_edges       = []
        self.basin_nodes       = []
        self.basin_borders     = []
        self.basin_connections = []

        self.plot_options = dict(basin_edge_color = 'k',
                                 basin_edge_ls    = '-',
                                 basin_node_color = 'b',
                                 basin_fill_alpha = 0.5,
                                 basin_conn_lw    = 3,
                                 basin_conn_color = "C0")


    def add_basin_edge(self,start,end,
                       n_points           = 10,
                       max_deviation_frac = 0.5):
        """Add an edge of a basin geometry."""

        start = np.array(start)
        end   = np.array(end)

        direction  = (end-start)/(n_points-1)
        orthogonal = np.array([direction[1],-direction[0]])
        t          = np.linspace(0,1,n_points)

        points         = start[None,:] + (end-start)[None,:] * t[:,None]
        deviation      = (2*np.random.rand(n_points-2,1)-1)*max_deviation_frac*orthogonal
        points[1:-1]  += deviation

        self.basin_edges.append(points)


    def get_basin_perimeter(self, basin_edge_indices, edge_flips):
        """Get all the points that define the perimeter of a basin,
        as a combination of basin edges."""

        points = []

        for i, edge_flip in zip(basin_edge_indices, edge_flips):

            new_points =self.basin_edges[i]

            if edge_flip:
                new_points = new_points[::-1]

            points.append(new_points[:-1])

        return np.concatenate(points)



    def add_basin(self, basin_edge_indices, edge_flips):
        """Specify which basin edges together form a basin."""

        points   = self.get_basin_perimeter(basin_edge_indices, edge_flips)
        centroid = get_polygon_centroid(points)

        self.basin_nodes.append(centroid)
        self.basin_borders.append([basin_edge_indices,edge_flips])



    def add_basin_connection(self, basin_1, basin_2, basin_edge_index):
        """Specify which basins must be connected in the graph."""

        self.basin_connections.append([basin_1,basin_2,basin_edge_index])



    def plot(self, axis_off = True, annotate_basin_edges = False,
             legend = True):
        """Show the model."""

        ax = self.ax

        for k, (basin_edge_indices, edge_flips) in enumerate(self.basin_borders):
            points = self.get_basin_perimeter(basin_edge_indices, edge_flips)
            ax.fill(points[:,0],
                    points[:,1],
                    zorder = 0,
                    alpha = self.plot_options["basin_fill_alpha"],
                    label = f"Basin {k+1} geometry")

        for basin_edge in self.basin_edges:
            ax.plot(basin_edge[:,0],
                basin_edge[:,1],
                color = self.plot_options["basin_edge_color"],
                ls    = self.plot_options["basin_edge_ls"],
                zorder = 1)

        basin_connection_nodes = np.zeros((len(self.basin_connections),2))

        for j, (basin_1, basin_2, basin_edge_index) in enumerate(self.basin_connections):
            basin_1_loc = self.basin_nodes[basin_1]
            basin_2_loc = self.basin_nodes[basin_2]

            basin_edge     = self.basin_edges[basin_edge_index]
            basin_edge_loc = basin_edge[len(basin_edge)//2]

            basin_connection_nodes[j] = basin_edge_loc

            ax.plot([basin_1_loc[0], basin_edge_loc[0], basin_2_loc[0]],
                    [basin_1_loc[1], basin_edge_loc[1], basin_2_loc[1]],
                    lw = self.plot_options["basin_conn_lw"],
                    color = self.plot_options["basin_conn_color"],
                    zorder = 2)

        basin_nodes = np.array(self.basin_nodes)
        ax.scatter(basin_nodes[:,0],
                   basin_nodes[:,1],
                   c = self.plot_options["basin_node_color"],
                   label = 'Basin nodes',
                   zorder = 3)

        ax.scatter(basin_connection_nodes[:,0],
                   basin_connection_nodes[:,1],
                   zorder = 3,
                   label = "Basin connection nodes")

        if annotate_basin_edges:
            for i,basin_edge in enumerate(self.basin_edges):
                loc = np.mean(basin_edge, axis = 0)
                ax.annotate(i, loc, bbox = dict(fc = 'w'))



        ax.set_aspect("equal")

        if axis_off:
            ax.axis('off')

        if legend:
            ax.legend(bbox_to_anchor=(1, 1))




if __name__ == "__main__":

    np.random.seed(31415)

    model_plotter = ModelPlotter()

    for start, end in [[( 0.0, 0.0),( 1.0, 1.0)],
                       [( 1.0, 1.0),( 2.0,-1.0)],
                       [( 2.0,-1.0),( 0.0, 0.0)],
                       [( 1.0, 1.0),( 3.0, 2.0)],
                       [( 3.0, 2.0),( 3.5, 0.5)],
                       [( 3.5, 0.5),( 2.0,-1.0)],
                       [( 0.0, 0.0),( 0.5, 2.0)],
                       [( 0.5, 2.0),( 3.0, 2.0)]]:
        model_plotter.add_basin_edge(start,end, max_deviation_frac=0.7)

    model_plotter.add_basin([0,1,2], 3*[False])
    model_plotter.add_basin([1,3,4,5], [True, False, False, False])
    model_plotter.add_basin([0,3,7,6], [False, False, True, True])

    for basin_connection in [[0,1,1],
                             [0,2,0]]:
        model_plotter.add_basin_connection(*basin_connection)

    model_plotter.plot(annotate_basin_edges = False,
                       legend = True)