Merge pull request #1004 from raphaeltimbo/add-campbell-mode-shape

Add campbell mode shape plot

ross/results.py

@@ -85,6 +85,12 @@ def save(self, file):
             data = {}
 
         data[f"{self.__class__.__name__}"] = args
+
+        try:
+            del data["CampbellResults"]["modal_results"]
+        except KeyError:
+            pass
+
         with open(file, "w") as f:
             toml.dump(data, f, encoder=toml.TomlNumpyEncoder())
 
@@ -148,6 +154,8 @@ def load(cls, file):
 
         data = toml.load(file)
         data = list(data.values())[0]
+        if cls == CampbellResults:
+            data["modal_results"] = None
         for key, value in data.items():
             if key == "rotor":
                 aux_file = str(file)[:-5] + "_rotor" + str(file)[-5:]
@@ -1354,19 +1362,24 @@ class CampbellResults(Results):
         Array with the Logarithmic decrement
     whirl_values : array
         Array with the whirl values (0, 0.5 or 1)
+    modal_results : dict
+        Dictionary with the modal results for each speed in the speed range.
 
     Returns
     -------
     fig : Plotly graph_objects.Figure()
         The figure object with the plot.
     """
 
-    def __init__(self, speed_range, wd, log_dec, damping_ratio, whirl_values):
+    def __init__(
+        self, speed_range, wd, log_dec, damping_ratio, whirl_values, modal_results
+    ):
         self.speed_range = speed_range
         self.wd = wd
         self.log_dec = log_dec
         self.damping_ratio = damping_ratio
         self.whirl_values = whirl_values
+        self.modal_results = modal_results
 
     @check_units
     def plot(
@@ -1595,6 +1608,87 @@ def plot(
 
         return fig
 
+    def plot_with_mode_shape(
+        self,
+        harmonics=[1],
+        frequency_units="rad/s",
+        damping_parameter="log_dec",
+        frequency_range=None,
+        damping_range=None,
+        fig=None,
+        **kwargs,
+    ):
+
+        try:
+            from ipywidgets import VBox
+        except ImportError:
+            raise ImportError("Please install ipywidgets to use this feature.")
+
+        def _plot_with_mode_shape_callback(trace, points, state):
+            point_idx = points.point_inds
+            if len(point_idx) > 0:
+                frequency = trace.x[point_idx][0]
+                natural_frequency = trace.y[point_idx][0]
+
+                # run modal analysis for desired frequency
+                modal = self.modal_results[Q_(frequency, frequency_units).to("rad/s").m]
+
+                # identify index of desired mode
+                idx = (
+                    np.abs(
+                        modal.wd - Q_(natural_frequency, frequency_units).to("rad/s").m
+                    )
+                ).argmin()
+
+                new_plot_mode_3d = modal.plot_mode_3d(
+                    idx,
+                    frequency_units=frequency_units,
+                    damping_parameter=damping_parameter,
+                )
+                with plot_mode_3d.batch_update():
+                    # update title
+                    plot_mode_3d.layout["title"]["text"] = new_plot_mode_3d.layout[
+                        "title"
+                    ]["text"]
+                    for data, new_data in zip(plot_mode_3d_data, new_plot_mode_3d.data):
+                        for param in data:
+                            data[param] = new_data[param]
+
+        camp_fig = self.plot(
+            harmonics=harmonics,
+            frequency_units=frequency_units,
+            damping_parameter=damping_parameter,
+            frequency_range=frequency_range,
+            damping_range=damping_range,
+            fig=fig,
+            **kwargs,
+        )
+        camp_fig = go.FigureWidget(camp_fig)
+        plot_mode_3d = self.modal_results[self.speed_range[0]].plot_mode_3d(
+            0, frequency_units=frequency_units, damping_parameter=damping_parameter
+        )
+        plot_mode_3d = go.FigureWidget(plot_mode_3d)
+        plot_mode_3d_data = plot_mode_3d.data
+
+        for scatter in camp_fig.data:
+            scatter.on_click(_plot_with_mode_shape_callback)
+
+        return VBox([camp_fig, plot_mode_3d])
+
+    def save(self, file):
+        # TODO save modal results
+        warn(
+            "The CampbellResults.save method is not saving the attribute 'modal_results' for now."
+        )
+        super().save(file)
+
+    @classmethod
+    def load(cls, file):
+        warn(
+            "The CampbellResults.save method is not saving the attribute 'modal_results' for now."
+        )
+        return super().load(file)
+
 
 class FrequencyResponseResults(Results):
     """Class used to store results and provide plots for Frequency Response.

ross/rotor_assembly.py

@@ -2101,8 +2101,10 @@ def run_campbell(self, speed_range, frequencies=6, frequency_type="wd"):
 
         results = np.zeros([len(speed_range), frequencies, 6])
 
+        modal_results = {}
         for i, w in enumerate(speed_range):
             modal = self.run_modal(speed=w, num_modes=2 * frequencies)
+            modal_results[w] = modal
 
             if frequency_type == "wd":
                 results[i, :, 0] = modal.wd[:frequencies]
@@ -2125,6 +2127,7 @@ def run_campbell(self, speed_range, frequencies=6, frequency_type="wd"):
             log_dec=results[..., 1],
             damping_ratio=results[..., 2],
             whirl_values=results[..., 3],
+            modal_results=modal_results,
         )
 
         return results