Repository: ross-rotordynamics/ross
Pull request #722 (2020-12-09 17:49:09)

Fix DoF error when using "n_link"
* Add `link_nodes` attribute to Rotor object that holds the nodes created using `n_link`
* Correct DoFs with `link_nodes` in `ForcedResponseResults` and `TimeResponseResults`
* Run Black


Patch:
diff --git a/ross/defects/abs_defect.py b/ross/defects/abs_defect.py
index 1d25c164..26fbf5f0 100644
--- a/ross/defects/abs_defect.py
+++ b/ross/defects/abs_defect.py
@@ -20,12 +20,10 @@ class Defect(ABC):
 
     def run_time_response(self):
         results = TimeResponseResults(
+            rotor=self,
             t=self.time_vector,
             yout=self.response.T,
             xout=[],
-            nodes_list=self.rotor.nodes,
-            nodes_pos=self.rotor.nodes_pos,
-            number_dof=self.rotor.number_dof,
         )
         return results
 
diff --git a/ross/results.py b/ross/results.py
index 4025040c..9ee5f899 100644
--- a/ross/results.py
+++ b/ross/results.py
@@ -149,6 +149,7 @@ class Results(ABC):
             if key == "rotor":
                 aux_file = str(file)[:-5] + "_rotor" + str(file)[-5:]
                 from ross.rotor_assembly import Rotor
+
                 data[key] = Rotor.load(aux_file)
 
             elif isinstance(value, Iterable):
@@ -1498,15 +1499,29 @@ class FrequencyResponseResults(Results):
             inp, out, frequency_units, amplitude_units, None, **mag_kwargs
         )
         fig1 = self.plot_phase(
-            inp, out, frequency_units, amplitude_units, phase_units, None, **phase_kwargs
+            inp,
+            out,
+            frequency_units,
+            amplitude_units,
+            phase_units,
+            None,
+            **phase_kwargs,
         )
         fig2 = self.plot_polar_bode(
-            inp, out, frequency_units, amplitude_units, phase_units, None, **polar_kwargs
+            inp,
+            out,
+            frequency_units,
+            amplitude_units,
+            phase_units,
+            None,
+            **polar_kwargs,
         )
 
         if fig is None:
             fig = make_subplots(
-                rows=2, cols=2, specs=[[{}, {"type": "polar", "rowspan": 2}], [{}, None]]
+                rows=2,
+                cols=2,
+                specs=[[{}, {"type": "polar", "rowspan": 2}], [{}, None]],
             )
 
         for data in fig0["data"]:
@@ -2023,7 +2038,9 @@ class ForcedResponseResults(Results):
             major_axis_vector[3, :] = axis vector response for the input angle
             major_axis_vector[4, :] = phase response for the input angle
         """
-        number_dof = self.rotor.number_dof
+        ndof = self.rotor.number_dof
+        nodes = self.rotor.nodes
+        link_nodes = self.rotor.link_nodes
 
         unit_type = str(Q_(1, amplitude_units).dimensionality)
         try:
@@ -2034,8 +2051,10 @@ class ForcedResponseResults(Results):
             )
 
         major_axis_vector = np.zeros((5, len(self.speed_range)), dtype=complex)
-        dofx = number_dof * node
-        dofy = number_dof * node + 1
+
+        fix_dof = (node - nodes[-1] - 1) * ndof // 2 if node in link_nodes else 0
+        dofx = ndof * node - fix_dof
+        dofy = ndof * node + 1 - fix_dof
 
         # Relative angle between probes (90°)
         Rel_ang = np.exp(1j * np.pi / 2)
@@ -2111,7 +2130,7 @@ class ForcedResponseResults(Results):
             major_axis_vector[4, :] = absolute values for major axes vectors
         """
         nodes = self.rotor.nodes
-        number_dof = self.rotor.number_dof
+        ndof = self.rotor.number_dof
 
         major_axis_vector = np.zeros((5, len(nodes)), dtype=complex)
         idx = np.where(np.isclose(self.speed_range, speed, atol=1e-6))[0][0]
@@ -2125,8 +2144,8 @@ class ForcedResponseResults(Results):
             )
 
         for i, n in enumerate(nodes):
-            dofx = number_dof * n
-            dofy = number_dof * n + 1
+            dofx = ndof * n
+            dofy = ndof * n + 1
 
             # Relative angle between probes (90°)
             Rel_ang = np.exp(1j * np.pi / 2)
@@ -3507,18 +3526,14 @@ class TimeResponseResults(Results):
 
     Parameters
     ----------
+    rotor : Rotor.object
+        The Rotor object
     t : array
         Time values for the output.
     yout : array
         System response.
     xout : array
         Time evolution of the state vector.
-    nodes_list : array
-        list with nodes from a rotor model.
-    nodes_pos : array
-        Rotor nodes axial positions.
-    number_dof : int
-        Number of degrees of freedom per shaft element's node
 
     Returns
     -------
@@ -3526,13 +3541,11 @@ class TimeResponseResults(Results):
         The figure object with the plot.
     """
 
-    def __init__(self, t, yout, xout, nodes_list, nodes_pos, number_dof):
+    def __init__(self, rotor, t, yout, xout):
         self.t = t
         self.yout = yout
         self.xout = xout
-        self.nodes_list = nodes_list
-        self.nodes_pos = nodes_pos
-        self.number_dof = number_dof
+        self.rotor = rotor
 
     def plot_1d(
         self,
@@ -3578,12 +3591,18 @@ class TimeResponseResults(Results):
         fig : Plotly graph_objects.Figure()
             The figure object with the plot.
         """
+        nodes = self.rotor.nodes
+        link_nodes = self.rotor.link_nodes
+        ndof = self.rotor.number_dof
+
         if fig is None:
             fig = go.Figure()
 
         for i, p in enumerate(probe):
-            dofx = p[0] * self.number_dof
-            dofy = p[0] * self.number_dof + 1
+            fix_dof = (p[0] - nodes[-1] - 1) * ndof // 2 if p[0] in link_nodes else 0
+            dofx = ndof * p[0] - fix_dof
+            dofy = ndof * p[0] + 1 - fix_dof
+
             angle = Q_(p[1], probe_units).to("rad").m
 
             # fmt: off
@@ -3648,19 +3667,23 @@ class TimeResponseResults(Results):
         fig : Plotly graph_objects.Figure()
             The figure object with the plot.
         """
+        nodes = self.rotor.nodes
+        link_nodes = self.rotor.link_nodes
+        ndof = self.rotor.number_dof
+
+        fix_dof = (node - nodes[-1] - 1) * ndof // 2 if node in link_nodes else 0
+        dofx = ndof * node - fix_dof
+        dofy = ndof * node + 1 - fix_dof
+
         if fig is None:
             fig = go.Figure()
 
         fig.add_trace(
             go.Scatter(
-                x=Q_(self.yout[:, self.number_dof * node], "m")
-                .to(displacement_units)
-                .m,
-                y=Q_(self.yout[:, self.number_dof * node + 1], "m")
-                .to(displacement_units)
-                .m,
+                x=Q_(self.yout[:, dofx], "m").to(displacement_units).m,
+                y=Q_(self.yout[:, dofy], "m").to(displacement_units).m,
                 mode="lines",
-                name=f"Orbit",
+                name="Orbit",
                 legendgroup="Orbit",
                 showlegend=False,
                 hovertemplate=(
@@ -3704,20 +3727,20 @@ class TimeResponseResults(Results):
         fig : Plotly graph_objects.Figure()
             The figure object with the plot.
         """
+        nodes_pos = self.rotor.nodes_pos
+        nodes = self.rotor.nodes
+        ndof = self.rotor.number_dof
+
         if fig is None:
             fig = go.Figure()
 
-        for n in self.nodes_list:
-            x_pos = np.ones(self.yout.shape[0]) * self.nodes_pos[n]
+        for n in nodes:
+            x_pos = np.ones(self.yout.shape[0]) * nodes_pos[n]
             fig.add_trace(
                 go.Scatter3d(
                     x=Q_(x_pos, "m").to(rotor_length_units).m,
-                    y=Q_(self.yout[:, self.number_dof * n], "m")
-                    .to(displacement_units)
-                    .m,
-                    z=Q_(self.yout[:, self.number_dof * n + 1], "m")
-                    .to(displacement_units)
-                    .m,
+                    y=Q_(self.yout[:, ndof * n], "m").to(displacement_units).m,
+                    z=Q_(self.yout[:, ndof * n + 1], "m").to(displacement_units).m,
                     mode="lines",
                     line=dict(color=tableau_colors["blue"]),
                     name="Mean",
@@ -3731,11 +3754,11 @@ class TimeResponseResults(Results):
             )
 
         # plot center line
-        line = np.zeros(len(self.nodes_pos))
+        line = np.zeros(len(nodes_pos))
 
         fig.add_trace(
             go.Scatter3d(
-                x=Q_(self.nodes_pos, "m").to(rotor_length_units).m,
+                x=Q_(nodes_pos, "m").to(rotor_length_units).m,
                 y=line,
                 z=line,
                 mode="lines",
diff --git a/ross/rotor_assembly.py b/ross/rotor_assembly.py
index c15c2afb..03708228 100644
--- a/ross/rotor_assembly.py
+++ b/ross/rotor_assembly.py
@@ -313,6 +313,11 @@ class Rotor(object):
         self.nodes = list(range(len(self.nodes_pos)))
         self.L = nodes_pos[-1]
 
+        if "n_link" in df.columns:
+            self.link_nodes = list(df["n_link"].dropna().unique().astype(int))
+        else:
+            self.link_nodes = []
+
         # rotor mass can also be calculated with self.M()[::4, ::4].sum()
         self.m_disks = np.sum([disk.m for disk in self.disk_elements])
         self.m_shaft = np.sum([sh_el.m for sh_el in self.shaft_elements])
@@ -2398,9 +2403,7 @@ class Rotor(object):
         """
         t_, yout, xout = self.time_response(speed, F, t)
 
-        results = TimeResponseResults(
-            t, yout, xout, self.nodes, self.nodes_pos, self.number_dof
-        )
+        results = TimeResponseResults(self, t, yout, xout)
 
         return results
 


Test patch:
diff --git a/ross/tests/test_results.py b/ross/tests/test_results.py
index eee27d4d..384506d4 100644
--- a/ross/tests/test_results.py
+++ b/ross/tests/test_results.py
@@ -138,6 +138,4 @@ def test_save_load_timeresponse(rotor1):
     assert response2.t.all() == response.t.all()
     assert response2.yout.all() == response.yout.all()
     assert response2.xout.all() == response.xout.all()
-    assert np.array(response2.nodes_list).all() == np.array(response.nodes_list).all()
-    assert np.array(response2.nodes_pos).all() == np.array(response.nodes_pos).all()
-    assert response2.number_dof == response.number_dof
+    assert response2.rotor == response.rotor