text: import numpy as np
  from scipy.linalg import lu_factor, lu_solve
  
  
  class Truss:
      """
      Simple static equilibrium solver for truss structures
  
      Args
      ----
      points : ndarray, shape (n_points, dim)
          Point coordinates with spatial dimension ``dim``
      lines : ndarray, shape (n_lines, 2)
          Connectivity of bars
      """
  
      def __init__(self, points, lines):
          self.points = np.asarray(points)
          self.lines = np.asarray(lines)
          assert self.lines.shape[1] == 2
  
          self.n_points, self.dim = self.points.shape
          self.n_lines = self.lines.shape[0]
          self.n_ddl = self.n_points * self.dim
  
          # Solving
          self.B = None  # gradient matrix
          self.K = None
          self.K_lu = None
          self.F = None
  
      def integrate(self, fun):
          """
          Integrate a function defined on bars
  
          Args
          ----
          fun : ndarray, shape (n_lines, ...)
              Function to be integrated
          """
          return np.sum(self.A * self.L * fun, axis=0)
  
      def solve(
          self, A=1, E=1, U_dict={}, F_dict={}, sig0=None, construct_K=False
      ):
          """
          Solve the static equilibrium problem for the truss structure
  
          Args
          ----
          A : float or ndarray, shape (n_lines, )
              Cross section area
          E : float
              Young's modulus
          U_dict : dict
              Prescribed displacement ``{point_id: (Ux, Uy), ...}``, use ``None`` for ``Ux`` or ``Uy`` when this component is not concerned
          F_dict : dict
              Prescribed nodal force ``{point_id: (Fx, Fy), ...}``
          sig0 : ndarray, shape (n_lines, )
              Initial stress
          construct_K : bool
              Whether force reconstructing and performing LU factorization of the stiffness matrix
  
          Returns
          -------
          u : ndarray, shape (n_ddl, )
              Displacement solution
          eps : ndarray, shape (n_lines, )
              Strain
          sig : ndarray, shape (n_lines, )
              Stress
          """
          # Geometric and material properties
          try:
              assert len(A) == self.n_lines
              self.A = A
          except TypeError:
              self.A = A * np.ones(self.n_lines)
          except AssertionError:
              raise RuntimeError("Length of A must match the number of bars")
          self.E = E
  
          # Stiffness matrix
          if self.K is None or construct_K:
              self._compute_elementary_quantities()
              self._construct_K()
  
              # Apply Dirichlet conditions to K
              if len(U_dict) > 0:
                  self.K_Dirichlet = self.K.copy()
                  self._apply_Dirichlet(U_dict, K=self.K_Dirichlet)
              else:
                  self.K_Dirichlet = self.K
  
          # LU factorization
          if self.K_lu is None or construct_K:
              self.K_lu, self.K_piv = lu_factor(self.K_Dirichlet, check_finite=False)
  
          # Right hand side
          self._construct_F(F_dict, sig0)
          self._apply_Dirichlet(U_dict, F=self.F)
  
          # Solve and post-processing
          u = lu_solve((self.K_lu, self.K_piv), self.F, check_finite=False)
          eps = self._strain(u)
          return u, eps, self.E * eps
  
      def _global_ddl_indices(self, point_id):
          return np.array([2 * point_id, 2 * point_id + 1], dtype=int)
  
      def _compute_elementary_quantities(self):
          if self.B is None:
              self.L = np.zeros(self.n_lines)
              self.B = np.zeros((self.n_lines, 4))
              for i, line in enumerate(self.lines):
                  coords = self.points[line]
                  x, y = coords[:, 0], coords[:, 1]
                  self.L[i] = np.linalg.norm(coords[1] - coords[0])
                  cos = (x[1] - x[0]) / self.L[i]
                  sin = (y[1] - y[0]) / self.L[i]
                  self.B[i] = np.array([-cos, -sin, cos, sin]) / self.L[i]
  
      def _construct_K(self):
          self.K = np.zeros((self.n_ddl, self.n_ddl))
          for i, line in enumerate(self.lines):
              global_indices = np.hstack(
                  [self._global_ddl_indices(line[0]), self._global_ddl_indices(line[1])]
              )
              Ke = self.A[i] * self.E * np.outer(self.B[i], self.B[i]) * self.L[i]
              self.K[global_indices[:, np.newaxis], global_indices] += Ke
  
      def _apply_Dirichlet(self, U_dict={}, K=None, F=None):
          for point_id, U in U_dict.items():
              assert len(U) == 2
              idx = [i for i in range(len(U)) if U[i] is not None]
              if len(idx) == 0:
                  continue
              U = np.asarray(U)[idx]
              global_indices = self._global_ddl_indices(point_id)[idx]
              if K is not None:
                  K[global_indices, :] = 0
                  K[global_indices[:, np.newaxis], global_indices] = np.eye(len(idx))
              if F is not None:
                  F[global_indices] = U
  
      def _construct_F(self, F_dict={}, sig0=None):
          self.F = np.zeros(self.n_ddl)
  
          # Prescribed nodal forces
          for point_id, F in F_dict.items():
              assert len(F) == 2
              global_indices = self._global_ddl_indices(point_id)
              self.F[global_indices] += F
  
          # Initial stress
          if sig0 is not None:
              assert len(sig0) == self.n_lines
              self._compute_elementary_quantities()
              for i, line in enumerate(self.lines):
                  global_indices = np.hstack(
                      [
                          self._global_ddl_indices(line[0]),
                          self._global_ddl_indices(line[1]),
                      ]
                  )
                  Fe = -self.A[i] * sig0[i] * self.B[i] * self.L[i]
                  self.F[global_indices] += Fe
  
      def _strain(self, u):
          eps = np.zeros(self.n_lines)
          for i, line in enumerate(self.lines):
              global_indices = np.hstack(
                  [self._global_ddl_indices(line[0]), self._global_ddl_indices(line[1])]
              )
              u_ = u[global_indices]
              eps[i] = self.B[i] @ u_
          return eps
  
      def plot(self, ax=None, u=None, eps=None, points_id=False, lines_id=False):
          import matplotlib.pyplot as plt
  
          if ax is None:
              _, ax = plt.subplots()
  
          if points_id:
              for i, point in enumerate(self.points):
                  ax.text(point[0], point[1], f"{i:d}")
  
          if eps is not None:
              eps_absmax = np.abs(eps).max()
              norm = plt.cm.colors.Normalize(vmin=-eps_absmax, vmax=eps_absmax)
              colors = plt.cm.coolwarm(norm(eps))
  
          for i, line in enumerate(self.lines):
              coords = self.points[line]
              style = "--"
              color = "gray"
              if eps is not None and u is None:
                  style = "-"
                  color = colors[i]
              ax.plot(coords[:, 0], coords[:, 1], style, color=color)
              if lines_id:
                  centroid = np.mean(coords, axis=0)
                  ax.text(centroid[0], centroid[1], f"({i:d})")
  
          if u is not None:
              deformed = self.points + u.reshape((-1, 2))
              for i, line in enumerate(self.lines):
                  coords = deformed[line]
                  color = "C0"
                  if eps is not None and u is not None:
                      color = colors[i]
                  ax.plot(coords[:, 0], coords[:, 1], color=color)
  