# Copyright (c) 2020. Robin Thibaut, Ghent University
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import cm, colorbar, colors
from matplotlib.collections import PatchCollection
from matplotlib.patches import Polygon
from matplotlib.path import Path
from matplotlib.ticker import LogFormatter
from matplotlib.widgets import PolygonSelector, TextBox
[docs]def find_norm(l): # Convert values to linear space to facilitate visualization!
uv = list(dict.fromkeys(l))
uv.sort()
ls = np.linspace(0, 1, len(uv))
nl = []
for i in range(len(l)):
c = 0
for j in range(len(uv)):
if l[i] <= uv[j] and not c:
nl.append(ls[j])
c += 1
return nl
[docs]class ModelMaker(object):
[docs] def __init__(
self,
model_name=None,
centerxy=np.array([]),
blocks=np.array([]),
values=np.array([]),
values_log=0,
bck=1,
):
"""
:param model_name: path to the file to be created containing the output
:param centerxy: x-y coordinates of the center of the cells
:param blocks: coordinates of the corners of the different blocks
:param values: array containing the value assigned to each block
:param values_log: flag indicating if values should be log transformed or not
:param bck: background value
"""
if centerxy.any(): # If block center coordinates are provided, plot them.
xs = centerxy[:, 0]
ys = centerxy[:, 1]
fig, ax = plt.subplots()
pts = ax.scatter(xs, ys, c="black", alpha=0.5)
else: # Else, if only blocks are provided, the centers are computed based on the mean of the coordinates of
# their corners
try:
centerxy = np.array(
[np.mean(b, axis=0) for b in blocks]
) # Mean computed
xs = centerxy[:, 0]
ys = centerxy[:, 1]
fig, ax = plt.subplots()
pts = ax.scatter(xs, ys, c="black", alpha=0.5)
except Exception as e:
print(e)
exit()
# I use this option to put the points in first plan, as to always see them.
pts.set_zorder(2)
self.points = centerxy
self.bck = bck
self.model_name = model_name
self.ax = ax # Axes object
self.blocks = blocks
self.ax.set_title("Now setting value: 0.0") # Initial title
self.ax.set_facecolor((0.86, 0.86, 0.86)) # Set gray background
# Adjusts plt dimensions to insert colorbar, textbox...
plt.subplots_adjust(bottom=0.2, right=0.8)
self.vals = (
[]
) # List that will contain the values assigned to the different polygons
self.cmap = cm.get_cmap("jet") # Color map function to be later used
# It is necessary to define a different axes for the box (normalized)
axbox = plt.axes([0.4, 0.07, 0.1, 0.07])
# 4 - tuple of floats
# rect = [left, bottom, width, height]. A new axes is added
# with dimensions rect in normalized (0, 1) units using ~.Figure.add_axes on the current figure.
# Location of colorbar for the user-defined polygons
self.axcb = plt.axes([0.82, 0.07, 0.015, 0.83])
bounds = [0] + [1] # Adding one bound more to have enough intervals
ticks = [0] # Ticks - default value
cols = colors.ListedColormap([self.cmap(v) for v in ticks])
cbnorm = colors.BoundaryNorm(bounds, cols.N)
mcb = colorbar.ColorbarBase(
self.axcb,
cmap=cols,
norm=cbnorm,
boundaries=bounds,
ticks=ticks,
ticklocation="right",
orientation="vertical",
)
#
# textstr = """Select points in the figure by enclosing them within a polygon.
# Press the 'esc' key to start a new polygon.
# Try holding the 'shift' key to move all of the vertices.
# Try holding the 'ctrl' key to move a single vertex."""
#
# #axtxt = plt.axes([0.15, 0.0, 0.2, 0.15])
# props = dict(boxstyle='round', facecolor='green', alpha=0.5)
# ax.text(0, -0.2, textstr, transform=ax.transAxes, fontsize=10, bbox=props)
# Text box to input the value to input
self.vinput = TextBox(axbox, label=None, initial="0")
# What happens when pressing enter
self.vinput.on_submit(self.button_submit)
self.index = [] # List that will contain the final results !
# Creates a canvas from the ax of the scatter plot
self.canvas = self.ax.figure.canvas
self.collection = pts # 'collection' is the scatter plot
# Necessary for later to define if points enclosed by polygon - basically
self.xys = pts.get_offsets()
# equals to the - x-y coordinates of the different points.
self.Npts = len(self.xys) # Number of points
# Ensure that we have separate colors for each object
self.fc = pts.get_facecolors() # Gets the rgb of the points
if not values.any():
facecolors = "gray"
alpha = 0.35 # Opacity of the polygons - if no values assigned - soft gray
else:
cmap2 = cm.get_cmap("coolwarm")
if values_log:
# Making a nice linear space
itv = 10 ** np.linspace(
np.log10(min(values)), np.log10(max(values)), 12
)
# out ouf log values to represent some ticks on the colorbar
norm2 = colors.LogNorm(
vmin=min(values), vmax=max(values)
) # Log norm for color bar
# Necessary arg to produce the color scale
formatter = LogFormatter(10, labelOnlyBase=False)
else:
itv = np.linspace(min(values), max(values), 8) # Linear space
norm2 = colors.Normalize(vmin=min(values), vmax=max(values))
formatter = None
# Individual color of each polygon
facecolors = [cmap2(norm2(v)) for v in values]
alpha = 0.6 # Opacity of the polygons
# Colorbar if initial values present - plotting a nice color bar
ticks2 = [round(v, 1) for v in itv]
plt.subplots_adjust(left=0.2, bottom=0.2, right=0.8)
axcb1 = plt.axes([0.15, 0.07, 0.015, 0.83])
cb1 = colorbar.ColorbarBase(
axcb1,
cmap=cmap2,
norm=norm2,
ticks=ticks2,
boundaries=None,
ticklocation="left",
format=formatter,
orientation="vertical",
)
cb1.set_ticklabels(ticks2) # Setting the proper labels
if (
self.blocks.any()
): # If blocks are provided. I should change this as the direction this code is going is
# to provide blocks by default
xs = self.blocks[:, :, 0] # x-coordinates blocks corners
ys = self.blocks[:, :, 1] # y-coordinates blocks corners
patches = [] # Coloring the blocks, in gray or with different colors
for b in blocks:
polygon = Polygon(b, closed=True)
patches.append(polygon)
p = PatchCollection(
patches, alpha=alpha, facecolors=facecolors, edgecolors="black"
)
p.set_zorder(0)
self.ax.add_collection(p)
# 5% padding in x-direction for visualization
padx = (xs.max() - xs.min()) * 0.05
pady = (ys.max() - ys.min()) * 0.05 # 5% padding
self.ax.set_xlim(xs.min() - padx, xs.max() + padx)
self.ax.set_ylim(ys.min() - pady, ys.max() + pady)
self.collection.set_facecolors(facecolors) # Coloring the points
# Polygon selector object
self.poly = PolygonSelector(self.ax, self.onselect)
self.ind = [] # Initiates the ind list for a new polygon!
def handle_close(evt): # Function that disconnects the PolygonSelector
self.disconnect()
# When closing window, finishes the job
self.canvas.mpl_connect("close_event", handle_close)
# Final results, array filled with background value.
self.final_results = np.ones(len(self.points)) * self.bck
[docs] def onselect(self, verts):
"""
Defines what happens when a polygon is completed
:param verts: polygon vertices automatically returned by the Polygon Selector
:return:
"""
if len(self.vals) == 0: # If the user forgets to press enter to initiate value
self.vals.append(0)
print("Press enter in the text box")
# When a polygon is completed or modified after completion, the `onselect`
path = Path(verts)
# function is called and passed a list of the vertices as ``(xdata, ydata)`` tuples.
self.ind = np.nonzero(path.contains_points(self.xys))[0] #
# xycenter = self.xys[self.ind]
# xycenter.data
self.index.append([self.vals[-1], self.ind, verts])
# OP : Value, Index #, XY Coordinates selected points, Polygon vtx
# Lists all values assigned to each polygon,
current_vals = [self.index[i][0] for i in range(len(self.index))]
# stored in the 'index' list
vals_cb = list(dict.fromkeys(current_vals)) # Removes duplicates
lsps = find_norm(current_vals) # Get colors
cc = [self.cmap(v) for v in lsps] # Colors used for the points/blocks
ccb = lsps.copy() # Copy list
ccb = list(dict.fromkeys(ccb)) # Removes duplicates
ccb.sort()
# List of colors for the colorbar
cols = colors.ListedColormap([self.cmap(v) for v in ccb])
if (
len(vals_cb) > 1
): # Creating a colormap fn to automatically rescale the colors of selected zones
try:
vals_cb.sort()
# Adding one bound more to have enough intervals
bounds = vals_cb + [max(vals_cb) * 1.1]
ticks = vals_cb # Ticks labels
cbnorm = colors.BoundaryNorm(bounds, cols.N)
mcb = colorbar.ColorbarBase(
self.axcb,
cmap=cols,
norm=cbnorm,
boundaries=bounds,
ticks=ticks,
ticklocation="right",
orientation="vertical",
)
except:
pass
else:
pass
if (
not self.blocks.any()
): # If no blocks, simply coloring the points accordingly
for i in range(len(self.index)):
self.fc[self.index[i][1]] = self.cmap(lsps[i])
self.collection.set_facecolors(self.fc) # updates points color
else:
if len(self.index) > 0:
for i in range(len(self.index)): # If blocks
# Concerned mesh blocks
bid = self.blocks[self.index[i][1]]
patches = []
for b in bid:
polygon = Polygon(b, closed=True)
patches.append(polygon)
p = PatchCollection(patches, alpha=1, facecolors=cc[i])
p.set_zorder(1)
self.ax.add_collection(p)
else:
pass
# self.collection.set_facecolors(self.fc) # updates points color
self.canvas.draw()
# Necessary to copy to not mess with the original array
idx = np.copy(self.index)
# What you see is what you get
# Removes duplicates from older selections, scanning for the latest
for i in range(len(idx) - 1, 0, -1):
# polygons first to the last.
indc = idx[i][1]
indp = idx[i - 1][1]
upd = np.array(list(set(indp) - set(indc)))
# Updates the oldest polygons values, removing duplicates index values
self.index[i - 1][1] = upd
# final_results = np.ones(len(self.points))*self.bck # Multiplying by the background value
for item in self.index:
for i in item[1]:
self.final_results[i] = item[0]
# Now save the results
if self.model_name:
print("Saving file...")
# The final output is a 1-column file with the values assigned,
with open(self.model_name, "w") as mm:
# in the same order as the blocks provided. The file is opened and saved every time a modification is
# made.
mm.write(str(int(self.Npts)) + "\n")
[mm.write(str(fr) + "\n") for fr in self.final_results]
mm.close()
[docs] def disconnect(self): # When finished
"""
Disconnect the Polygon selector
:return:
"""
print("Job done")
self.poly.disconnect_events()