From 44159426dfa9924b24f0ddaa1ec1a77f67e1656f Mon Sep 17 00:00:00 2001
From: Chris McComb <ccmcc2012@gmail.com>
Date: Wed, 3 Jan 2024 15:04:48 -0500
Subject: [PATCH] New plotting options for truss

---
 tests/test_construction_and_io.py |  6 +++-
 trussme/report.py                 |  9 +++++-
 trussme/visualize.py              | 46 +++++++++++++++++++++++++++----
 3 files changed, 54 insertions(+), 7 deletions(-)

diff --git a/tests/test_construction_and_io.py b/tests/test_construction_and_io.py
index 70877de..115aeeb 100644
--- a/tests/test_construction_and_io.py
+++ b/tests/test_construction_and_io.py
@@ -21,7 +21,11 @@ def test_demo_report(self):
         # Build truss from file
         truss_from_file = trussme.read_trs(TEST_TRUSS_FILENAME)
 
-        trussme.report_to_md("asdf.md", truss_from_file, trussme.Goals())
+        trussme.report_to_md(
+            os.path.join(os.path.dirname(__file__), "asdf.md"),
+            truss_from_file,
+            trussme.Goals(),
+        )
 
     def test_build_methods(self):
         goals = trussme.Goals(
diff --git a/trussme/report.py b/trussme/report.py
index 36cfb37..59fbe71 100644
--- a/trussme/report.py
+++ b/trussme/report.py
@@ -394,6 +394,10 @@ def __generate_stress_analysis(truss, goals, with_figures: bool = True) -> str:
 
     # Print information about members
     analysis += "\n## FORCES AND STRESSES\n"
+
+    if with_figures:
+        analysis += trussme.visualize.plot_truss(truss, starting_shape="force") + "\n"
+
     data = []
     rows = []
     for m in truss.members:
@@ -430,7 +434,10 @@ def __generate_stress_analysis(truss, goals, with_figures: bool = True) -> str:
     analysis += "\n## DEFLECTIONS\n"
 
     if with_figures:
-        analysis += trussme.visualize.plot_truss(truss, deflected_shape=True) + "\n"
+        analysis += (
+            trussme.visualize.plot_truss(truss, starting_shape="k", deflected_shape="m")
+            + "\n"
+        )
 
     data = []
     rows = []
diff --git a/trussme/visualize.py b/trussme/visualize.py
index ef3cedc..987bac6 100644
--- a/trussme/visualize.py
+++ b/trussme/visualize.py
@@ -1,10 +1,20 @@
 import matplotlib.pyplot
 import io
 import re
+import numpy
+from typing import Union, Literal
 
 
 def plot_truss(
-    truss, deflected_shape: bool = False, exaggeration_factor: float = 10
+    truss,
+    starting_shape: Union[
+        None, Literal["fos", "force", "b", "g", "r", "c", "m", "y", "k"]
+    ] = "k",
+    deflected_shape: Union[
+        None, Literal["fos", "force", "b", "g", "r", "c", "m", "y", "k"]
+    ] = None,
+    exaggeration_factor: float = 10,
+    fos_threshold: float = 1.0,
 ) -> str:
     """Plot the truss.
 
@@ -12,10 +22,18 @@ def plot_truss(
     ----------
     truss: Truss
         The truss to plot.
-    deflected_shape: bool, default=False
-        Whether to plot the deflected shape.
+    starting_shape: : Union[None, Literal["fos", "force", "b", "g", "r", "c", "m", "y", "k"]], default="k"
+        How to show the starting shape. If None, the starting shape is not shown. If "fos", the members are colored
+        green if the factor of safety is above the threshold and red if it is below. If "force", the members are colored
+        according to the force in the member. If a color, the members are colored that color.
+    deflected_shape: : Union[None, Literal["fos", "force", "b", "g", "r", "c", "m", "y", "k"]], default = None
+        How to show the deflected shape. If None, the starting shape is not shown. If "fos", the members are colored
+        green if the factor of safety is above the threshold and red if it is below. If "force", the members are colored
+        according to the force in the member. If a color, the members are colored that color.
     exaggeration_factor: float, default=10
         The factor by which to exaggerate the deflected shape.
+    fos_threshold: float, default=1.0
+        The threshold for the factor of safety. If the factor of safety is below this value, the member is colored red.
 
     Returns
     -------
@@ -30,15 +48,33 @@ def plot_truss(
     ax.axis("equal")
     ax.set_axis_off()
 
+    scaler: float = numpy.max(numpy.abs([member.force for member in truss.members]))
+
     for member in truss.members:
+        if starting_shape == "fos":
+            color = "g" if member.fos > fos_threshold else "r"
+        elif starting_shape == "force":
+            color = matplotlib.pyplot.cm.bwr(member.force / (2 * scaler) + 0.5)
+        elif starting_shape is None:
+            break
+        else:
+            color = starting_shape
         ax.plot(
             [member.begin_joint.coordinates[0], member.end_joint.coordinates[0]],
             [member.begin_joint.coordinates[1], member.end_joint.coordinates[1]],
-            color="k",
+            color=color,
         )
 
     if deflected_shape:
         for member in truss.members:
+            if starting_shape == "fos":
+                color = "g" if member.fos > fos_threshold else "r"
+            elif starting_shape == "force":
+                color = matplotlib.pyplot.cm.bwr(member.force / (2 * scaler) + 0.5)
+            elif starting_shape is None:
+                break
+            else:
+                color = starting_shape
             ax.plot(
                 [
                     member.begin_joint.coordinates[0]
@@ -52,7 +88,7 @@ def plot_truss(
                     member.end_joint.coordinates[1]
                     + exaggeration_factor * member.end_joint.deflections[1],
                 ],
-                color="m",
+                color=color,
             )
 
     imgdata = io.StringIO()