Repository: ross-rotordynamics/ross (python)

Issue:
CampbellResults.plot() - damping ratio and filtering
I have started using Ross 1.0.0, and have some two suggestions for improving plotting of Campbell results. My goal is to incorporate changes that may be of interest to other users, so I do not have to change the modules manually every time I update.

Damping ratio: We use damping ratio, not log dec. There should be an option for plotting using damping ratio:
1. The damping ratio could be passed to CampbellResults, or calculated in plot. damping = 1.0/np.sqrt(1+(2*np.pi/self.log_dec)**2)
2. The color axis should be reversed when using damping ratio: coloraxis_colorscale="rdbu"
3. Change text: coloraxis_colorbar=dict(title=dict(text="<b>Damping ratio</b>", side="right"))

Filtering: The RAPPID rotordynamic software has options for filtering the modes shown on the campbell diagram. For example, modes with high damping may not be of interest, so we can supply options for a range of damping ratios that should be included (or left out) of the plot. Also, I had some 0.5 cpm modes that I didn't want to show, so filtering of frequencies could also be optional. Filtering of machine speeds is easily obtained by limiting the X-axis, which may be done outside of the plot function. 
I did this filtering manually, something like:
                whirl_mask = whirl_i == whirl_dir
                # lpa 2021-02-25 added a second masking for high damping, and zero modes
                damping_mask = log_dec_i < lpa_damping
                freq_mask = w_i > lpa_low_freq
                lpa_mask = whirl_mask & damping_mask & freq_mask

Pull request description:
Add damping parameter and filtering to campbell
This PR adds three additional arguments to the Campbell plot:
- damping_parameter: choose between log_dec or damping_ratio;
- frequency_range: filter out modes that are out of this frequency range;
- damping_range: filter out modes that are out of this damping range;

Closes #783.

Patch:
diff --git a/ross/__init__.py b/ross/__init__.py
index d83454d0..37790111 100644
--- a/ross/__init__.py
+++ b/ross/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "1.0.0"
+__version__ = "1.0.1"
 from plotly import io as _pio
 
 import ross.plotly_theme
@@ -11,6 +11,7 @@ from .point_mass import *
 from .results import *
 from .rotor_assembly import *
 from .shaft_element import *
+from .units import Q_
 from .utils import get_data_from_figure, visualize_matrix
 
 _pio.templates.default = "ross"
diff --git a/ross/results.py b/ross/results.py
index 4672ea47..d14a772c 100644
--- a/ross/results.py
+++ b/ross/results.py
@@ -16,7 +16,7 @@ from plotly.subplots import make_subplots
 from scipy import linalg as la
 
 from ross.plotly_theme import tableau_colors
-from ross.units import Q_
+from ross.units import Q_, check_units
 from ross.utils import intersection
 
 __all__ = [
@@ -952,13 +952,24 @@ class CampbellResults(Results):
         The figure object with the plot.
     """
 
-    def __init__(self, speed_range, wd, log_dec, whirl_values):
+    def __init__(self, speed_range, wd, log_dec, damping_ratio, whirl_values):
         self.speed_range = speed_range
         self.wd = wd
         self.log_dec = log_dec
+        self.damping_ratio = damping_ratio
         self.whirl_values = whirl_values
 
-    def plot(self, harmonics=[1], frequency_units="rad/s", fig=None, **kwargs):
+    @check_units
+    def plot(
+        self,
+        harmonics=[1],
+        frequency_units="rad/s",
+        damping_parameter="log_dec",
+        frequency_range=None,
+        damping_range=None,
+        fig=None,
+        **kwargs,
+    ):
         """Create Campbell Diagram figure using Plotly.
 
         Parameters
@@ -969,6 +980,18 @@ class CampbellResults(Results):
         frequency_units : str, optional
             Frequency units.
             Default is "rad/s"
+        damping_parameter : str, optional
+            Define which value to show for damping. We can use "log_dec" or "damping_ratio".
+            Default is "log_dec".
+        frequency_range : tuple, pint.Quantity(tuple), optional
+            Tuple with (min, max) values for the frequencies that will be plotted.
+            Frequencies that are not within the range are filtered out and are not plotted.
+            It is possible to use a pint Quantity (e.g. Q_((2000, 1000), "RPM")).
+            Default is None (no filter).
+        damping_range : tuple, optional
+            Tuple with (min, max) values for the damping parameter that will be plotted.
+            Damping values that are not within the range are filtered out and are not plotted.
+            Default is None (no filter).
         fig : Plotly graph_objects.Figure()
             The figure object with the plot.
         kwargs : optional
@@ -980,12 +1003,39 @@ class CampbellResults(Results):
         -------
         fig : Plotly graph_objects.Figure()
             The figure object with the plot.
+
+        Examples
+        --------
+        >>> import ross as rs
+        >>> import numpy as np
+        >>> Q_ = rs.Q_
+        >>> rotor = rs.rotor_example()
+        >>> speed = np.linspace(0, 400, 101)
+        >>> camp = rotor.run_campbell(speed)
+        >>> fig = camp.plot(
+        ...     harmonics=[1, 2],
+        ...     damping_parameter="damping_ratio",
+        ...     frequency_range=Q_((2000, 10000), "RPM"),
+        ...     damping_range=(-0.1, 100),
+        ...     frequency_units="RPM",
+        ... )
         """
-        wd = Q_(self.wd, "rad/s").to(frequency_units).m
+        if damping_parameter == "log_dec":
+            damping_values = self.log_dec
+            title_text = "<b>Log Dec</b>"
+        elif damping_parameter == "damping_ratio":
+            damping_values = self.damping_ratio
+            title_text = "<b>Damping Ratio</b>"
+        else:
+            raise ValueError(
+                f"damping_parameter can be 'log_dec' or 'damping_ratio'. {damping_parameter} is not valid"
+            )
+
+        wd = self.wd
         num_frequencies = wd.shape[1]
-        log_dec = self.log_dec
+
         whirl = self.whirl_values
-        speed_range = Q_(self.speed_range, "rad/s").to(frequency_units).m
+        speed_range = self.speed_range
 
         if fig is None:
             fig = go.Figure()
@@ -994,7 +1044,7 @@ class CampbellResults(Results):
             coloraxis_cmin=0.0,
             coloraxis_cmax=1.0,
             coloraxis_colorscale="rdbu",
-            coloraxis_colorbar=dict(title=dict(text="<b>Log Dec</b>", side="right")),
+            coloraxis_colorbar=dict(title=dict(text=title_text, side="right")),
         )
         for k, v in default_values.items():
             kwargs.setdefault(k, v)
@@ -1014,11 +1064,22 @@ class CampbellResults(Results):
                 crit_x.extend(x)
                 crit_y.extend(y)
 
+        # filter frequency range
+        if frequency_range is not None:
+            crit_x_filtered = []
+            crit_y_filtered = []
+            for x, y in zip(crit_x, crit_y):
+                if frequency_range[0] < y < frequency_range[1]:
+                    crit_x_filtered.append(x)
+                    crit_y_filtered.append(y)
+            crit_x = crit_x_filtered
+            crit_y = crit_y_filtered
+
         if len(crit_x) and len(crit_y):
             fig.add_trace(
                 go.Scatter(
-                    x=crit_x,
-                    y=crit_y,
+                    x=Q_(crit_x, "rad/s").to(frequency_units).m,
+                    y=Q_(crit_y, "rad/s").to(frequency_units).m,
                     mode="markers",
                     marker=dict(symbol="x", color="black"),
                     name="Crit. Speed",
@@ -1037,17 +1098,29 @@ class CampbellResults(Results):
             for i in range(num_frequencies):
                 w_i = wd[:, i]
                 whirl_i = whirl[:, i]
-                log_dec_i = log_dec[:, i]
+                damping_values_i = damping_values[:, i]
 
                 whirl_mask = whirl_i == whirl_dir
-                if any(check for check in whirl_mask):
+                mask = whirl_mask
+                if frequency_range is not None:
+                    frequency_mask = (w_i > frequency_range[0]) & (
+                        w_i < frequency_range[1]
+                    )
+                    mask = mask & frequency_mask
+                if damping_range is not None:
+                    damping_mask = (damping_values_i > damping_range[0]) & (
+                        damping_values_i < damping_range[1]
+                    )
+                    mask = mask & damping_mask
+
+                if any(check for check in mask):
                     fig.add_trace(
                         go.Scatter(
-                            x=speed_range[whirl_mask],
-                            y=w_i[whirl_mask],
+                            x=Q_(speed_range[mask], "rad/s").to(frequency_units).m,
+                            y=Q_(w_i[mask], "rad/s").to(frequency_units).m,
                             marker=dict(
                                 symbol=mark,
-                                color=log_dec_i[whirl_mask],
+                                color=damping_values_i[mask],
                                 coloraxis="coloraxis",
                             ),
                             mode="markers",
@@ -1063,8 +1136,8 @@ class CampbellResults(Results):
         for j, h in enumerate(harmonics):
             fig.add_trace(
                 go.Scatter(
-                    x=speed_range,
-                    y=h * speed_range,
+                    x=Q_(speed_range, "rad/s").to(frequency_units).m,
+                    y=h * Q_(speed_range, "rad/s").to(frequency_units).m,
                     mode="lines",
                     line=dict(dash="dashdot", color=list(tableau_colors)[j]),
                     name="{}x speed".format(h),
@@ -1089,12 +1162,15 @@ class CampbellResults(Results):
 
         fig.update_xaxes(
             title_text=f"Frequency ({frequency_units})",
-            range=[np.min(speed_range), np.max(speed_range)],
+            range=[
+                np.min(Q_(speed_range, "rad/s").to(frequency_units).m),
+                np.max(Q_(speed_range, "rad/s").to(frequency_units).m),
+            ],
             exponentformat="none",
         )
         fig.update_yaxes(
             title_text=f"Natural Frequencies ({frequency_units})",
-            range=[0, 1.1 * np.max(wd)],
+            range=[0, 1.1 * np.max(Q_(wd, "rad/s").to(frequency_units).m)],
         )
         fig.update_layout(
             legend=dict(
diff --git a/ross/rotor_assembly.py b/ross/rotor_assembly.py
index 25df9a08..fbec14dd 100644
--- a/ross/rotor_assembly.py
+++ b/ross/rotor_assembly.py
@@ -2082,7 +2082,7 @@ class Rotor(object):
         # whirl[2](y axis), 3]
         self._check_frequency_array(speed_range)
 
-        results = np.zeros([len(speed_range), frequencies, 5])
+        results = np.zeros([len(speed_range), frequencies, 6])
 
         for i, w in enumerate(speed_range):
             modal = self.run_modal(speed=w, num_modes=2 * frequencies)
@@ -2090,21 +2090,24 @@ class Rotor(object):
             if frequency_type == "wd":
                 results[i, :, 0] = modal.wd[:frequencies]
                 results[i, :, 1] = modal.log_dec[:frequencies]
-                results[i, :, 2] = modal.whirl_values()[:frequencies]
+                results[i, :, 2] = modal.damping_ratio[:frequencies]
+                results[i, :, 3] = modal.whirl_values()[:frequencies]
             else:
                 idx = modal.wn.argsort()
                 results[i, :, 0] = modal.wn[idx][:frequencies]
                 results[i, :, 1] = modal.log_dec[idx][:frequencies]
-                results[i, :, 2] = modal.whirl_values()[idx][:frequencies]
+                results[i, :, 2] = modal.damping_ratio[idx][:frequencies]
+                results[i, :, 3] = modal.whirl_values()[idx][:frequencies]
 
-            results[i, :, 3] = w
-            results[i, :, 4] = modal.wn[:frequencies]
+            results[i, :, 4] = w
+            results[i, :, 5] = modal.wn[:frequencies]
 
         results = CampbellResults(
             speed_range=speed_range,
             wd=results[..., 0],
             log_dec=results[..., 1],
-            whirl_values=results[..., 2],
+            damping_ratio=results[..., 2],
+            whirl_values=results[..., 3],
         )
 
         return results
diff --git a/ross/units.py b/ross/units.py
index 9883b29e..eb643fe0 100644
--- a/ross/units.py
+++ b/ross/units.py
@@ -28,6 +28,7 @@ units = {
     "odr": "meter",
     "speed": "radian/second",
     "frequency": "radian/second",
+    "frequency_range": "radian/second",
     "m": "kg",
     "mx": "kg",
     "my": "kg",


Test patch:
diff --git a/ross/tests/test_results.py b/ross/tests/test_results.py
index 94aa364f..675b9e32 100644
--- a/ross/tests/test_results.py
+++ b/ross/tests/test_results.py
@@ -5,6 +5,7 @@ import numpy as np
 import pytest
 from numpy.testing import assert_allclose, assert_almost_equal, assert_equal
 
+from ross import Q_
 from ross.results import *
 from ross.rotor_assembly import *
 
@@ -142,3 +143,23 @@ def test_save_load_timeresponse(rotor1):
     assert response2.yout.all() == response.yout.all()
     assert response2.xout.all() == response.xout.all()
     assert response2.rotor == response.rotor
+
+
+def test_campbell_plot(rotor1):
+    speed = np.linspace(0, 400, 101)
+    camp = rotor1.run_campbell(speed)
+    fig = camp.plot(
+        harmonics=[1, 2],
+        damping_parameter="damping_ratio",
+        frequency_range=Q_((2000, 10000), "RPM"),
+        damping_range=(-0.1, 100),
+        frequency_units="RPM",
+    )
+    crit_array_x = np.array(
+        [2590.2641754, 1306.51513941, 2868.14592367, 1420.76907353, 3264.81334336]
+    )
+    crit_array_y = np.array(
+        [2590.2641754, 2613.03027882, 2868.14592367, 2841.53814705, 6529.62668672]
+    )
+    assert_allclose(fig.data[0]["x"], crit_array_x)
+    assert_allclose(fig.data[0]["y"], crit_array_y)
diff --git a/ross/tests/test_units.py b/ross/tests/test_units.py
index dbc4adb5..e9884b58 100644
--- a/ross/tests/test_units.py
+++ b/ross/tests/test_units.py
@@ -13,11 +13,11 @@ def test_new_units_loaded():
 def auxiliary_function():
     # fmt: off
     @check_units
-    def func(E, G_s, rho, L, idl, idr, odl, odr, speed, frequency, m, mx, my, Ip, Id,
+    def func(E, G_s, rho, L, idl, idr, odl, odr, speed, frequency, frequency_range, m, mx, my, Ip, Id,
              width, depth, i_d, o_d, kxx, kxy, kxz, kyx, kyy, kyz, kzx, kzy, kzz, cxx, cxy,
              cxz, cyx, cyy, cyz, czx, czy, czz, unbalance_magnitude, unbalance_phase):
         return (
-            E, G_s, rho, L, idl, idr, odl, odr, speed, frequency, m, mx, my, Ip, Id,
+            E, G_s, rho, L, idl, idr, odl, odr, speed, frequency, frequency_range,m, mx, my, Ip, Id,
             width, depth, i_d, o_d, kxx, kxy, kxz, kyx, kyy, kyz, kzx, kzy, kzz, cxx, cxy,
             cxz, cyx, cyy, cyz, czx, czy, czz, unbalance_magnitude, unbalance_phase
         )
@@ -28,7 +28,7 @@ def auxiliary_function():
 def test_units(auxiliary_function):
     # fmt: off
     results = auxiliary_function(
-        E=1, G_s=1, rho=1, L=1, idl=1, idr=1, odl=1, odr=1, speed=1, frequency=1,
+        E=1, G_s=1, rho=1, L=1, idl=1, idr=1, odl=1, odr=1, speed=1, frequency=1, frequency_range=(1, 1),
         m=1, mx=1, my=1, Ip=1, Id=1, width=1,  depth=1, i_d=1, o_d=1, kxx=1, kxy=1, kxz=1, kyx=1,
         kyy=1, kyz=1, kzx=1, kzy=1, kzz=1, cxx=1, cxy=1, cxz=1, cyx=1, cyy=1, cyz=1,
         czx=1, czy=1, czz=1, unbalance_magnitude=1, unbalance_phase=1
@@ -36,7 +36,7 @@ def test_units(auxiliary_function):
     # check if all available units are tested
     assert len(results) == len(units)
 
-    (E, G_s, rho, L, idl, idr, odl, odr, speed, frequency, m, mx, my, Ip, Id,
+    (E, G_s, rho, L, idl, idr, odl, odr, speed, frequency, frequency_range,m, mx, my, Ip, Id,
      width, depth, i_d, o_d, kxx, kxy, kxz, kyx, kyy, kyz, kzx, kzy, kzz, cxx, cxy,
      cxz, cyx, cyy, cyz, czx, czy, czz, unbalance_magnitude, unbalance_phase) = results
     # fmt: on
@@ -51,6 +51,7 @@ def test_units(auxiliary_function):
     assert odr == 1
     assert speed == 1
     assert frequency == 1
+    assert_allclose(frequency_range, (1, 1))
     assert m == 1
     assert mx == 1
     assert my == 1
@@ -94,6 +95,7 @@ def test_unit_Q_(auxiliary_function):
         odr=Q_(1, "meter"),
         speed=Q_(1, "radian/second"),
         frequency=Q_(1, "radian/second"),
+        frequency_range=Q_((1, 1), "radian/second"),
         m=Q_(1, "kg"),
         mx=Q_(1, "kg"),
         my=Q_(1, "kg"),
@@ -128,7 +130,7 @@ def test_unit_Q_(auxiliary_function):
     # check if all available units are tested
     assert len(results) == len(units)
     # fmt: off
-    (E, G_s, rho, L, idl, idr, odl, odr, speed, frequency, m, mx, my, Ip, Id,
+    (E, G_s, rho, L, idl, idr, odl, odr, speed, frequency, frequency_range,m, mx, my, Ip, Id,
      width, depth, i_d, o_d, kxx, kxy, kxz, kyx, kyy, kyz, kzx, kzy, kzz, cxx, cxy,
      cxz, cyx, cyy, cyz, czx, czy, czz, unbalance_magnitude, unbalance_phase) = results
     # fmt: on
@@ -143,6 +145,7 @@ def test_unit_Q_(auxiliary_function):
     assert odr == 1
     assert speed == 1
     assert frequency == 1
+    assert_allclose(frequency_range, (1, 1))
     assert m == 1
     assert mx == 1
     assert my == 1
@@ -186,6 +189,7 @@ def test_unit_Q_conversion(auxiliary_function):
         odr=Q_(1, "inches"),
         speed=Q_(1, "RPM"),
         frequency=Q_(1, "RPM"),
+        frequency_range=Q_((1, 1), "RPM"),
         m=Q_(1, "lb"),
         mx=Q_(1, "lb"),
         my=Q_(1, "lb"),
@@ -221,7 +225,7 @@ def test_unit_Q_conversion(auxiliary_function):
     assert len(results) == len(units)
 
     # fmt: off
-    (E, G_s, rho, L, idl, idr, odl, odr, speed, frequency, m, mx, my, Ip, Id,
+    (E, G_s, rho, L, idl, idr, odl, odr, speed, frequency, frequency_range,m, mx, my, Ip, Id,
      width, depth, i_d, o_d, kxx, kxy, kxz, kyx, kyy, kyz, kzx, kzy, kzz, cxx, cxy,
      cxz, cyx, cyy, cyz, czx, czy, czz, unbalance_magnitude, unbalance_phase) = results
     # fmt: on
@@ -236,6 +240,7 @@ def test_unit_Q_conversion(auxiliary_function):
     assert odr == 0.0254
     assert speed == 0.10471975511965977
     assert frequency == 0.10471975511965977
+    assert_allclose(frequency_range, (0.10471975511965977, 0.10471975511965977))
     assert m == 0.4535923700000001
     assert mx == 0.4535923700000001
     assert my == 0.4535923700000001