187 lines
7.6 KiB
Python
187 lines
7.6 KiB
Python
# FEM.py
|
||
|
||
import numpy as np
|
||
|
||
from Bar import Bar
|
||
from DiscreteMethod import DiscreteMethod
|
||
|
||
class FEM(DiscreteMethod):
|
||
def __init__(self, bar:'Bar', desired_order="2"):
|
||
super().__init__(bar, desired_order)
|
||
|
||
def get_kappa(self, row:int, col:int):
|
||
dx = self.dx
|
||
alpha = self.alpha
|
||
|
||
if (row == 1 and col == 1) or (row == 2 and col == 2):
|
||
return -1/dx + alpha**2*dx/ 3.0
|
||
elif row == 3 and col == 3:
|
||
return -1/(3*dx) + alpha**2*dx/30.0
|
||
elif (row == 1 and col == 2) or (row == 2 and col == 1):
|
||
return 1/dx + alpha**2*dx/6.0
|
||
elif (row == 1 and col == 3) or (row == 3 and col == 1) or (row == 2 and col == 3) or (row == 3 and col == 2):
|
||
return alpha**2*dx/12.0
|
||
|
||
def get_all_kappa(self):
|
||
'''Returns the kappa matrix for p=1 or p=2 FEM.
|
||
Note that both matrices are equivalent regardless of p=1 or p=2, but p=1 does not use row 3 or col 3'''
|
||
|
||
K = np.zeros((3, 3))
|
||
for row in range(3):
|
||
for col in range(3):
|
||
K[row, col] = self.get_kappa(row+1, col+1)
|
||
return K
|
||
|
||
|
||
def build_stiffness_matrix(self):
|
||
'''Given a number of sections, build the stiffness_matrix for P=1 FEM'''
|
||
num_sections = self.num_sections
|
||
if self.desired_order == "2": # p=1
|
||
size = num_sections + 1 # Number of nodes for P=1
|
||
self.stiffness_matrix = np.zeros((size, size))
|
||
|
||
# Get local stiffness matrix (K) for a single section
|
||
K_local = self.get_all_kappa()
|
||
|
||
# Extract local stiffness terms for P=1
|
||
k11 = K_local[0, 0]
|
||
k12 = K_local[0, 1]
|
||
k21 = K_local[1, 0]
|
||
k22 = K_local[1, 1]
|
||
|
||
# Loop through each section and assemble the global stiffness matrix
|
||
for section in range(num_sections):
|
||
self.stiffness_matrix[section, section - 1] = k21
|
||
self.stiffness_matrix[section, section] = k22 + k11
|
||
self.stiffness_matrix[section, section + 1] = k12
|
||
|
||
# Enforce boundary conditions by fixing the first and last nodes
|
||
self.stiffness_matrix[0, :] = 0 # Set first row to zero
|
||
self.stiffness_matrix[0, 0] = 1 # Fix the left boundary
|
||
|
||
self.stiffness_matrix[-1, :] = 0 # Set last row to zero
|
||
self.stiffness_matrix[-1, -1] = 1 # Fix the right boundary
|
||
|
||
elif self.desired_order == "4": # p=2
|
||
size = 2 * num_sections + 1 # Number of linear + quadratic nodes
|
||
self.stiffness_matrix = np.zeros((size, size))
|
||
|
||
# Get local stiffness matrix (K) for a single section
|
||
K_local = self.get_all_kappa()
|
||
|
||
# Extract local stiffness terms for P=2
|
||
k11 = K_local[0, 0]
|
||
k12 = K_local[0, 1]
|
||
k13 = K_local[0, 2]
|
||
k21 = K_local[1, 0]
|
||
k22 = K_local[1, 1]
|
||
k23 = K_local[1, 2]
|
||
k31 = K_local[2, 0]
|
||
k32 = K_local[2, 1]
|
||
k33 = K_local[2, 2]
|
||
|
||
# Loop through each section and assemble the global stiffness matrix for similar to p=1
|
||
for section in range(num_sections):
|
||
# Assemble contributions from K_local
|
||
self.stiffness_matrix[section, section - 1] = k21
|
||
self.stiffness_matrix[section, section] = k22 + k11
|
||
self.stiffness_matrix[section, section + 1] = k12
|
||
|
||
# Add the k23 and k13 terms in that correspond to Ei and Ei+1
|
||
self.stiffness_matrix[section, num_sections+section] = k23
|
||
self.stiffness_matrix[section, num_sections+section+1] = k13
|
||
|
||
|
||
# Add the quadratic element relations
|
||
for section in range(num_sections):
|
||
# Assemble contributions from K_local
|
||
self.stiffness_matrix[num_sections+section+1, section] = k31
|
||
self.stiffness_matrix[num_sections+section+1, section+1] = k32
|
||
self.stiffness_matrix[num_sections+section+1, num_sections+section+1] = k33
|
||
|
||
# Enforce boundary conditions by fixing the first and last nodes
|
||
self.stiffness_matrix[0, :] = 0 # Set first row to zero
|
||
self.stiffness_matrix[0, 0] = 1 # Fix the left boundary
|
||
|
||
self.stiffness_matrix[num_sections, :] = 0 # Set last temperature row to 0
|
||
self.stiffness_matrix[num_sections, num_sections] = 1 # Fix the right boundary
|
||
|
||
def build_constant_vector(self):
|
||
'''Given a number of sections, build the constant_vector where the
|
||
first and last elements come from boundary conditions'''
|
||
num_sections = self.num_sections
|
||
if self.desired_order == "2": # p=1
|
||
self.constant_vector = np.zeros(num_sections + 1) # Include endpoints
|
||
self.constant_vector[0] = self.U_0 # Left boundary
|
||
self.constant_vector[-1] = self.U_L # Right boundary
|
||
if self.desired_order == "4": # p=2
|
||
self.constant_vector = np.zeros(2*num_sections + 1) # Include endpoints
|
||
self.constant_vector[0] = self.U_0 # Left boundary
|
||
self.constant_vector[num_sections] = self.U_L # Right boundary
|
||
|
||
def run(self, forcing_freq:float, num_sections:int=4):
|
||
'''Runs FDM for the given forcing frequency and num_sections.
|
||
Stores displacment ampltitude data in self.disp_amp
|
||
Stores strain amplitude data in self.strain_amp'''
|
||
# Save the appropriate values
|
||
self.num_sections = num_sections
|
||
self.dx = 1 / num_sections
|
||
self.alpha = forcing_freq*(self.bar.density / self.bar.E)**0.5
|
||
self.forcing_freq = forcing_freq
|
||
|
||
self.build_stiffness_matrix()
|
||
self.build_constant_vector()
|
||
|
||
# Solve the system
|
||
if self.desired_order == "2": # p=1
|
||
self.disp_amp = np.linalg.solve(self.stiffness_matrix, self.constant_vector)
|
||
elif self.desired_order == "4": # p=3
|
||
self.disp_amp = np.linalg.solve(self.stiffness_matrix, self.constant_vector)
|
||
self.disp_amp = self.disp_amp[0:self.num_sections+1]
|
||
|
||
# Calculate all of the nodal strains
|
||
self.calc_all_nodal_strain()
|
||
|
||
# Save the x values for the nodal values
|
||
self.x_vals = np.linspace(0, 1, self.num_sections+1) # sections + endpoints
|
||
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import matplotlib.pyplot as plt
|
||
from common import freq_from_alpha
|
||
bar = Bar()
|
||
|
||
# Initialize FDM
|
||
fem = FEM(bar, desired_order="4")
|
||
|
||
# Define alpha values for testing
|
||
alpha_values = np.linspace(0, 20, 6)
|
||
|
||
# Set boundary conditions
|
||
fem.set_boundary_conditions(U_0=0, U_L=100)
|
||
|
||
# Plot displacement amplitudes for varying alpha values
|
||
plt.figure(figsize=(12, 6))
|
||
for alpha in [20]:
|
||
forcing_freq = freq_from_alpha(bar, alpha)
|
||
fem.run(forcing_freq, num_sections=100)
|
||
plt.plot(fem.x_vals, fem.disp_amp, label=f'α = {alpha:.2f}')
|
||
|
||
plt.title("FEM Displacement Amplitude vs x for Various α Values")
|
||
plt.xlabel("x (Normalized Position)")
|
||
plt.ylabel("Displacement Amplitude")
|
||
plt.legend()
|
||
plt.grid(True)
|
||
plt.show()
|
||
|
||
# Set boundary conditions and run the solver
|
||
fem.set_boundary_conditions(U_0=0, U_L=100)
|
||
freq = freq_from_alpha(bar, 4)
|
||
num_sections = 100
|
||
fem.run(freq, num_sections)
|
||
|
||
# Interpolate at a specific x_val
|
||
x_val = 0.31
|
||
interpolated_strain = fem.get_interpolated_strain_amp(x_val)
|
||
print(f"Interpolated Strain Amplitude at x = {x_val}: {interpolated_strain}") |