Add plot_warping_function(), enforce Lagrangian Multiplier constraint

src/sectionproperties/analysis/fea.py

@@ -25,28 +25,33 @@
 def _assemble_torsion(
     k_el: np.ndarray,
     f_el: np.ndarray,
+    c_el: np.ndarray,
+    n: np.ndarray,
     b: np.ndarray,
     weight: float,
     nx: float,
     ny: float,
-) -> tuple[np.ndarray, np.ndarray]:
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
     """Utility function for calculating the torsion stiffness matrix and load vector.
 
     Args:
         k_el: Element stiffness matrix
         f_el: Element load vector
+        c_el: Element constraint vector
+        n: Shape function
         b: Strain matrix
         weight: Effective weight
         nx: Global x-coordinate
         ny: Global y-coordinate
 
     Returns:
-        Torsion stiffness matrix and load vector (``k_el``, ``f_el``)
+        Torsion stiffness matrix and load vector (``k_el``, ``f_el``, ``c_el``)
     """
     # calculated modulus weighted stiffness matrix and load vector
     k_el += weight * b.transpose() @ b
     f_el += weight * b.transpose() @ np.array([ny, -nx])
-    return k_el, f_el
+    c_el += weight * n
+    return k_el, f_el, c_el
 
 
 @njit(cache=True, nogil=True)
@@ -272,30 +277,32 @@ def geometric_properties(
             self.material.density,
         )
 
-    def torsion_properties(self) -> tuple[np.ndarray, np.ndarray]:
+    def torsion_properties(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
         """Calculates the element warping stiffness matrix and the torsion load vector.
 
         Returns:
-            Element stiffness matrix ``k_el`` and element torsion load vector ``f_el``
+            Element stiffness matrix ``k_el``, element torsion load vector ``f_el``
+            and element constraint vector ``c_el``
         """
         # initialise stiffness matrix and load vector
         k_el = np.zeros(shape=(6, 6), dtype=float)
         f_el = np.zeros(shape=6, dtype=float)
+        c_el = np.zeros(shape=6, dtype=float)
 
         # Gauss points for 4 point Gaussian integration
         gps = gauss_points(n=4)
 
         for gp in gps:
             # determine shape function, shape function derivative,
             # jacobian and global coordinates
-            _, b, j, nx, ny = shape_function(coords=self.coords, gauss_point=gp)
+            n, b, j, nx, ny = shape_function(coords=self.coords, gauss_point=gp)
 
             weight = gp[0] * j * self.material.elastic_modulus
 
             # calculated modulus weighted stiffness matrix and load vector
-            k_el, f_el = _assemble_torsion(k_el, f_el, b, weight, nx, ny)
+            k_el, f_el, c_el = _assemble_torsion(k_el, f_el, c_el, n, b, weight, nx, ny)
 
-        return k_el, f_el
+        return k_el, f_el, c_el
 
     def shear_load_vectors(
         self,

src/sectionproperties/analysis/section.py

@@ -1411,17 +1411,19 @@ def assemble_torsion(
         col = []  # list holding column indices
         data = []  # list holding stiffness matrix entries
         f_torsion = np.zeros(n_size)  # force vector array
+        c_torsion = np.zeros(n_size)  # constraint vector array
 
         # loop through all elements in the mesh
         for el in self.elements:
             # determine number of nodes in the current element
             n = len(el.node_ids)
 
             # calculate the element stiffness matrix and torsion load vector
-            k_el, f_el = el.torsion_properties()
+            k_el, f_el, c_el = el.torsion_properties()
 
             # assemble the torsion load vector
             f_torsion[el.node_ids] += f_el
+            c_torsion[el.node_ids] += c_el
 
             for node_id in el.node_ids:
                 row.extend([node_id] * n)
@@ -1432,15 +1434,15 @@ def assemble_torsion(
                 progress.update(task_id=task, advance=1)
 
         # construct Lagrangian multiplier matrix:
-        # column vector of ones
+        # column vector
         row.extend(range(n_size))
         col.extend([n_size] * n_size)
-        data.extend([1] * n_size)
+        data.extend(c_torsion)
 
-        # row vector of ones
+        # row vector
         row.extend([n_size] * n_size)
         col.extend(range(n_size))
-        data.extend([1] * n_size)
+        data.extend(c_torsion)
 
         k_lg = coo_matrix(
             (data, (row, col)), shape=(n_size + 1, n_size + 1), dtype=float
@@ -1673,6 +1675,61 @@ def plot_centroids(
 
         return ax
 
+    def plot_warping_function(
+        self,
+        title: str = "Warping Function",
+        level: int = 20,
+        cmap: str = "viridis",
+        with_lines: bool = True,
+        **kwargs,
+    ):
+        r"""Plots the warping function over the mesh.
+
+        Args:
+            title: Plot title
+            level: Number of contour levels
+            cmap: Colormap
+            with_lines: If set to True, contour lines are displayed
+            kwargs: Passed to :func:`~sectionproperties.post.post.plotting_context`
+
+        Raises:
+            RuntimeError: If a warping analysis has not been performed
+
+        Returns:
+            Matplotlib axes object
+        """
+        if self.section_props.omega is None:
+            raise RuntimeError(
+                "Perform a warping analysis before plotting the warping function."
+            )
+
+        # create plot and setup the plot
+        with post.plotting_context(title=title, **kwargs) as (fig, ax):
+            assert ax
+
+            loc = self.mesh["vertices"]
+
+            if with_lines:
+                ax.tricontour(
+                    loc[:, 0],
+                    loc[:, 1],
+                    self.section_props.omega,
+                    colors="k",
+                    levels=level,
+                )
+
+            ax.tricontourf(
+                loc[:, 0],
+                loc[:, 1],
+                self.section_props.omega,
+                cmap=cmap,
+                levels=level,
+            )
+            ax.set_xlabel("X")
+            ax.set_ylabel("Y")
+
+        return ax
+
     def display_mesh_info(self) -> None:
         """Prints mesh statistics to the command line."""
         console = Console()