codexmomentum-collision-before-after.py
momentum-collision-before-after.py
"""Compare stipulated isolated one-dimensional collision outcomes.
Tested with Python 3.14.4, NumPy 2.3.5, Matplotlib 3.10.7.
Write momentum-collision-before-after.png to the current working directory.
"""
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
import numpy as np
def main():
masses = np.array([2.0, 1.0])
velocities = np.array([[3.0, 0.0], [1.0, 4.0], [2.0, 2.0]])
momenta = velocities * masses
kinetic = 0.5 * masses * velocities**2
internal_gain = kinetic[0].sum() - kinetic.sum(axis=1)
np.testing.assert_allclose(momenta.sum(axis=1), [6.0, 6.0, 6.0])
np.testing.assert_allclose(kinetic.sum(axis=1), [9.0, 9.0, 6.0])
np.testing.assert_allclose(internal_gain, [0.0, 0.0, 3.0])
print("Momentum totals (kg m/s):", momenta.sum(axis=1).tolist())
print("Kinetic totals (J):", kinetic.sum(axis=1).tolist())
print("Internal energy gains (J):", internal_gain.tolist())
plt.rcParams.update({"font.size": 12, "savefig.facecolor": "white"})
blue, orange, grey = "#276b9b", "#cf7323", "#e2e5e9"
colors = [blue, orange]
labels = ["Before", "Elastic\nafter", "Sticking\nafter"]
positions = np.arange(3)
fig, axes = plt.subplots(1, 3, figsize=(11.5, 6.0), facecolor="white")
fig.subplots_adjust(left=0.065, right=0.985, bottom=0.27,
top=0.73, wspace=0.38)
fig.suptitle("Same total momentum, different collision outcomes",
fontsize=18, y=0.965)
fig.text(0.5, 0.87,
"Isolated pair: mass 1 = 2 kg, mass 2 = 1 kg; rightward is positive",
ha="center", fontsize=12)
fig.text(0.5, 0.82,
"Incoming velocities: 3 m/s and 0 m/s. These are stipulated ideal examples.",
ha="center", fontsize=11)
width = 0.32
for i, color in enumerate(colors):
bars = axes[0].bar(positions + (i - 0.5) * width,
velocities[:, i], width, color=color)
axes[0].bar_label(bars, labels=[f"{v:g}" for v in velocities[:, i]],
padding=4, fontsize=12)
axes[0].set(title="Velocities (m/s)", ylim=(0, 4.9), yticks=[0, 1, 2, 3, 4])
for ax, values in zip(axes[1:], [momenta, kinetic], strict=True):
bottoms = np.zeros(3)
for i, color in enumerate(colors):
ax.bar(positions, values[:, i], 0.62, bottom=bottoms, color=color)
for x, bottom, value in zip(positions, bottoms, values[:, i], strict=True):
if value > 0:
ax.text(x, bottom + value / 2, f"{value:g}", ha="center",
va="center", color="white", fontsize=13, weight="bold")
bottoms += values[:, i]
axes[1].set(title="Momentum (kg m/s)", ylim=(0, 7.8), yticks=[0, 2, 4, 6])
for x in positions:
axes[1].text(x, 6.35, "Total 6", ha="center", fontsize=11)
axes[2].bar(positions, internal_gain, 0.62, bottom=kinetic.sum(axis=1),
color=grey, edgecolor="#69727e", hatch="///")
axes[2].text(2, 7.5, "3", ha="center", va="center", fontsize=13, weight="bold")
axes[2].set(title="Energy accounting (J)", ylim=(0, 11.4), yticks=[0, 3, 6, 9])
for x, energy in zip(positions, kinetic.sum(axis=1), strict=True):
axes[2].text(x, 9.5, f"K = {energy:g}", ha="center", fontsize=11)
for ax in axes:
ax.set_xticks(positions, labels)
ax.set_xlim(-0.65, 2.65)
ax.set_facecolor("white")
ax.set_axisbelow(True)
ax.grid(axis="y", alpha=0.17)
ax.spines[["top", "right"]].set_visible(False)
ax.tick_params(axis="x", length=0, pad=8)
ax.title.set_fontsize(13)
fig.legend(handles=[Patch(facecolor=blue, label="Body 1 (2 kg)"),
Patch(facecolor=orange, label="Body 2 (1 kg)"),
Patch(facecolor=grey, edgecolor="#69727e", hatch="///",
label="Gain in internal energy")],
loc="lower center", bbox_to_anchor=(0.5, 0.095), ncol=3,
frameon=False, fontsize=11)
fig.text(0.5, 0.04,
"Sticking transfers 3 J from kinetic to internal energy; the isolated pair keeps all its energy.",
ha="center", fontsize=11)
fig.savefig("momentum-collision-before-after.png", dpi=120,
facecolor="white", transparent=False)
plt.close(fig)
if __name__ == "__main__":
main()