Rewritten plot routines to take into account a user-provided unit for…

… the frequency
RubendeBruin · Aug 17, 2024 · b97fa2f · b97fa2f
1 parent d333e01
commit b97fa2f
Show file tree

Hide file tree

Showing 7 changed files with 173 additions and 53 deletions.
diff --git a/src/mafredo/__init__.py b/src/mafredo/__init__.py
@@ -17,5 +17,6 @@
 
 from mafredo.hyddb1 import Hyddb1, Symmetry
 from mafredo.rao import Rao
+from mafredo.helpers import FrequencyUnit, Symmetry, MotionMode
 
-__all__ = ['Hyddb1', 'Symmetry', 'Rao']
+__all__ = ['Hyddb1', 'Symmetry', 'Rao', 'FrequencyUnit', 'Symmetry', 'MotionMode']
diff --git a/src/mafredo/helpers.py b/src/mafredo/helpers.py
@@ -3,6 +3,21 @@
 from scipy.optimize import fsolve
 from enum import Enum
 
+class FrequencyUnit(Enum):
+    Hz = 0
+    rad_s = 1
+    seconds = 2
+
+    def to_unit(self, omega):
+        if self == FrequencyUnit.Hz:
+            return 'Hz', omega / (2*np.pi)
+        elif self == FrequencyUnit.rad_s:
+            return 'rad/s', omega
+        elif self == FrequencyUnit.seconds:
+            return 's', (2*np.pi) / omega
+        else:
+            raise ValueError('Unknown unit')
+
 class Symmetry(Enum):
     No = 0
     XZ = 1

diff --git a/src/mafredo/hyddb1.py b/src/mafredo/hyddb1.py
@@ -10,11 +10,13 @@
     MotionMode,
     Symmetry,
     MotionModeToStr,
+    FrequencyUnit
 )
 
 
 
 
+
 class Hyddb1(object):
     """
     Hydrodynamic Database 1st Order
@@ -1006,7 +1008,53 @@ def fixed_format(ident, sections):
             f.write(fixed_format("PARA2", [0, 0]))
             f.write("END\n")
 
-    def plot(self, adm=True, damp=True, amp=True, phase=True, do_show=True):
+    def _plot_amass_or_damping(self, data, ylab, unit : FrequencyUnit):
+
+        fig, axes = plt.subplots(3, 2, figsize=(10, 15))
+        axes = axes.flatten()
+
+        x_label, x = unit.to_unit(self._mass.omega.values)
+
+        for i in range(6):
+
+            for other in range(6):
+                if i == other:
+                    lw = 2
+                else:
+                    lw = 1
+
+                data = self._mass.sel(radiating_dof=i, influenced_dof=other).values
+
+                axes[i].plot(x, data, lw=lw, label=self._modes[other])
+
+            axes[i].set_title(self._modes[i])
+            axes[i].set_xlabel(f'[{x_label}]')
+            axes[i].set_ylabel(ylab)
+            if i == 5:
+                axes[i].legend()
+
+        return fig
+
+
+    def plot_added_mass(self, unit = FrequencyUnit.rad_s):
+        """Plots the added mass matrix"""
+
+        fig = self._plot_amass_or_damping(self._mass, "Added mass", unit)
+
+        fig.suptitle("Added mass\nDiagonal terms shown with thicker line")
+
+        return fig
+
+    def plot_damping(self, unit = FrequencyUnit.rad_s):
+        """Plots the damping matrix"""
+
+        fig = self._plot_amass_or_damping(self._damping, "Damping", unit)
+
+        fig.suptitle("Damping\nDiagonal terms shown with thicker line")
+
+        return fig
+
+    def plot(self, adm=True, damp=True, amp=True, phase=True, do_show=True, unit=FrequencyUnit.rad_s):
         """Produces a plot of the contents of the database
 
         Args:
@@ -1034,10 +1082,7 @@ def plot(self, adm=True, damp=True, amp=True, phase=True, do_show=True):
             axes = axes.flatten()
             for i in range(6):
                 force = self._force[i]
-                if force.n_wave_directions > 1:
-                    force._data["amplitude"].plot(ax=axes[i], cmap=plt.cm.GnBu)
-                else:
-                    force._data["amplitude"].plot(ax=axes[i])
+                force.plot_amplitude(ax=axes[i],unit=unit)
                 axes[i].set_title(self._modes[i])
             fig.suptitle("Force RAO amplitudes")
 
@@ -1051,10 +1096,7 @@ def plot(self, adm=True, damp=True, amp=True, phase=True, do_show=True):
             axes = axes.flatten()
             for i in range(6):
                 force = self._force[i]
-                if force.n_wave_directions > 1:
-                    force._data["phase"].plot(ax=axes[i], cmap=plt.cm.twilight_shifted)
-                else:
-                    force._data["phase"].plot(ax=axes[i])
+                force.plot_surface('phase',ax=axes[i],unit=unit)
                 axes[i].set_title(self._modes[i])
             fig.suptitle("Force RAO phase [rad]")
 
@@ -1063,51 +1105,13 @@ def plot(self, adm=True, damp=True, amp=True, phase=True, do_show=True):
         # Added mass
 
         if adm:
-            fig, axes = plt.subplots(3, 2, figsize=(10, 15))
-            axes = axes.flatten()
-            for i in range(6):
-
-                mode = self._modes[i]
-
-                for other in range(6):
-                    if i == other:
-                        lw = 2
-                    else:
-                        lw = 1
-                    self._mass.sel(radiating_dof=i, influenced_dof=other).plot(
-                        ax=axes[i], lw=lw, label=self._modes[other]
-                    )
-
-                axes[i].set_title(self._modes[i])
-                if i == 5:
-                    axes[i].legend()
-            fig.suptitle("Added mass\nDiagonal terms shown with thicker line")
-
+            fig = self.plot_added_mass(unit=unit)
             figs.append(fig)
 
         # Damping
 
         if damp:
-            fig, axes = plt.subplots(3, 2, figsize=(10, 15))
-            axes = axes.flatten()
-            for i in range(6):
-
-                mode = self._modes[i]
-
-                for other in range(6):
-                    if i == other:
-                        lw = 2
-                    else:
-                        lw = 1
-                    self._damping.sel(radiating_dof=i, influenced_dof=other).plot(
-                        ax=axes[i], lw=lw, label=self._modes[other]
-                    )
-
-                axes[i].set_title(self._modes[i])
-                if i == 5:
-                    axes[i].legend()
-            fig.suptitle("Damping \nDiagonal terms shown with thicker line")
-
+            fig = self.plot_damping(unit=unit)
             figs.append(fig)
 
         if do_show:

diff --git a/src/mafredo/rao.py b/src/mafredo/rao.py
@@ -2,7 +2,7 @@
 
 import xarray as xr
 import numpy as np
-from mafredo.helpers import expand_omega_dim_const, expand_direction_to_full_range
+from mafredo.helpers import expand_omega_dim_const, expand_direction_to_full_range, FrequencyUnit
 from mafredo.helpers import MotionMode, Symmetry, MotionModeToStr
 
 
@@ -89,9 +89,9 @@ class Rao(object):
 
     Plotting:
 
-    For plotting just use xarray:
+    .. method:: plot_amplitude
+    .. method:: plot_phase
 
-    >>> my_rao['ampltiude'].plot()
 
     others:
 
@@ -448,3 +448,67 @@ def __getitem__(self, key):
     def __str__(self):
         return str(self._data)
 
+
+    def plot_amplitude(self, ax=None, unit=FrequencyUnit.rad_s):
+        """Plots the amplitude"""
+        self.plot('amplitude', ax=ax, unit=unit)
+
+    def plot_phase(self, ax=None, unit=FrequencyUnit.rad_s):
+        """Plots the phase"""
+        self.plot('phase', ax=ax, unit=unit)
+
+    def plot(self, what = 'amplitude', ax=None, unit=FrequencyUnit.rad_s):
+        """Plots the amplitude"""
+        import matplotlib.pyplot as plt
+
+        if ax is None:
+            ax = plt.gca()
+
+        omega = self.omega
+
+        unit_label,x = unit.to_unit(omega)
+
+        headings = self._data['wave_direction'].values
+
+        for i, heading in enumerate(headings):
+            data = self._data[what].sel(wave_direction=heading).values
+            ax.plot(x, data, label='{}'.format(heading))
+
+        if self.n_wave_directions > 1:
+            ax.legend()
+
+        ax.legend()
+
+        ax.set_ylabel('Amplitude')
+        ax.set_xlabel(f'Frequency [{unit_label}]')
+
+    def plot_surface(self, what = 'amplitude', ax=None, unit=FrequencyUnit.rad_s, cmap=None):
+        """Plots amplitude or phase as a surface plot"""
+
+        if self.n_wave_directions == 1:
+            self.plot(what=what, ax=ax, unit=unit)
+
+        if cmap is None: # default colormaps
+            if what == 'amplitude':
+                cmap = 'Greys'
+            else:
+                cmap = 'hsv' # cyclic colormap
+
+        import matplotlib.pyplot as plt
+
+        if ax is None:
+            ax = plt.gca()
+
+        omega = self.omega
+        unit_label,x = unit.to_unit(omega)
+
+        headings = self._data['wave_direction'].values
+
+        data = self._data[what].values
+
+        ax.imshow(data, extent=[min(headings), max(headings), min(x), max(x) ], aspect='auto', cmap=cmap)
+
+        ax.set_xlabel('Heading [deg]')
+        ax.set_ylabel(f'Frequency [{unit_label}]')
+
+
diff --git a/tests/files/grid_t20.dhyd b/tests/files/grid_t20.dhyd
diff --git a/tests/test_plot.py b/tests/test_plot.py
@@ -0,0 +1,12 @@
+from mafredo.hyddb1 import Hyddb1, FrequencyUnit
+
+def gimme():
+    hyd = Hyddb1.create_from(r'files/grid_t20.dhyd')
+    return hyd
+
+if __name__ == '__main__':
+    h = gimme()
+    h.plot(unit=FrequencyUnit.seconds)
+
+    import matplotlib.pyplot as plt
+    plt.show()
diff --git a/tests/test_unit_conversion.py b/tests/test_unit_conversion.py
@@ -0,0 +1,24 @@
+import numpy as np
+from mafredo import FrequencyUnit
+
+def test_FrequencyUnit():
+
+    omega = 4*np.pi
+
+    unit = FrequencyUnit.Hz
+    unit_name, x = unit.to_unit(omega)
+
+    assert unit_name == 'Hz'
+    assert x == 2.0
+
+    unit = FrequencyUnit.rad_s
+    unit_name, x = unit.to_unit(omega)
+
+    assert unit_name == 'rad/s'
+    assert x == omega
+
+    unit = FrequencyUnit.seconds
+    unit_name, x = unit.to_unit(omega)
+
+    assert unit_name == 's'
+    assert x == 0.5