Draw figures for an opposed-roll mechanic

clash.py

#!/usr/bin/python3

# January 2024
# Ian McDougall

# Plotting some outcomes of opposed rolls from 1d6 vs 1d6 to 5d6*5 vs 5d6*5

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from dyce import H
from fractions import Fraction

nfig = 0
d6 = H(6)

### Plotting Functions (re-used from Rod Reel & Fist script)

# from https://matplotlib.org/stable/gallery/images_contours_and_fields/image_annotated_heatmap.html
def heatmap(data, row_labels, col_labels, ax=None, cbar_kw={}, cbarlabel="", **kwargs):
	if not ax:
		ax = plt.gca()
	im = ax.imshow(data, **kwargs)
	ax.set_xticks(np.arange(data.shape[1]))
	ax.set_yticks(np.arange(data.shape[0]))
	ax.set_xticklabels(col_labels)
	ax.set_yticklabels(row_labels)
	ax.tick_params(length=0)
	for edge, spine in ax.spines.items():
		spine.set_visible(False)
	return im

# ibid.
def annotate_heatmap(im, data=None, valfmt="{x:.2f}", textcolors=("black", "white"), threshold=None, **textkw):
	if not isinstance(data, (list, np.ndarray)):
		data = im.get_array()
	if threshold is not None:
		threshold = im.norm(threshold)
	else:
		threshold = im.norm(data.max())/2
	kw = dict(horizontalalignment="center", verticalalignment="center")
	kw.update(textkw)
	if isinstance(valfmt, str):
		valfmt = mpl.ticker.StrMethodFormatter(valfmt)
	texts = []
	for i in range(data.shape[0]):
		for j in range(data.shape[1]):
			kw.update(color=textcolors[int(im.norm(data[i, j]) >= threshold)])
			text = im.axes.text(j, i, valfmt(data[i, j], None), **kw)
			texts.append(text)
	return texts


### Plot nd6 vs 1d6 over n=[1,5]

n = range(1,6)
wins = []
ties = []

for i in range(len(n)):
	roll = (n[i]@d6).vs(d6)
	wins.append(Fraction(roll[1], roll.total))
	ties.append(Fraction(roll[0], roll.total))

nlabels = ["n={}".format(x) for x in n]

nfig += 1
fig, ax = plt.subplots()
ax.bar(nlabels, wins, label="Wins")
ax.bar(nlabels, ties, label="Ties", bottom=wins)
ax.set_ylim(0,1)
ax.legend()
fig.suptitle("nd6 vs. 1d6")
fig.tight_layout()
fig.savefig('figure{:02d}.png'.format(nfig))

### Plot 1d6*m vs 1d6 over m=[1,5]

m = range(1,6)
wins = []
ties = []

for i in range(len(m)):
	roll = (d6*m[i]).vs(d6)
	wins.append(Fraction(roll[1], roll.total))
	ties.append(Fraction(roll[0], roll.total))

mlabels = ["m={}".format(x) for x in m]

nfig += 1
fig, ax = plt.subplots()
ax.bar(mlabels, wins, label="Wins")
ax.bar(mlabels, ties, label="Ties", bottom=wins)
ax.set_ylim(0,1)
ax.legend()
fig.suptitle("1d6*m vs. 1d6")
fig.tight_layout()
fig.savefig('figure{:02d}.png'.format(nfig))

### Heatmap of nd6 vs. 1d6*m

ndice = [n[x]@d6 for x in range(len(n))]
mdice = [d6*m[x] for x in range(len(m))]
matches = [[ndice[x].vs(mdice[y]) for x in range(len(n))] for y in range(len(m))]

results = [[np.divide(max(matches[y][x].values()),matches[y][x].total) for x in range(len(n))] for y in range(len(m))]	# not sure I have m and n right here, but they're the same for now

colors = [[sum(np.multiply(list(matches[y][x].values()),list(matches[y][x].keys())))/matches[y][x].total for x in range(len(n))] for y in range(len(m))]	# same caveat

nfig += 1
fig, ax = plt.subplots()
im = ax.imshow(colors, norm=mpl.colors.Normalize(-1,1), cmap="PiYG")
ax.set_xticks(np.arange(len(nlabels)))
ax.set_yticks(np.arange(len(mlabels)))
ax.set_xticklabels(nlabels)
ax.set_yticklabels(mlabels)
ax.xaxis.tick_top()
ax.tick_params(length=0)
for edge, spine in ax.spines.items():
	spine.set_visible(False)
texts = annotate_heatmap(im, data=np.array(results), threshold=-0.5, textcolors=("white", "black"))
# Abandoning for now the idea that I can make the text colors do what I want.
fig.suptitle("nd6 vs. 1d6*m")
fig.tight_layout()
fig.savefig('figure{:02d}.png'.format(nfig))