🔀 Merge pull request #332 from bldrkamal/add-stiffness-matrix-feature

Add stiffness matrix feature to SystemElements class

anastruct/fem/system.py

@@ -2292,6 +2292,31 @@ def apply_load_case(self, loadcase: LoadCase) -> None:
 
             exec(f"self.{method}({kwargs})")  # pylint: disable=exec-used
 
+    def get_stiffness_matrix(self, element_id: int) -> Optional[np.ndarray]:
+        """
+        Return the stiffness matrix for a specific element by its ID.
+        Args:
+            element_id (int): ID of the element.
+
+        Returns:
+        Optional[Union[list, None]]: The stiffness matrix of the element if it exists, otherwise None.
+        """
+        try:
+            element = self.element_map[element_id]
+        except KeyError:
+            print(f"Element with ID {element_id} does not exist.")
+            return None
+
+        if isinstance(element, Element):
+            if hasattr(element, "stiffness_matrix"):
+                print(f"Stiffness Matrix for Element ID {element_id}:")
+                return element.stiffness_matrix
+
+            print(f"Element ID {element_id} does not have a stiffness matrix.")
+        else:
+            print(f"Invalid element type for element ID {element_id}.")
+        return None
+
     def __deepcopy__(self, _: str) -> "SystemElements":
         """Deepcopy the SystemElements object.
 

tests/test_stiffness.py

@@ -0,0 +1,35 @@
+import unittest
+
+from anastruct.fem import system as se
+
+
+class TestSystemElements(unittest.TestCase):
+    def test_get_stiffness_matrix(self):
+        # Create a system
+        system = se.SystemElements()
+
+        # Add a simple beam element (make sure it has a stiffness matrix)
+        system.add_element(
+            location=[[0, 0], [5, 0]]
+        )  # Adjust as necessary based on your implementation
+
+        # Here, we assume that the added element has a stiffness matrix defined
+        # Check if stiffness matrix is correctly returned for element 1
+        stiffness_matrix = system.get_stiffness_matrix(1)
+
+        # Assert that a stiffness matrix was returned
+        self.assertIsNotNone(stiffness_matrix)
+        print("Stiffness matrix returned successfully.")
+
+    def test_invalid_element(self):
+        # Create a system without adding any elements
+        system = se.SystemElements()
+
+        # Try to get the stiffness matrix of an element that doesn't exist
+        stiffness_matrix = system.get_stiffness_matrix(999)  # Nonexistent ID
+        self.assertIsNone(stiffness_matrix)
+        print("Handled invalid element ID correctly.")
+
+
+if __name__ == "__main__":
+    unittest.main()