杆+python有限元

import numpy as np
import pygame
import sys
from pygame.locals import *
from typing import List

# Define constants
EA = 1000
P = -10
WIDTH, HEIGHT = 640, 360
# Colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
BLUE = (0, 0, 255)


class Node:
"""
类表示桁架结构中的一个节点.
"""
def __init__(self, node_id: str, coordinate: List[float]):
self.id = node_id
self.coordinate = coordinate


class RodElement:
"""
类表示桁架结构中的一个杆元素.
"""
def __init__(self, element_id: str, node_ids: List[str]):
self.id = element_id
self.node_ids = node_ids
self.u_element = [0] * 4 # 杆元素的位移向量
self.node_coord = None

def length(self):
"""
计算杆元素的长度.
"""
coord1 = np.array(self.node_coord[0])[0]
coord2 = np.array(self.node_coord[1])[0]
return np.linalg.norm(coord2 - coord1)

def k_element(self):
"""
计算杆元素的刚度矩阵.
"""
dx = self.node_coord[1][0][0] - self.node_coord[0][0][0]
dy = self.node_coord[1][0][1] - self.node_coord[0][0][1]
L = self.length()
c = dx / L
s = dy / L
k = np.array([
[c**2, c*s, -c**2, -c*s],
[c*s, s**2, -c*s, -s**2],
[-c**2, -c*s, c**2, c*s],
[-c*s, -s**2, c*s, s**2]
]) * EA / L
return k


class Main:
"""
类表示整个桁架结构.
"""
def __init__(self):
self.nodes = {}
self.elements = {}
self.u = None

def create_structure(self):
"""
创建桁架结构，包括一组节点和相互连接它们的杆元素.
"""
num_nodes = int(input("请输入桁架结构的节点数: "))
distance = eval(input("请输入桁架结构相邻两个节点之间的距离: "))

# 创建节点
for i in range(num_nodes):
node_id = str(i)
coordinate = [i * distance, 0]
self.nodes[node_id] = Node(node_id, coordinate)

# 创建杆元素
element_id = 0
for i in range(num_nodes - 1):
node_id1 = str(i)
node_id2 = str(i + 1)
self.add_rod_element(str(element_id), [node_id1, node_id2])
element_id += 1

def add_rod_element(self, element_id: str, node_ids: List[str]):
"""
往桁架结构中添加一个杆元素.
"""
try:
node1 = self.nodes[node_ids[0]]
node2 = self.nodes[node_ids[1]]
except KeyError:
print(f"节点{node_ids[0]}或节点{node_ids[1]}不存在.")
return

rod_element = RodElement(element_id, node_ids)
rod_element.node_coord = np.mat([node1.coordinate, node2.coordinate])
self.elements[element_id] = rod_element

def calculate_force_and_stiffness(self):
"""
计算桁架结构的受力向量和刚度矩阵.
"""
num_nodes = len(self.nodes)
F = np.zeros((2 * num_nodes,))
K = np.zeros((2 * num_nodes, 2 * num_nodes))
for i, node in self.nodes.items():
if node.coordinate[1] == 0:
F[2 * int(i) + 1] = P # z轴方向上的受力

for element in self.elements.values():
k = element.k_element()
u_element = np.mat(element.u_element).T
f_element = k.dot(u_element)
node1_id, node2_id = element.node_ids
F[2 * int(node1_id):2 * int(node1_id) + 2] += f_element[:2].flatten()
F[2 * int(node2_id):2 * int(node2_id) + 2] += f_element[2:].flatten()
K[2 * int(node1_id):2 * int(node1_id) + 2, 2 * int(node1_id):2 * int(node1_id) + 2] += k[:2, :2]
K[2 * int(node1_id):2 * int(node1_id) + 2, 2 * int(node2_id):2 * int(node2_id) + 2] += k[:2, 2:]
K[2 * int(node2_id):2 * int(node2_id) + 2, 2 * int(node1_id):2 * int(node1_id) + 2] += k[2:, :2]
K[2 * int(node2_id):2 * int(node2_id) + 2, 2 * int(node2_id):2 * int(node2_id) + 2] += k[2:, 2:]

# 移除约束节点的受力和刚度贡献，设定位移,解方程组
constraint_nodes = [node for node in self.nodes.values() if node.coordinate[1] == 0]
num_constraints = len(constraint_nodes)
for i, node in enumerate(constraint_nodes):
j = 2 * int(node.id)
F = np.delete(F, j - i)
F = np.delete(F, j - i)
K = np.delete(K, j - i, 0)
K = np.delete(K, j - i, 1)

U = np.linalg.solve(K, F)
for i, node in self.nodes.items():
indices = [2 * int(i), 2 * int(i) + 1]
self.u = np.append(self.u, U[indices])
self.u = self.u.reshape((-1, 2))

def plot_structure(self):
"""
绘制桁架结构.
"""
pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Truss")
clock = pygame.time.Clock()

while True:
for event in pygame.event.get():
if event.type == QUIT:
pygame.quit()
sys.exit()

screen.fill(WHITE)
for element in self.elements.values():
node1 = self.nodes[element.node_ids[0]]
node2 = self.nodes[element.node_ids[1]]
pygame.draw.line(screen, BLACK, node1.coordinate, node2.coordinate, 2)

pygame.display.update()
clock.tick(60)


if __name__ == "__main__":

main = Main()
main.create_structure()
main.calculate_force_and_stiffness()
main.plot_structure()