Truss solver

truss-solver.py

import sys
import re
import numpy as np

trussfile = open(sys.argv[-1])

joints = {}  # keys are joint names, values are positions
members = {} # keys and values are joint names
memberlist = [] # items are two-char strings of member names
forces = {}  # keys are joint names, values are directions ('x', 'y', or 'z')
forcelist = [] # items are two-char strings of force names
loads = {}   # keys are joint names, values are vectors
dimensions = None # number of dimensions; either 2 or 3. detected automatically

mode_line_matcher = re.compile(r'([A-Z]+)_MODE')
# mode line:
# JOINT_MODE
joint_line_matcher = re.compile(r'\s*([A-Z]),\s*\((-?[0-9.]+),\s*(-?[0-9.]+)(,\s*(-?[0-9.]+))?\)')
# joint line:
# A, (0,1)
member_line_matcher = re.compile(r'\s*([A-Z])\s*([A-Z])$')
# member line:
# AB
force_line_matcher = re.compile(r'\s*([A-Z])\s*([xyz])$')
# force line:
# Ax

# load line (just use joint matcher):
# A, (100,200)

mode = "OUTSIDE"
for line in trussfile:
    line = line.rstrip()
    if not line or not mode:
        pass
    elif mode == "OUTSIDE":
        matches = mode_line_matcher.match(line)
        if not matches:
            pass
        else:
            mode = matches.group(1)
    elif mode == "JOINT":
        matches = joint_line_matcher.match(line)
        if not matches:
            if line == "END":
                mode = "OUTSIDE"
            else:
                pass
        else:
            if matches.group(4):
                dimensions = 3
                joints[matches.group(1)] = np.array([float(matches.group(2)),
                                                     float(matches.group(3)),
                                                     float(matches.group(5))])
            else:
                dimensions = 2
                joints[matches.group(1)] = np.array([float(matches.group(2)),
                                                     float(matches.group(3))])
            members[matches.group(1)] = []
            forces[matches.group(1)] = []
    elif mode == "MEMBER":
        matches = member_line_matcher.match(line)
        if not matches:
            if line == "END":
                mode = "OUTSIDE"
            else:
                pass
        else:
            members[matches.group(1)].append(matches.group(2))
            members[matches.group(2)].append(matches.group(1))
            memberlist.append(matches.group(1)+matches.group(2))
    elif mode == "FORCE":
        matches = force_line_matcher.match(line)
        if not matches:
            if line == "END":
                mode = "OUTSIDE"
            else:
                pass
        else:
            forces[matches.group(1)].append(matches.group(2))
            forcelist.append(matches.group(1)+matches.group(2))
    elif mode == "LOAD":
        matches = joint_line_matcher.match(line)
        if not matches:
            if line == "END":
                mode = "OUTSIDE"
            else:
                pass
        else:
            if matches.group(4):
                dimensions = 3
                loads[matches.group(1)] = np.array([float(matches.group(2)),
                                                    float(matches.group(3)),
                                                    float(matches.group(5))])
            else:
                dimensions = 2
                loads[matches.group(1)] = np.array([float(matches.group(2)),
                                                    float(matches.group(3))])

trussfile.close()

joint_names = list(joints.keys())
joint_names.sort()
memberlist.sort()
forcelist.sort()

## construct the member length list
member_length_list = []
for member in memberlist:
    member_length_list.append(
            np.linalg.norm(joints[member[0]]-joints[member[1]]))

## construct the system matrix

# system shape:
# members, forces

system = None
for joint_name in joint_names:
    # x equation
    eqn = np.array([])
    for index,member in enumerate(memberlist):
        coeff = 0
        if joint_name in member:
            other_joint = None
            if joint_name == member[0]:
                other_joint = member[1]
            else:
                other_joint = member[0]
            coeff = (joints[joint_name][0]-joints[other_joint][0])/member_length_list[index]
        eqn = np.concatenate((eqn,np.array([coeff])))
    for force in forcelist:
        coeff = 0
        if force == joint_name+"x":
            coeff = 1
        eqn = np.concatenate((eqn,np.array([coeff])))
    if system == None:
        system = np.array([eqn])
    else:
        system = np.concatenate((system,np.array([eqn])))
    # y equation
    eqn = np.array([])
    for index,member in enumerate(memberlist):
        coeff = 0
        if joint_name in member:
            other_joint = None
            if joint_name == member[0]:
                other_joint = member[1]
            else:
                other_joint = member[0]
            coeff = (joints[joint_name][1]-joints[other_joint][1])/member_length_list[index]
        eqn = np.concatenate((eqn,np.array([coeff])))
    for force in forcelist:
        coeff = 0
        if force == joint_name+"y":
            coeff = 1
        eqn = np.concatenate((eqn,np.array([coeff])))
    system = np.concatenate((system,np.array([eqn])))
    # z equation
    if dimensions == 3:
        eqn = np.array([])
        for index,member in enumerate(memberlist):
            coeff = 0
            if joint_name in member:
                other_joint = None
                if joint_name == member[0]:
                    other_joint = member[1]
                else:
                    other_joint = member[0]
                coeff = (joints[joint_name][2]-joints[other_joint][2])/member_length_list[index]
            eqn = np.concatenate((eqn,np.array([coeff])))
        for force in forcelist:
            coeff = 0
            if force == joint_name+"z":
                coeff = 1
            eqn = np.concatenate((eqn,np.array([coeff])))
        system = np.concatenate((system,np.array([eqn])))

## create the vector to solve for

sysvec = np.array([])
for joint_name in joint_names:
    # x eqn
    try:
        sysvec = np.concatenate((sysvec,np.array([-loads[joint_name][0]])))
    except KeyError:
        sysvec = np.concatenate((sysvec,np.array([0])))
    # y eqn
    try:
        sysvec = np.concatenate((sysvec,np.array([-loads[joint_name][1]])))
    except KeyError:
        sysvec = np.concatenate((sysvec,np.array([0])))
    # z eqn
    if dimensions == 3:
        try:
            sysvec = np.concatenate((sysvec,np.array([-loads[joint_name][2]])))
        except KeyError:
            sysvec = np.concatenate((sysvec,np.array([0])))

## solve and print out the solution
solution = list(np.linalg.lstsq(system,sysvec)[0])
labels = memberlist+forcelist
for index,num in enumerate(solution):
    print(labels[index],num)