# Simple combination of FEM elastic data for overlapping nodes using SciPy'c COO format
import numpy as np
from scipy.sparse import coo_array

# Lists use in to construct the COO format array
rows = []
cols = []
data = []

# --- ELEMENT 1 ---
elem1_nodes = [0, 1] # global node indices
elem1_stiffness = np.array([
    [ 100, -100],
    [-100,  100]
])
print("Element 1 stiffness matrix:\n", elem1_stiffness)

for i in range(2):
    for j in range(2):
        rows.append(elem1_nodes[i])
        cols.append(elem1_nodes[j])
        data.append(elem1_stiffness[i, j])

# --- ELEMENT 2 ---
elem2_nodes = [1, 2] # global node indices
elem2_stiffness = np.array([
    [ 200, -200],
    [-200,  200]
])
print("Element 2 stiffness matrix:\n", elem2_stiffness)

for i in range(2):
    for j in range(2):
        rows.append(elem2_nodes[i])
        cols.append(elem2_nodes[j])
        data.append(elem2_stiffness[i, j])

# --- STIFFNESS MATRIX ASSEMBLY USING COO ---
# Create the sparse matrix for 3 nodes (3x3 shape)
stiffness = coo_array((data, (rows, cols)), shape=(3, 3))

# Convert to a dense array to see the automatic summation at overlapping Node 1
print("\nStiffness Matrix:")
print(stiffness.toarray())