Problem: Axial Bar Stiffness and Displacement Analysis

The problem involves creating a Python workflow for axial bar structural analysis: compute stable element stiffness tensors, assemble a global matrix from two elements via Kahan summation, and solve for displacements with input validation and singularity checks to predict structural responses accurately.

Background: ### Industry Background & Technical Context for Structural Finite Element Analysis (FEA) Axial Bar Systems

#### 1. Structural Mechanics Fundamentals: Axial Bar Stiffness
In structural engineering, axial bar elements are core components of truss, frame, and pipeline systems. The classical axial stiffness matrix for a single bar element connecting two nodes is derived from Hooke's Law and equilibrium principles:
\[
[k] = \frac{EA}{L} \begin{bmatrix} 1 & -1 \\ -1 & 1 \end{bmatrix}
\]
where \(E\) = Young's modulus (elastic stiffness), \(A\) = cross-sectional area, \(L\) = bar length. For viscoelastic materials (e.g., polymer composites, damping treatments), \(E\) is extended to a complex-valued modulus (\(E = E' + iE''\)) to model frequency-dependent damping and energy dissipation.

#### 2. Batch Processing & Tensor Support
Modern FEA workflows require handling multiple elements simultaneously (batch processing). Tensor inputs (multi-dimensional arrays) allow vectorized computation of stiffness matrices for large datasets, reducing computational overhead. The `element_stiffness` function supports broadcast-compatible input shapes to enable efficient batch processing of multiple bar elements.

#### 3. Numerical Stability Enhancements
- **Log-Transform Stiffness Calculation**: Direct computation \(k = EA/L\) can cause numerical overflow/underflow for extreme values (e.g., very high \(E\) or very small \(L\)). Using log-transformation (\(\log(k) = \log(E) + \log(A) - \log(L)\)) followed by exponentiation avoids these issues. Iterative refinement (2 iterations) corrects for small arithmetic errors introduced by the log-exp transformation.
- **Kahan Summation**: When assembling global matrices, floating-point accumulation errors can degrade precision. Kahan summation (a compensated summation algorithm) minimizes these errors by tracking and correcting residuals, ensuring accurate accumulation of overlapping matrix entries.
- **Row Pivoting**: For linear system solving, row swapping (when \(|a21| > |a11|\)) improves matrix conditioning, reducing amplification of numerical errors during determinant calculation and solution.
- **Singularity & Condition Number Checks**: A near-zero determinant indicates a singular/unstable structure (e.g., unconstrained system). The condition number (\(\sigma_{max}/\sigma_{min}\), where \(\sigma\) are singular values) quantifies solution sensitivity to input perturbations; values >1e10 trigger warnings for potential instability.

#### 4. Global Matrix Assembly
For two connected axial bars (nodes 0-1 and 1-2), the global stiffness matrix is 3x3. Element matrices are mapped to global node indices:
- \(k1\) (element 0-1) contributes to global positions (0,0)-(1,1)
- \(k2\) (element1-2) contributes to global positions (1,1)-(2,2)
Overlapping entries (e.g., (1,1)) are summed using Kahan summation for precision. Complex values are converted to floats if their imaginary part is negligible (to simplify output for cases with minimal damping).

#### 5. Linear System Solution for Displacements
Assuming node 0 is fixed (displacement \(u_0=0\)), we solve the reduced system for \(u_2\) and \(u_3\) (displacements at nodes1 and2). Using Cramer's rule for the 2x2 submatrix extracted from the global matrix:
\[
u_2 = \frac{F_2a_{22} - F_3a_{12}}{\det}, \quad u_3 = \frac{a_{11}F_3 - a_{21}F_2}{\det}
\]
where \(\det = a_{11}a_{22} - a_{12}a_{21}\), and \(F_2, F_3\) are applied forces at nodes1 and2.

import numpy as np
from typing import List, Tuple, Union
from numpy.typing import ArrayLike
import math
from scipy.sparse import issparse, spmatrix
import warnings

Input/Output: The problem takes as inputs: E1 (elastic modulus of first bar, Union[float, ArrayLike] with positive real part >1e-12), A1 (cross-sectional area of first bar, Union[float, ArrayLike] with positive real part >1e-12), L1 (length of first bar, Union[float, ArrayLike] with positive real part >1e-12), E2, A2, L2 (same properties for second bar), F2 (applied force at node 2, Union[int, float]), and F3 (applied force at node 3, Union[int, float]). It generates as outputs: two 2x2 complex-valued element stiffness tensors (k1, k2) computed from E, A, L via log-transformed calculation of (E*A/L) with iterative refinement; a 3x3 global stiffness matrix K assembled by adding k1 to positions (0,0)-(1,1) and k2 to (1,1)-(2,2) using Kahan summation (converted to a list with real elements where imaginary parts are negligible); and a tuple of floats (u2, u3) representing displacements at nodes 2 and 3, solved by extracting the 2x2 submatrix from K at indices (1,1)-(2,2), swapping rows if needed to avoid small pivots, checking for singular matrices via determinant (raising error if near-zero), computing the condition number (warning for values >1e10), and solving the linear system using Cramer's rule.

Required dependencies:
import numpy as np
from typing import List, Tuple, Union
from numpy.typing import ArrayLike
import math
from scipy.sparse import issparse, spmatrix
import warnings

Sub-steps:
step_number: 1
step_description_prompt: Please write a Python function named element_stiffness. It computes axial bar element stiffness tensor. Steps: convert inputs E, A, L to complex128 arrays, validate shape compatibility, ensure positive real parts (>1e-12), calculate k using log transform plus iterative refinement, form 2x2 matrix tensor and return it.
function_header: def element_stiffness(E, A, L):
test_cases: test_stiffness = [
    (200, 0.01, 2, [[1.0, -1.0], [-1.0, 1.0]]),   # k = 1.0
    (100, 1.0, 10, [[10.0, -10.0], [-10.0, 10.0]]), # k = 10.0
    (3000, 0.5, 15, [[100.0, -100.0], [-100.0, 100.0]]) # k = 100.0
]
for E, A, L, expected in test_stiffness:
    res = element_stiffness(E, A, L)
    # Handle complex values - compare real parts only
    for i in range(2):
        for j in range(2):
            r_val = complex(res[i][j]).real if isinstance(res[i][j], (complex, np.complex128)) else res[i][j]
            e_val = complex(expected[i][j]).real if isinstance(expected[i][j], (complex, np.complex128)) else expected[i][j]
            assert math.isclose(r_val, e_val, rel_tol=1e-3) or abs(r_val - e_val) < 1e-6

step_number: 2
step_description_prompt: Please write Python code to generate a function named assemble_global_matrix that takes k1 and k2 as inputs. First validate k1 and k2 are 2x2 numeric matrices. Convert them to complex arrays, then build a 3x3 matrix K by adding k1 to positions (0,0)-(1,1) and k2 to (1,1)-(2,2) using Kahan summation. Convert result to list, replacing complex numbers with zero imaginary parts to floats.
function_header: def assemble_global_matrix(k1, k2):
test_cases: k1 = [[1.0, -1.0], [-1.0, 1.0]]
k2 = [[2.0, -2.0], [-2.0, 2.0]]
test_assembly = [
    (k1, k1, [[1.0, -1.0, 0.0], [-1.0, 2.0, -1.0], [0.0, -1.0, 1.0]]), 
    (k1, k2, [[1.0, -1.0, 0.0], [-1.0, 3.0, -2.0], [0.0, -2.0, 2.0]]), 
    (k1, [[0,0],[0,0]], [[1.0, -1.0, 0.0], [-1.0, 1.0, 0.0], [0.0, 0.0, 0.0]]) 
]
for mat1, mat2, expected in test_assembly:
    res = assemble_global_matrix(mat1, mat2)
    assert all(math.isclose(res[i][j], expected[i][j]) for i in range(3) for j in range(3))

step_number: 3
step_description_prompt: Please write a Python function named solve_displacements. It takes K_global (numpy array or scipy sparse matrix), F2, F3 as inputs. Check input validity, extract submatrix elements at indices (1,1), (1,2), (2,1), (2,2). Handle row swap if needed, compute determinant (raise error if singular), calculate condition number (warn if high), then return u2 and u3 as a tuple of floats.
function_header: def solve_displacements(K_global, F2, F3):
test_cases: test_solve = [
    ([[0,0,0], [0, 2, -1], [0, -1, 1]], 1.0, 0.0, (1.0, 1.0)), # 2u2-u3=1, -u2+u3=0 -> u2=1, u3=1
    ([[0,0,0], [0, 10, -5], [0, -5, 5]], 15.0, 0.0, (3.0, 3.0)), # 10u2-5u3=15, -5u2+5u3=0 -> u2=3, u3=3
    ([[0,0,0], [0, 4, 0], [0, 0, 2]], 8.0, 4.0, (2.0, 2.0))      # 4u2=8, 2u3=4 -> u2=2, u3=2
]
for K, f2, f3, expected in test_solve:
    u2, u3 = solve_displacements(K, f2, f3)
    # Handle tuple comparison
    if isinstance(expected, tuple):
        assert abs(u2 - expected[0]) < 1e-3 and abs(u3 - expected[1]) < 1e-3
    else:
        assert abs(u2 - expected) < 1e-3
    # Compare tuple results
    actual_val = (u2, u3)
    assert isinstance(expected, tuple) and len(expected) == 2
    assert abs(u2 - expected[0]) < 1e-3 and abs(u3 - expected[1]) < 1e-3

General solution:

General tests: