## Jupyter Snippet CB2nd 04_energy

Jupyter Snippet CB2nd 04_energy

# 9.4. Finding the equilibrium state of a physical system by minimizing its potential energy

``````import numpy as np
import scipy.optimize as opt
import matplotlib.pyplot as plt
%matplotlib inline
``````
``````g = 9.81  # gravity of Earth
m = .1  # mass, in kg
n = 20  # number of masses
e = .1  # initial distance between the masses
l = e  # relaxed length of the springs
k = 10000  # spring stiffness
``````
``````P0 = np.zeros((n, 2))
P0[:, 0] = np.repeat(e * np.arange(n // 2), 2)
P0[:, 1] = np.tile((0, -e), n // 2)
``````
``````A = np.eye(n, n, 1) + np.eye(n, n, 2)
# We display a graphic representation of
# the matrix.
f, ax = plt.subplots(1, 1)
ax.imshow(A)
ax.set_axis_off()
``````

``````L = l * (np.eye(n, n, 1) + np.eye(n, n, 2))
for i in range(n // 2 - 1):
L[2 * i + 1, 2 * i + 2] *= np.sqrt(2)
``````
``````I, J = np.nonzero(A)
``````
``````def dist(P):
return np.sqrt((P[:,0]-P[:,0][:,np.newaxis])**2 +
(P[:,1]-P[:,1][:,np.newaxis])**2)
``````
``````def show_bar(P):
fig, ax = plt.subplots(1, 1, figsize=(5, 4))

# Wall.
ax.axvline(0, color='k', lw=3)

# Distance matrix.
D = dist(P)

# Get normalized elongation in [-1, 1].
elong = np.array([D[i, j] - L[i, j]
for i, j in zip(I, J)])
elong_max = np.abs(elong).max()

# The color depends on the spring tension, which
# is proportional to the spring elongation.
colors = np.zeros((len(elong), 4))
colors[:, -1] = 1  # alpha channel is 1

# Use two different sequentials colormaps for
# positive and negative elongations, to show
# compression and extension in different colors.
if elong_max > 1e-10:
# We don't use colors if all elongations are
# zero.
elong /= elong_max
pos, neg = elong > 0, elong < 0
colors[pos] = plt.cm.copper(elong[pos])
colors[neg] = plt.cm.bone(-elong[neg])

# We plot the springs.
for i, j, c in zip(I, J, colors):
ax.plot(P[[i, j], 0],
P[[i, j], 1],
lw=2,
color=c,
)

# We plot the masses.
ax.plot(P[[I, J], 0], P[[I, J], 1], 'ok',)

# We configure the axes.
ax.axis('equal')
ax.set_xlim(P[:, 0].min() - e / 2,
P[:, 0].max() + e / 2)
ax.set_ylim(P[:, 1].min() - e / 2,
P[:, 1].max() + e / 2)
ax.set_axis_off()

return ax
``````
``````ax = show_bar(P0)
ax.set_title("Initial configuration")
``````

``````def energy(P):
# The argument P is a vector (flattened matrix).
# We convert it to a matrix here.
P = P.reshape((-1, 2))
# We compute the distance matrix.
D = dist(P)
# The potential energy is the sum of the
# gravitational and elastic potential energies.
return (g * m * P[:, 1].sum() +
.5 * (k * A * (D - L)**2).sum())
``````
``````energy(P0.ravel())
``````
``````-0.981
``````
``````bounds = np.c_[P0[:2, :].ravel(),
P0[:2, :].ravel()].tolist() + \
[[None, None]] * (2 * (n - 2))
``````
``````P1 = opt.minimize(energy, P0.ravel(),
method='L-BFGS-B',
bounds=bounds).x.reshape((-1, 2))
``````
``````ax = show_bar(P1)
ax.set_title("Equilibrium configuration")
``````