Repository: ross-rotordynamics/ross
Pull request #809 (2021-06-25 12:24:29)

Fix disk scale factor and add test


Patch:
diff --git a/ross/disk_element.py b/ross/disk_element.py
index 1de657c8..89dd0baf 100644
--- a/ross/disk_element.py
+++ b/ross/disk_element.py
@@ -124,7 +124,7 @@ class DiskElement(Element):
             f"{self.__class__.__name__}"
             f"(Id={self.Id:{0}.{5}}, Ip={self.Ip:{0}.{5}}, "
             f"m={self.m:{0}.{5}}, color={self.color!r}, "
-            f"n={self.n}, tag={self.tag!r})"
+            f"n={self.n}, scale_factor={self.scale_factor}, tag={self.tag!r})"
         )
 
     def __str__(self):
@@ -477,7 +477,7 @@ class DiskElement(Element):
         >>> file_path = os.path.dirname(os.path.realpath(__file__)) + '/tests/data/shaft_si.xls'
         >>> list_of_disks = DiskElement.from_table(file_path, sheet_name="More")
         >>> list_of_disks[0]
-        DiskElement(Id=0.0, Ip=0.0, m=15.12, color='Firebrick', n=3, tag=None)
+        DiskElement(Id=0.0, Ip=0.0, m=15.12, color='Firebrick', n=3, scale_factor=1, tag=None)
         """
         parameters = read_table_file(file, "disk", sheet_name=sheet_name)
         if tag is None:
diff --git a/ross/rotor_assembly.py b/ross/rotor_assembly.py
index 5b9c3401..e46ab624 100644
--- a/ross/rotor_assembly.py
+++ b/ross/rotor_assembly.py
@@ -1948,8 +1948,20 @@ class Rotor(object):
 
         mean_od = np.mean(nodes_o_d)
         # plot disk elements
+
+        # calculate scale factor if disks have scale_factor='mass'
+        if all([disk.scale_factor == "mass" for disk in self.disk_elements]):
+            max_mass = max([disk.m for disk in self.disk_elements])
+            for disk in self.disk_elements:
+                f = disk.m / max_mass
+                disk._scale_factor_calculated = (1 - f) * 0.5 + f * 1.0
+
         for disk in self.disk_elements:
-            step = disk.scale_factor * mean_od
+            scale_factor = disk.scale_factor
+            if scale_factor == "mass":
+                scale_factor = disk._scale_factor_calculated
+            step = scale_factor * mean_od
+
             position = (nodes_pos[disk.n], nodes_o_d[disk.n] / 2, step)
             fig = disk._patch(position, fig)
 
@@ -3259,13 +3271,6 @@ class CoAxialRotor(Rotor):
             if disk.tag is None:
                 disk.tag = "Disk " + str(i)
 
-        # calculate scale factor if disks have scale_factor='mass'
-        if all([disk.scale_factor == "mass" for disk in disk_elements]):
-            max_mass = max([disk.m for disk in disk_elements])
-            for disk in disk_elements:
-                f = disk.m / max_mass
-                disk.scale_factor = (1 - f) * 0.5 + f * 1.0
-
         for i, brg in enumerate(bearing_elements):
             if brg.__class__.__name__ == "BearingElement" and brg.tag is None:
                 brg.tag = "Bearing " + str(i)


Test patch:
diff --git a/ross/tests/test_rotor_assembly.py b/ross/tests/test_rotor_assembly.py
index 12c7b3af..6d694415 100644
--- a/ross/tests/test_rotor_assembly.py
+++ b/ross/tests/test_rotor_assembly.py
@@ -1639,3 +1639,60 @@ def test_save_load(rotor8):
     rotor8_loaded = Rotor.load(file)
 
     rotor8 == rotor8_loaded
+
+
+def test_plot_rotor(rotor8):
+    fig = rotor8.plot_rotor()
+
+    for d in fig.data:
+        if d["name"] == "Disk 0":
+            actual_x = d["x"]
+            actual_y = d["y"]
+    expected_x = [
+        0.5,
+        0.5083333333333333,
+        0.49166666666666664,
+        0.5,
+        None,
+        0.5,
+        0.5083333333333333,
+        0.49166666666666664,
+        0.5,
+    ]
+    expected_y = [0.025, 0.125, 0.125, 0.025, None, -0.025, -0.125, -0.125, -0.025]
+    assert_allclose(actual_x[:4], expected_x[:4])
+    assert_allclose(actual_y[:4], expected_y[:4])
+
+    # mass scale factor
+    for disk in rotor8.disk_elements:
+        disk.scale_factor = "mass"
+
+    fig = rotor8.plot_rotor()
+    for d in fig.data:
+        if d["name"] == "Disk 0":
+            actual_x = d["x"]
+            actual_y = d["y"]
+    expected_x = [
+        0.5,
+        0.5068020833333333,
+        0.4931979166666667,
+        0.5,
+        None,
+        0.5,
+        0.5068020833333333,
+        0.4931979166666667,
+        0.5,
+    ]
+    expected_y = [
+        0.025,
+        0.106625,
+        0.106625,
+        0.025,
+        None,
+        -0.025,
+        -0.106625,
+        -0.106625,
+        -0.025,
+    ]
+    assert_allclose(actual_x[:4], expected_x[:4])
+    assert_allclose(actual_y[:4], expected_y[:4])