codexconditional-information-xor.py
conditional-information-xor.py
"""Draw the joint XOR law and the dependence revealed by conditioning.
Uses only matplotlib and numpy from the root pyproject.toml; saves PNG to cwd.
"""
from pathlib import Path
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch, FancyBboxPatch
import numpy as np
def draw_joint(ax, matrix, title, subtitle):
ax.imshow(matrix, cmap='Blues', vmin=0, vmax=0.5)
for x in range(2):
for y in range(2):
probability = matrix[x, y]
label = {0: '0', 0.25: '1/4', 0.5: '1/2'}[probability]
ax.text(y, x, label, ha='center', va='center', fontsize=19,
color='white' if probability > 0.3 else '#213547')
ax.set(xticks=[0, 1], yticks=[0, 1], xlabel='Y', ylabel='X')
ax.set_title(title, fontsize=12)
ax.set_xticks(np.arange(-0.5, 2, 1), minor=True)
ax.set_yticks(np.arange(-0.5, 2, 1), minor=True)
ax.grid(which='minor', color='white', linewidth=2)
ax.tick_params(which='minor', bottom=False, left=False)
ax.text(0.5, -0.29, subtitle, ha='center', va='top', transform=ax.transAxes,
fontsize=11)
def main():
joint = np.zeros((2, 2, 2))
states = []
for x in (0, 1):
for y in (0, 1):
z = x ^ y
joint[x, y, z] = 0.25
states.append([str(x), str(y), str(z), '1/4'])
marginal = joint.sum(axis=2)
conditional = [joint[:, :, z] / joint[:, :, z].sum() for z in (0, 1)]
assert np.allclose(marginal, 0.25)
assert all(np.allclose(m.sum(axis=0), [0.5, 0.5]) for m in conditional)
assert all(np.allclose(m.sum(axis=1), [0.5, 0.5]) for m in conditional)
plt.rcParams.update({'font.size': 12})
fig = plt.figure(figsize=(10, 7.5), layout='constrained', facecolor='white')
grid = fig.add_gridspec(2, 6, height_ratios=[1, 1.28])
graph = fig.add_subplot(grid[0, :3])
graph.set(xlim=(-1.5, 1.5), ylim=(-0.2, 1.7))
graph.axis('off')
graph.set_title('Independent fair X, Y; Z is their XOR', fontsize=12, pad=12)
positions = {'X': (-0.9, 1.05), 'Y': (0.9, 1.05), 'Z': (0, 0.10)}
for name, (x, y) in positions.items():
graph.add_patch(FancyBboxPatch((x - 0.19, y - 0.16), 0.38, 0.32,
boxstyle='round,pad=0.035',
facecolor='#edf2f7', edgecolor='#41576b'))
graph.text(x, y, name, ha='center', va='center', fontsize=16, weight='bold')
for x in (-0.9, 0.9):
graph.add_patch(FancyArrowPatch((x * 0.83, 0.88), (x * 0.18, 0.31),
arrowstyle='-|>', mutation_scale=18,
linewidth=2, color='#425e79'))
graph.text(0, 1.42, 'No dependence between the inputs', ha='center', fontsize=10)
graph.text(0, 0.59, r'$Z=X\oplus Y$', ha='center', fontsize=16)
table_ax = fig.add_subplot(grid[0, 3:])
table_ax.axis('off')
table_ax.set_title('All four possible joint states', fontsize=12, pad=12)
table = table_ax.table(cellText=states, colLabels=['X', 'Y', 'Z', 'P(X,Y,Z)'],
cellLoc='center', bbox=[0.03, 0.02, 0.94, 0.90],
colWidths=[0.15, 0.15, 0.15, 0.35])
table.auto_set_font_size(False)
table.set_fontsize(13)
for (row, col), cell in table.get_celld().items():
cell.set_edgecolor('#c2cbd4')
cell.set_facecolor('#e7eef5' if row == 0 else 'white')
if row == 0:
cell.set_text_props(weight='bold')
matrices = [marginal, *conditional]
titles = ['Before observing Z', 'Given Z = 0\nP(Z = 0) = 1/2',
'Given Z = 1\nP(Z = 1) = 1/2']
subtitles = ['P(X,Y)\nIndependent: I(X;Y) = 0',
'P(X,Y | Z = 0)\nY = X: one bit of dependence',
'P(X,Y | Z = 1)\nY = 1 - X: one bit of dependence']
for index, (matrix, title, subtitle) in enumerate(zip(matrices, titles, subtitles)):
ax = fig.add_subplot(grid[1, 2*index:2*index+2])
draw_joint(ax, matrix, title, subtitle)
fig.suptitle('Conditioning can reveal dependence: I(X;Y | Z) = 1 bit',
fontsize=16, weight='bold')
fig.savefig(Path.cwd() / 'conditional-information-xor.png', dpi=120,
facecolor='white', transparent=False)
plt.close(fig)
if __name__ == '__main__':
main()