Repository: ross-rotordynamics/ross
Pull request #732 (2020-12-17 18:26:58)

Add function to collect and organize data from figures
This function takes a `go.Figure()` object from Plotly and converts the data to a `DataFrame` from Pandas. It works only for Scatter Plots (XY) and does no support subplots or polar plots.

This function is suitable to the following analyzes:
- `run_freq_response()`
- `run_unbalance_response()`
- `run_time_response()`


Patch:
diff --git a/ross/__init__.py b/ross/__init__.py
index 4aa5e388..70d762bd 100644
--- a/ross/__init__.py
+++ b/ross/__init__.py
@@ -11,6 +11,6 @@ from .point_mass import *
 from .results import *
 from .rotor_assembly import *
 from .shaft_element import *
-from .utils import visualize_matrix
+from .utils import visualize_matrix, get_data_from_figure
 
 _pio.templates.default = "ross"
diff --git a/ross/results.py b/ross/results.py
index c502a028..9ee5f899 100644
--- a/ross/results.py
+++ b/ross/results.py
@@ -165,56 +165,21 @@ class CriticalSpeedResults(Results):
 
     Parameters
     ----------
-    _wn : array
+    wn : array
         Undamped critical speeds array.
-    _wd : array
+    wd : array
         Undamped critical speeds array.
     log_dec : array
         Logarithmic decrement for each critical speed.
     damping_ratio : array
         Damping ratio for each critical speed.
-    whirl_direction : array
-        Whirl direction for each critical speed. Can be forward, backward or mixed.
     """
 
-    def __init__(self, _wn, _wd, log_dec, damping_ratio, whirl_direction):
-        self._wn = _wn
-        self._wd = _wd
+    def __init__(self, wn, wd, log_dec, damping_ratio):
+        self.wn = wn
+        self.wd = wd
         self.log_dec = log_dec
         self.damping_ratio = damping_ratio
-        self.whirl_direction = whirl_direction
-
-    def wn(self, frequency_units="rad/s"):
-        """Convert units for undamped critical speeds.
-
-        Parameters
-        ----------
-        frequency_units : str, optional
-            Critical speeds units.
-            Default is rad/s
-
-        Returns
-        -------
-        wn : array
-            Undamped critical speeds array.
-        """
-        return Q_(self.__dict__["_wn"], "rad/s").to(frequency_units).m
-
-    def wd(self, frequency_units="rad/s"):
-        """Convert units for damped critical speeds.
-
-        Parameters
-        ----------
-        frequency_units : str, optional
-            Critical speeds units.
-            Default is rad/s
-
-        Returns
-        -------
-        wd : array
-            Undamped critical speeds array.
-        """
-        return Q_(self.__dict__["_wd"], "rad/s").to(frequency_units).m
 
 
 class ModalResults(Results):
@@ -1115,7 +1080,7 @@ class FrequencyResponseResults(Results):
     speed_range : array
         Array with the speed range in rad/s.
     number_dof : int
-        Number of degrees of freedom per node.
+        Number of degrees of freedom per node
 
     Returns
     -------
@@ -1308,8 +1273,8 @@ class FrequencyResponseResults(Results):
 
         if fig is None:
             fig = go.Figure()
-        idx = len(fig.data)
 
+        idx = len(fig.data)
         fig.add_trace(
             go.Scatter(
                 x=frequency_range,
@@ -3591,10 +3556,9 @@ class TimeResponseResults(Results):
         fig=None,
         **kwargs,
     ):
-        """Plot time response.
+        """Plot time response for a single DoF using Plotly.
 
-        This method plots the time response given a tuple of probes with their nodes
-        and orientations.
+        This function will take a rotor object and plot its time response using Plotly.
 
         Parameters
         ----------
diff --git a/ross/rotor_assembly.py b/ross/rotor_assembly.py
index b7276031..03708228 100644
--- a/ross/rotor_assembly.py
+++ b/ross/rotor_assembly.py
@@ -637,7 +637,7 @@ class Rotor(object):
 
         return modal_results
 
-    def run_critical_speed(self, speed_range=None, num_modes=12, rtol=0.005):
+    def run_critical_speed(self, num_modes=12, sparse=True, rtol=0.005):
         """Calculate the critical speeds and damping ratios for the rotor model.
 
         This function runs an iterative method over "run_modal()" to minimize
@@ -657,16 +657,19 @@ class Rotor(object):
 
         Parameters
         ----------
-        speed_range : tuple
-            Tuple (start, end) with the desired range of frequencies (rad/s).
-            The function returns all eigenvalues within this range.
         num_modes : int, optional
             The number of eigenvalues and eigenvectors to be calculated using ARPACK.
             If sparse=True, it determines the number of eigenvalues and eigenvectors
             to be calculated. It must be smaller than Rotor.ndof - 1. It is not
             possible to compute all eigenvectors of a matrix with ARPACK.
-            If speed_range is not None, num_modes is overrided.
+            If sparse=False, num_modes does not have any effect over the method.
             Default is 12.
+        sparse : bool, optional
+            If True, ARPACK is used to calculate a desired number (according to
+            num_modes) or eigenvalues and eigenvectors.
+            If False, scipy.linalg.eig() is used to calculate all the eigenvalues and
+            eigenvectors.
+            Default is True.
         rtol : float, optional
             Tolerance (relative) for termination. Applied to scipy.optimize.newton.
             Default is 0.005 (0.5%).
@@ -681,65 +684,33 @@ class Rotor(object):
 
         Examples
         --------
-        >>> import ross as rs
-        >>> rotor = rs.rotor_example()
-
-        Finding the first Nth critical speeds
+        >>> rotor = rotor_example()
         >>> results = rotor.run_critical_speed(num_modes=8)
-        >>> np.round(results.wd())
+        >>> np.round(results.wd)
         array([ 92.,  96., 271., 300.])
-        >>> np.round(results.wn())
+        >>> np.round(results.wn)
         array([ 92.,  96., 271., 300.])
-
-        Finding the first critical speeds within a speed range
-        >>> results = rotor.run_critical_speed(speed_range=(100, 1000))
-        >>> np.round(results.wd())
-        array([271., 300., 636., 867.])
-
-        Changing output units
-        >>> np.round(results.wd("rpm"))
-        array([2590., 2868., 6074., 8278.])
-
-        Retrieving whirl directions
-        >>> results.whirl_direction # doctest: +ELLIPSIS
-        array([...
         """
-        num_modes = (self.ndof - 4) * 2 if speed_range is not None else num_modes
-
-        modal = self.run_modal(0, num_modes)
-        _wn = modal.wn
-        _wd = modal.wd
+        _wn = self.run_modal(0, num_modes, sparse).wn
         wn = np.zeros_like(_wn)
-        wd = np.zeros_like(_wd)
+        wd = np.zeros_like(_wn)
+        log_dec = np.zeros_like(_wn)
+        damping_ratio = np.zeros_like(_wn)
 
         for i in range(len(wn)):
-            wn_func = lambda s: (s - self.run_modal(s, num_modes).wn[i])
+            wn_func = lambda s: (s - self.run_modal(s, num_modes, sparse).wn[i])
             wn[i] = newton(func=wn_func, x0=_wn[i], rtol=rtol)
 
-        for i in range(len(wd)):
-            wd_func = lambda s: (s - self.run_modal(s, num_modes).wd[i])
-            wd[i] = newton(func=wd_func, x0=_wd[i], rtol=rtol)
+        for i in range(len(wn)):
+            wd_func = lambda s: (s - self.run_modal(s, num_modes, sparse).wd[i])
+            wd[i] = newton(func=wd_func, x0=wn[i], rtol=rtol)
 
-        log_dec = np.zeros_like(wn)
-        damping_ratio = np.zeros_like(wn)
-        whirl_direction = list(np.zeros_like(wn))
         for i, s in enumerate(wd):
-            modal = self.run_modal(s, num_modes)
+            modal = self.run_modal(s, num_modes, sparse)
             log_dec[i] = modal.log_dec[i]
             damping_ratio[i] = modal.damping_ratio[i]
-            whirl_direction[i] = modal.whirl_direction()[i]
-
-        whirl_direction = np.array(whirl_direction)
-        if speed_range is not None:
-            vmin, vmax = speed_range
-            idx = np.where((wd >= vmin) & (wd <= vmax))
-            wn = wn[idx]
-            wd = wd[idx]
-            log_dec = log_dec[idx]
-            damping_ratio = damping_ratio[idx]
-            whirl_direction = whirl_direction[idx]
 
-        return CriticalSpeedResults(wn, wd, log_dec, damping_ratio, whirl_direction)
+        return CriticalSpeedResults(wn, wd, log_dec, damping_ratio)
 
     def convergence(self, n_eigval=0, err_max=1e-02):
         """Run convergence analysis.
@@ -1127,7 +1098,7 @@ class Rotor(object):
         (61,)
         """
         critical_speeds = self.run_critical_speed(num_modes=num_modes, rtol=rtol)
-        omega = critical_speeds._wd
+        omega = critical_speeds.wd
         damping = critical_speeds.damping_ratio
         damping = np.array([d if d >= 0.005 else 0.005 for d in damping])
 
diff --git a/ross/utils.py b/ross/utils.py
index cd930f43..0e98aa2a 100644
--- a/ross/utils.py
+++ b/ross/utils.py
@@ -537,3 +537,64 @@ def intersection(x1, y1, x2, y2):
     xy0 = T[2:, in_range]
     xy0 = xy0.T
     return xy0[:, 0], xy0[:, 1]
+
+
+def get_data_from_figure(fig):
+    """Get data from a Plotly ccatter plot (XY) and convert to a pd.DataFrame.
+
+    This function takes a go.Figure() object from Plotly and converts the data to a
+    DataFrame from Pandas. It works only for Scatter Plots (XY) and does no support
+    subplots or polar plots.
+
+    This function is suitable to the following analyzes:
+        - run_freq_response()
+        - run_unbalance_response()
+        - run_time_response()
+
+    Parameters
+    ----------
+    fig : go.Figure
+        A Plotly figure generated by go.Figure() command.
+
+    Returns
+    -------
+    df : pd.DataFrame
+        A DataFrame with the fig.data.
+
+    Examples
+    --------
+    >>> import ross as rs
+    >>> rotor = rs.rotor_example()
+
+    Get post-processed data from an unbalance response
+    >>> resp = rotor.run_unbalance_response(
+    ...     3, 0.001, 0.0, np.linspace(0, 1000, 101)
+    ... )
+    >>> fig = resp.plot_magnitude(probe=[(3, 0.0, "probe1"), (3, np.pi/2, "probe2")])
+    >>> df = rs.get_data_from_figure(fig)
+
+    Use the probe tag to navigate through pandas data
+    Index 0 for frequency array
+    >>> df["probe1"][0] # doctest: +ELLIPSIS
+    array([   0.,   10.,   20.,   30.,...
+
+    Index 1 for amplitude array
+    >>> df["probe1"][1] # doctest: +ELLIPSIS
+    array([0.00000000e+00,...
+
+    Or use "iloc" to obtain the desired array from pandas
+    >>> df.iloc[1, 0] # doctest: +ELLIPSIS
+    array([0.00000000e+00, 1.60570205e-07, 6.63435909e-07,...
+    """
+    dict_data = {data["name"]: {} for data in fig.data}
+
+    xaxis_label = fig.layout.xaxis.title.text
+    yaxis_label = fig.layout.yaxis.title.text
+
+    for i, data in enumerate(fig.data):
+        dict_data[data["name"]][xaxis_label] = data.x
+        dict_data[data["name"]][yaxis_label] = data.y
+
+    df = pd.DataFrame(dict_data)
+
+    return df


Test patch:
diff --git a/ross/tests/test_results.py b/ross/tests/test_results.py
index 94aa364f..384506d4 100644
--- a/ross/tests/test_results.py
+++ b/ross/tests/test_results.py
@@ -5,8 +5,8 @@ import numpy as np
 import pytest
 from numpy.testing import assert_allclose, assert_almost_equal, assert_equal
 
-from ross.results import *
 from ross.rotor_assembly import *
+from ross.results import *
 
 
 @pytest.fixture
@@ -35,8 +35,8 @@ def test_save_load_criticalspeed(rotor1):
     response.save(file)
     response2 = CriticalSpeedResults.load(file)
 
-    assert response2._wn.all() == response._wn.all()
-    assert response2._wd.all() == response._wd.all()
+    assert response2.wn.all() == response.wn.all()
+    assert response2.wd.all() == response.wd.all()
     assert response2.log_dec.all() == response.log_dec.all()
     assert response2.damping_ratio.all() == response.damping_ratio.all()
 
@@ -58,10 +58,7 @@ def test_save_load_modal(rotor1):
     assert response2.ndof == response.ndof
     assert np.array(response2.nodes).all() == np.array(response.nodes).all()
     assert np.array(response2.nodes_pos).all() == np.array(response.nodes_pos).all()
-    assert (
-        np.array(response2.shaft_elements_length).all()
-        == np.array(response.shaft_elements_length).all()
-    )
+    assert np.array(response2.shaft_elements_length).all() == np.array(response.shaft_elements_length).all()
 
 
 def test_save_load_freqresponse(rotor1):
diff --git a/ross/tests/test_rotor_assembly.py b/ross/tests/test_rotor_assembly.py
index 14f7bf6f..7b0f1fd4 100644
--- a/ross/tests/test_rotor_assembly.py
+++ b/ross/tests/test_rotor_assembly.py
@@ -961,13 +961,13 @@ def test_run_critical_speed(rotor5, rotor6):
     log_dec6 = np.zeros_like(wd6)
     damping_ratio6 = np.zeros_like(wd6)
 
-    assert_almost_equal(results5._wd, wd5, decimal=4)
-    assert_almost_equal(results5._wn, wn5, decimal=4)
+    assert_almost_equal(results5.wd, wd5, decimal=4)
+    assert_almost_equal(results5.wn, wn5, decimal=4)
     assert_almost_equal(results5.log_dec, log_dec5, decimal=4)
     assert_almost_equal(results5.damping_ratio, damping_ratio5, decimal=4)
 
-    assert_almost_equal(results6._wd, wd6, decimal=4)
-    assert_almost_equal(results6._wn, wn6, decimal=4)
+    assert_almost_equal(results6.wd, wd6, decimal=4)
+    assert_almost_equal(results6.wn, wn6, decimal=4)
     assert_almost_equal(results6.log_dec, log_dec6, decimal=4)
     assert_almost_equal(results6.damping_ratio, damping_ratio6, decimal=4)