Lots of type hints

trussme/evaluate.py

@@ -1,5 +1,7 @@
 from typing import TypedDict
+
 import numpy
+from numpy.typing import NDArray
 
 
 TrussInfo = TypedDict("TrussInfo", {
@@ -12,13 +14,11 @@
 })
 
 
-def the_forces(truss_info: TrussInfo) -> tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray, float]:
+def the_forces(truss_info: TrussInfo) -> tuple[NDArray[float], NDArray[float], NDArray[float], float]:
     tj: numpy.ndarray = numpy.zeros([3, numpy.size(truss_info["connections"], axis=1)])
-    w: numpy.ndarray = numpy.array([
-        numpy.size(truss_info["reactions"], axis=0),
-        numpy.size(truss_info["reactions"], axis=1)
-    ])
-    dof: numpy.ndarray = numpy.zeros([3*w[1], 3*w[1]])
+    w: numpy.ndarray = numpy.array(
+        [numpy.size(truss_info["reactions"], axis=0), numpy.size(truss_info["reactions"], axis=1)])
+    dof: numpy.ndarray = numpy.zeros([3 * w[1], 3 * w[1]])
     deflections: numpy.ndarray = numpy.ones(w)
     deflections -= truss_info["reactions"]
 
@@ -28,19 +28,15 @@ def the_forces(truss_info: TrussInfo) -> tuple[numpy.ndarray, numpy.ndarray, num
     # Build the global stiffness matrix
     for i in range(numpy.size(truss_info["connections"], axis=1)):
         ends = truss_info["connections"][:, i]
-        length_vector = truss_info["coordinates"][:, ends[1]] \
-            - truss_info["coordinates"][:, ends[0]]
+        length_vector = truss_info["coordinates"][:, ends[1]] - truss_info["coordinates"][:, ends[0]]
         length = numpy.linalg.norm(length_vector)
-        direction = length_vector/length
+        direction = length_vector / length
         d2 = numpy.outer(direction, direction)
-        ea_over_l = truss_info["elastic_modulus"][i]*truss_info["area"][i]\
-            / length
-        ss = ea_over_l*numpy.concatenate((numpy.concatenate((d2, -d2), axis=1),
-                                          numpy.concatenate((-d2, d2), axis=1)),
-                                         axis=0)
-        tj[:, i] = ea_over_l*direction
-        e = list(range((3*ends[0]), (3*ends[0] + 3))) \
-            + list(range((3*ends[1]), (3*ends[1] + 3)))
+        ea_over_l = truss_info["elastic_modulus"][i] * truss_info["area"][i] / length
+        ss = ea_over_l * numpy.concatenate((numpy.concatenate((d2, -d2), axis=1), numpy.concatenate((-d2, d2), axis=1)),
+                                           axis=0)
+        tj[:, i] = ea_over_l * direction
+        e = list(range((3 * ends[0]), (3 * ends[0] + 3))) + list(range((3 * ends[1]), (3 * ends[1] + 3)))
         for ii in range(6):
             for j in range(6):
                 dof[e[ii], e[j]] += ss[ii, j]
@@ -56,15 +52,13 @@ def the_forces(truss_info: TrussInfo) -> tuple[numpy.ndarray, numpy.ndarray, num
     ff = numpy.where(deflections.T == 1)
     for i in range(len(ff[0])):
         deflections[ff[1][i], ff[0][i]] = flat_deflections[i]
-    forces = numpy.sum(numpy.multiply(
-        tj, deflections[:, truss_info["connections"][1, :]]
-        - deflections[:, truss_info["connections"][0, :]]), axis=0)
+    forces = numpy.sum(numpy.multiply(tj,
+        deflections[:, truss_info["connections"][1, :]] - deflections[:, truss_info["connections"][0, :]]), axis=0)
 
     # Check the condition number, and warn the user if it is out of range
     cond = numpy.linalg.cond(SSff)
 
     # Compute the reactions
-    reactions = numpy.sum(dof*deflections.T.flat[:], axis=1)\
-        .reshape([w[1], w[0]]).T
+    reactions = numpy.sum(dof * deflections.T.flat[:], axis=1).reshape([w[1], w[0]]).T
 
     return forces, deflections, reactions, cond

trussme/joint.py

@@ -1,9 +1,10 @@
 import numpy
-
+from numpy.typing import NDArray
+from typing import Literal
 
 class Joint(object):
 
-    def __init__(self, coordinates: numpy.ndarray):
+    def __init__(self, coordinates: NDArray[float]):
         # Save the joint id
         self.idx = -1
 
@@ -25,7 +26,7 @@ def __init__(self, coordinates: numpy.ndarray):
         # Loads
         self.deflections = numpy.zeros([3, 1])
 
-    def free(self, d:int = 3):
+    def free(self, d: int = 3):
         self.translation = numpy.zeros([3, 1])
         # If 2d, add out of plane support
         if d is 2:
@@ -35,8 +36,8 @@ def pinned(self, d: int = 3):
         # Restrict all translation
         self.translation = numpy.ones([3, 1])
 
-    def roller(self, axis: str = 'y', d: int = 3):
-        # Only support reaction along denotated axis
+    def roller(self, axis: Literal["x", "y"] = 'y', d: int = 3):
+        # Only support reaction along denoted axis
         self.translation = numpy.zeros([3, 1])
         self.translation[ord(axis)-120] = 1
 

trussme/member.py

@@ -2,44 +2,45 @@
 import typing
 import warnings
 from trussme.physical_properties import Material, MATERIALS
+from trussme.joint import Joint
 
 
 class Member(object):
 
     # Shape types
     shapes: list[str] = ["pipe", "bar", "square", "box"]
 
-    def __init__(self, joint_a: int, joint_b: int):
+    def __init__(self, joint_a: Joint, joint_b: Joint):
         # Save id number
         self.idx = -1
 
         # Shape independent variables
-        self.shape = ''
-        self.t = 0.0  # thickness
-        self.w = 0.0  # outer width
-        self.h = 0.0  # outer height
-        self.r = 0.0  # outer radius
+        self.shape: str = ''
+        self.t: float = 0.0  # thickness
+        self.w: float = 0.0  # outer width
+        self.h: float = 0.0  # outer height
+        self.r: float = 0.0  # outer radius
 
         # Material properties
-        self.material = ''  # string specifying material
-        self.elastic_modulus = 0.0        # Elastic modulus
-        self.Fy = 0.0       # yield strength
-        self.rho = 0.0      # material density
+        self.material: str = ''  # string specifying material
+        self.elastic_modulus: float = 0.0        # Elastic modulus
+        self.Fy: float = 0.0       # yield strength
+        self.rho: float = 0.0      # material density
 
         # Dependent variables
-        self.area = 0.0   # Cross-sectional area
-        self.I = 0.0   # Moment of inertia
-        self.LW = 0.0  # Linear weight
+        self.area: float = 0.0   # Cross-sectional area
+        self.I: float = 0.0   # Moment of inertia
+        self.LW: float = 0.0  # Linear weight
 
         # Variables to store information about truss state
-        self.force = 0
-        self.fos_yielding = 0
-        self.fos_buckling = 0
-        self.mass = 0
+        self.force: float = 0
+        self.fos_yielding: float = 0
+        self.fos_buckling: float = 0
+        self.mass: float = 0
 
         # Variable to store location in truss
         self.joints = [joint_a, joint_b]
-        self.length = 0
+        self.length: float = 0.0
         self.end_a = []
         self.end_b = []