pivots = [(0, 0), (1, 2), (2, 0), (3, 2), (4, 0)]

N_piv = len(pivots)
print('treillis à {} pivots (dont 2 pivots de fixation)'.format(N_piv))

# Conversion en tablea NumPy pour bénéficier des facilités du slicing
pivots_arr = np.array(pivots)

# Tracés en boucle:
for piv in pivots:
    plt.plot(piv[0], piv[1], 'o')

# Ajustements des limites :
xlim(-.5, 4.5)
ylim(-.5, 2.5);

# indice des points d'accroche:
iP_A = 0 # 1er point du treillis
iP_B = 1 # 2ème point

P_A = pivots[iP_A]
P_B = pivots[iP_B]

barres = [
 ((0, 0), (2, 0)), # barre entre le pt (0,0) et le pt (2,0)
 ((1, 2), (2, 0)),
 ((1, 2), (3, 2)),
 ((2, 0), (3, 2)),
 ((2, 0), (4, 0)),
 ((3, 2), (4, 0))]

N_bar = len(barres)
print('treillis à {} barres'.format(N_bar))

color=(0.6, 0.6, 0.8)

# Tracé des barres, en boucle :
for j, bj in enumerate(barres):
    # Coordonnées des 2 pivots d'accroche:
    (x1,y1), (x2, y2) = bj
    color = (.5, .5, .8)
    plt.plot((x1, x2), (y1, y2), '-', color = color, lw=4, zorder=1)

# Tracés des pivots :
for piv in pivots:
    plt.plot(piv[0], piv[1], 'ro')

# Ajustements des limites :
xlim(-.5, 4.5)
ylim(-.5, 2.5);

F_ext = np.zeros((N_piv, 2))
F_ext.T # .T signifie "transposée". (Utilisé ici pour l'esthétique)

# Appui sur le dernier pivot:

# Solution 1)
F_ext[N_piv-1, 0] = 0  # composante horizontale
F_ext[N_piv-1, 1] = -1 # composante verticale

# Solution 2) 
F_ext[-1] = (0., -1) # effort vers le bas

F_ext.T 

Inc_mat = np.zeros((N_piv, N_bar), dtype=int)
Inc_mat

print(pivots)
pivots.index((1,2))
# pivots.index((9,9)) # la recherche d'un élement inexistant génère une ValueError

for j, bar_j in enumerate(barres):
    P1, P2 = bar_j
    # Remarquons la convention de signe:
    i1 = pivots.index(P1)
    Inc_mat[i1,j] = -1 # la barre j "quitte" P1
    i2 = pivots.index(P2)
    Inc_mat[i2,j] = +1 # la barre j "arrive à" P2

print('Matrice d\'incidence:')
print(str(Inc_mat).replace('0','.')) # Méthode d'affichage cosmétique

# Étape 1 : construire une liste d'indice contenant les indices de tous les pivots sauf P_A et P_B
piv_ind = range(N_piv)
piv_ind.remove(iP_A)
piv_ind.remove(iP_B)
piv_ind

# Étape 2 : utiliser la liste `piv_ind` pour sélectionner les lignes de la matrice d'incidence :
Inc_mat_red = Inc_mat[piv_ind, :]

print(str(Inc_mat_red).replace('0','.'))

# Conversion des barres en trableau NumPy
barres_arr = np.array(barres)
print(barres_arr.shape) # On obtient un tableau 3D
barres_arr

barres_OP1 = barres_arr[:,0,:]
barres_OP2 = barres_arr[:,1,:]

# À quoi correspond ce barres_OP1 ?
barres_OP1

barres_P1P2 = barres_OP2 - barres_OP1
barres_P1P2

# Étape 1 : calcul de la longueur des barres :
barre_l = np.sqrt((barres_P1P2**2).sum(axis = 1))
barre_l

# Étape 2 : Transformation en vecteur colonne : 
barre_l = barre_l.reshape(-1,1)
barre_l

# Étape 3 : normalisation vectorisée
barre_dir = barres_P1P2 / barre_l

barre_dir

# 1) Matrice A
Ax = Inc_mat_red*barre_dir[:,0]
Ay = Inc_mat_red*barre_dir[:,1]
# ou bien: Ax = np.dot(Inc_mat_red, np.diag(barre_dir[:,0]))

Ax.round(2)

A = np.vstack((Ax, Ay))
# Image de la matrice:
# plt.imshow(A, interpolation='nearest')

print(A.round(2))

# 2) Vecteur b : force extérieure appliquée à chaque pivot:
# Empilement des composantes selon x et y:
b_ext = np.concatenate((F_ext[piv_ind,0], 
                        F_ext[piv_ind,1]))
b_ext

# On assemble A et b (on "borde" la matrice)
Ab = np.hstack((A,b_ext.reshape(-1,1))) # attention au reshape
print(np.round(Ab,2))

N = A.shape[0]

for k in range(N):
    # Trouver le pivot
    piv_candidats = Ab[k:,k]
    k_piv = np.argmax(np.abs(piv_candidats)) + k
    piv = Ab[k_piv,k]
    if piv == 0:
        print(np.round(Ab,2))
        raise ValueError('Systeme non inversible')
    
    # échange des lignes k et k_piv:
    Ab[[k,k_piv], :] = Ab[[k_piv,k], :]
    
    # Normalisation de la ligne k
    Ab[k,k:] /= piv
    # Mise à zéros de la colonne k
    for i in range(N):
        if i == k:
            continue
        Ab[i,:] -= Ab[k,:]*Ab[i,k]

print(np.round(Ab,2))

# Extraction du vecteur (colonne N+1) :
trac_barres_gauss = Ab[:,N]
trac_barres_gauss.round(2)

np.linalg. # taper TAB

# Sol 1 : Résolution avec linsolve
trac_barres = np.linalg.solve(A, b_ext)

# Sol2 : Résolution x = inv(A)*b
A_inv = np.linalg.inv(A)

trac_barres = np.dot(A_inv, b_ext)

print('Effort de traction sur chaque barre:')
print(trac_barres.round(2))

# Recherche de l'effort maximal
trac_max = np.max(np.abs(trac_barres))
print(' -> effort max (en val. absolue) : {:.1f}'.format(trac_max))

# Efforts sur les points d'accroche (de la part du treillis et de l'extérieur):
resul_A = -np.inner(Inc_mat[iP_A,:]*trac_barres, barre_dir.T) + F_ext[iP_A]
resul_B = -np.inner(Inc_mat[iP_B,:]*trac_barres, barre_dir.T) + F_ext[iP_B]

print('action du treillis sur pt A: {!s} ({:.2f})'.format(
      resul_A, np.linalg.norm(resul_A,2)) )
print('action du treillis sur pt B: {!s} ({:.2f})'.format(
      resul_B, np.linalg.norm(resul_B,2)) )

F_ext.sum(axis=0)

cmap = plt.cm.gray
#cmap = plt.cm.winter
#cmap = plt.cm.jet # couleurs par défaut, comme dans Matlab
plt.imshow(A, interpolation='none', cmap=cmap)
plt.title('matrice A');

col = plt.cm.jet(0.9)
col # On reconnait du rouge
#plot(0,0, 'D', color=col)


# Colormap pour colorer les barres selon l'effort subit: rouge-bleu
col_list = [(0.9,0,0.0), (0.7,0.7,0.7), (0,0,0.9)]
cmap = mpl.colors.LinearSegmentedColormap.from_list('red-blue', col_list)

# Tracé des barres et des efforts
for j, bj in enumerate(barres):
    # Coordonnées des 2 pivots d'accroche:
    (x1,y1), (x2, y2) = bj
    # direction :
    uj = barre_dir[j]
    # effort de traction sur la barre bj:
    trac = trac_barres[j]
    color = cmap(trac/(2*trac_max)+0.5)
    plt.plot((x1, x2), (y1, y2), '-', color = color, lw=4, zorder=1)


# Tracés des pivots :
for piv in pivots:
    plt.plot(piv[0], piv[1], 'wo')

# Ajustements des limites :
xlim(-.5, 4.5)
ylim(-.5, 2.5);

color=(0.8, 0.8, 0.8)
# Échelle pour le tracé des forces
F_scale = 0.3*barre_l.mean()/trac_max
# Couleur des efforts:
F_color = (1., 0, 0) # rouge
dx, dy  = 1,2

# Tracé des barres et des efforts
for j, bj in enumerate(barres):
    # Coordonnées des 2 pivots d'accroche:
    (x1,y1), (x2, y2) = bj
    # direction :
    uj = barre_dir[j]
    trac = trac_barres[j]
    
    plt.plot((x1, x2), (y1, y2), '-', color = color, lw=4, zorder=1)
    
    # Tracé des efforts barre -> pivot1 et pivot 2
    plt.arrow(x1, y1, +trac*uj[0]*F_scale, +trac*uj[1]*F_scale,
              zorder=2, head_width=0.05*dx, lw=0, width=0.02*dx, color=F_color)
    plt.arrow(x2, y2, -trac*uj[0]*F_scale, -trac*uj[1]*F_scale,
              zorder=2, head_width=0.05*dx, lw=0, width=0.02*dx, color=F_color)

# Tracés des pivots :
for piv in pivots:
    plt.plot(piv[0], piv[1], 'wo', zorder=3)

# Ajustements des limites :
xlim(-.5, 4.5)
ylim(-.5, 2.5)

title(u'Forces exercées par les barres sur les liasons')

fig = plt.figure('efforts treillis', figsize=(12,5))

ax = fig.add_subplot(111, title='efforts sur le treillis '
                                '({:d} pivots, {:d} barres)'.format(N_piv, N_bar))

# Échelle pour le tracé des forces
F_scale = 0.3*barre_l.mean()/trac_max
# Couleur des efforts:
F_color = (0,0.8,0) # vert
dx, dy  = 1,2

# Colormap pour colorer les barres selon l'effort subit: rouge-bleu
col_list = [(0.9,0,0.0), (0.7,0.7,0.7), (0,0,0.9)]
cm_rb = mpl.colors.LinearSegmentedColormap.from_list('red-blue', col_list)
cm = cm_rb
#cm = plt.cm.coolwarm_r

# Tracé des barres et des efforts
for j, bj in enumerate(barres):
    # Coordonnées des 2 pivots d'accroche:
    (x1,y1), (x2, y2) = bj
    # direction :
    uj = barre_dir[j]
    # effort de traction sur la barre bj:
    trac = trac_barres[j]
    color = cm(trac/(2*trac_max)+0.5)
    plt.plot((x1, x2), (y1, y2), '-', color = color, lw=4, zorder=1)
    
    # Tracé des efforts barre -> pivot1 et pivot 2
    plt.arrow(x1, y1, +trac*uj[0]*F_scale, +trac*uj[1]*F_scale,
              zorder=2, head_width=0.05*dx, lw=0, width=0.02*dx, color=F_color)
    plt.arrow(x2, y2, -trac*uj[0]*F_scale, -trac*uj[1]*F_scale,
              zorder=2, head_width=0.05*dx, lw=0, width=0.02*dx, color=F_color)
# end for each barre

# Couleur des pivots
piv_color = (1.,1.,1.) # blanc
piv_color_AB = (1.,1.,0.5) # jaune clair
piv_alpha = 1 # opaque


# Tracé des pivots
for i, Pi in enumerate(pivots):
    # Tracé du pivot:
    marker = 'D' if Pi in (P_A,P_B) else 'o' # marqueur Diamond 'D' ou disque 'o'
    color = piv_color_AB if Pi in (P_A,P_B) else piv_color
    plt.plot(Pi[0], Pi[1], marker, ms=8, c=color, alpha = piv_alpha, zorder=3)
    # Force extérieur s'appliquant sur le pivot
    Fi = F_ext[i]
    if Fi.any():
        plt.arrow(Pi[0], Pi[1], Fi[0]*F_scale, Fi[1]*F_scale,
                  zorder=2, head_width=0.05*dx, lw=0, width=0.02*dx, color=(1,0,0))
    if Pi in (P_A,P_B):
        F_soutien = -resul_A if Pi==P_A else -resul_B
        plt.arrow(Pi[0], Pi[1], F_soutien[0]*F_scale, F_soutien[1]*F_scale,
                  zorder=2, head_width=0.05*dx, lw=0, width=0.02*dx, color=(1,1,0))
# end for each pivot

# Limites du tracé
plt.xlim(min([x for (x,y) in pivots]) - dx*1,
         max([x for (x,y) in pivots]) + dx*1)
plt.ylim(min([y for (x,y) in pivots]) - dy*.3,
         max([y for (x,y) in pivots]) + dy*.3)

# Couleur de fond:
ax.patch.set_fc((0.9,)*3)
ax.set_aspect('equal')
plt.grid(False)

# "Maximisation" du graphique au sein de la figure :
fig.tight_layout()