Repository: petrobras/ross

Issue:
Add Probe class
Currently we use list of tuples to add probes to our results.
The problem is that, contrary to other parts of ROSS, we are not able to use pint.Quantity to the probe orientation.
Instead, we have to declare the unit in a separate argument:

```
        probe : list
            List with tuples (node, orientation angle, tag).

            node : int -> Indicate the node where the probe is located.

            orientation : float -> Probe orientation angle about the shaft.
            The 0 refers to +X direction.
            The strings 'major' and 'minor' can also be used to reference the major
            and minor axis.

            tag : str, optional -> Probe tag to be add a DataFrame column title.
        probe_units : str, option
            Units for probe orientation.
            Default is "rad".
```

With a probe class we would be able to do something like:

```python
probe_la = rs.Probe(15, Q_(45, 'deg'))
```

Another possibility would be to add filter (1x etc.) to the probe.

Comments:

Patch:
diff --git a/ross/__init__.py b/ross/__init__.py
index fd0f04f..7a42d0d 100644
--- a/ross/__init__.py
+++ b/ross/__init__.py
@@ -8,6 +8,7 @@ from .defects import *
 from .disk_element import *
 from .materials import *
 from .point_mass import *
+from .probe import *
 from .results import *
 from .rotor_assembly import *
 from .shaft_element import *
diff --git a/ross/probe.py b/ross/probe.py
new file mode 100644
index 0000000..78ecb9e
--- /dev/null
+++ b/ross/probe.py
@@ -0,0 +1,49 @@
+from ross.units import Q_, check_units
+
+
+class Probe:
+    """Class of a probe.
+
+    This class will create a probe object to be used in the rotor model.
+    The probe is used to measure the response of the rotor at a specific
+    location and orientation.
+
+    Parameters
+    ----------
+    node : int
+        Indicate the node where the probe is located.
+    angle : float, pint.Quantity
+        Probe orientation angle about the shaft (rad).
+    tag : str, optional
+        A tag to name the element.
+
+    Example
+    -------
+    >>> import ross as rs
+    >>> probe1 = rs.Probe(10, Q_(45, "degree"), "Probe Drive End - X")
+    >>> probe1.node
+    10
+    >>> probe1.angle
+    0.7853981633974483
+    """
+
+    @check_units
+    def __init__(self, node, angle, tag=None):
+        self.node = node
+        self.angle = angle
+        if tag is None:
+            self.tag = f"Probe - Node {self.node}, Angle {self.angle}"
+        else:
+            self.tag = tag
+
+    @property
+    def info(self):
+        return self.node, self.angle
+
+    def __str__(self):
+        return (
+            f"Probe {self.tag}"
+            f'\n{20*"-"}'
+            f"\nNode location           : {self.node}"
+            f"\nProbe orientation angle (rad) : {self.angle}"
+        )
diff --git a/ross/results.py b/ross/results.py
index e653403..129042b 100644
--- a/ross/results.py
+++ b/ross/results.py
@@ -7,6 +7,7 @@ import inspect
 from abc import ABC
 from collections.abc import Iterable
 from pathlib import Path
+from warnings import warn
 
 import numpy as np
 import pandas as pd
@@ -2083,16 +2084,7 @@ class ForcedResponseResults(Results):
         Parameters
         ----------
         probe : list
-            List with tuples (node, orientation angle, tag).
-
-            node : int -> Indicate the node where the probe is located.
-
-            orientation : float -> Probe orientation angle about the shaft.
-            The 0 refers to +X direction.
-            The strings 'major' and 'minor' can also be used to reference the major
-            and minor axis.
-
-            tag : str, optional -> Probe tag to be add a DataFrame column title.
+            List with rs.Probe objects.
         probe_units : str, option
             Units for probe orientation.
             Default is "rad".
@@ -2135,24 +2127,39 @@ class ForcedResponseResults(Results):
         for i, p in enumerate(probe):
             amplitude = []
             for speed_idx in range(len(self.speed_range)):
+                # first try to get the angle from the probe object
                 try:
-                    angle = Q_(p[1], probe_units).to("rad").m
-                except TypeError:
-                    angle = p[1]
+                    angle = p.angle
+                    node = p.node
+                # if it is a tuple, warn the user that the use of tuples is deprecated
+                except AttributeError:
+                    try:
+                        angle = Q_(p[1], probe_units).to("rad").m
+                        warn(
+                            "The use of tuples in the probe argument is deprecated. Use the Probe class instead.",
+                            DeprecationWarning,
+                        )
+                        node = p[0]
+                    except TypeError:
+                        angle = p[1]
+                        node = p[0]
 
                 ru_e, rv_e = response[:, speed_idx][
-                    self.rotor.number_dof * p[0] : self.rotor.number_dof * p[0] + 2
+                    self.rotor.number_dof * node : self.rotor.number_dof * node + 2
                 ]
                 orbit = Orbit(
-                    node=p[0], node_pos=self.rotor.nodes_pos[p[0]], ru_e=ru_e, rv_e=rv_e
+                    node=node, node_pos=self.rotor.nodes_pos[node], ru_e=ru_e, rv_e=rv_e
                 )
                 amp, phase = orbit.calculate_amplitude(angle=angle)
                 amplitude.append(amp)
 
             try:
-                probe_tag = p[2]
-            except IndexError:
-                probe_tag = f"Probe {i+1} - Node {p[0]}"
+                probe_tag = p.tag
+            except AttributeError:
+                try:
+                    probe_tag = p[2]
+                except IndexError:
+                    probe_tag = f"Probe {i+1} - Node {p[0]}"
 
             data[probe_tag] = Q_(amplitude, base_unit).to(amplitude_units).m
 
@@ -2173,16 +2180,7 @@ class ForcedResponseResults(Results):
         Parameters
         ----------
         probe : list
-            List with tuples (node, orientation angle, tag).
-
-            node : int -> Indicate the node where the probe is located.
-
-            orientation : float -> Probe orientation angle about the shaft.
-            The 0 refers to +X direction.
-            The strings 'major' and 'minor' can also be used to reference the major
-            and minor axis.
-
-            tag : str, optional -> Probe tag to be add a DataFrame column title.
+            List with rs.Probe objects.
         probe_units : str, option
             Units for probe orientation.
             Default is "rad".
@@ -2228,24 +2226,39 @@ class ForcedResponseResults(Results):
         for i, p in enumerate(probe):
             phase_values = []
             for speed_idx in range(len(self.speed_range)):
+                # first try to get the angle from the probe object
                 try:
-                    angle = Q_(p[1], probe_units).to("rad").m
-                except TypeError:
-                    angle = p[1]
+                    angle = p.angle
+                    node = p.node
+                # if it is a tuple, warn the user that the use of tuples is deprecated
+                except AttributeError:
+                    try:
+                        angle = Q_(p[1], probe_units).to("rad").m
+                        warn(
+                            "The use of tuples in the probe argument is deprecated. Use the Probe class instead.",
+                            DeprecationWarning,
+                        )
+                        node = p[0]
+                    except TypeError:
+                        angle = p[1]
+                        node = p[0]
 
                 ru_e, rv_e = response[:, speed_idx][
-                    self.rotor.number_dof * p[0] : self.rotor.number_dof * p[0] + 2
+                    self.rotor.number_dof * node : self.rotor.number_dof * node + 2
                 ]
                 orbit = Orbit(
-                    node=p[0], node_pos=self.rotor.nodes_pos[p[0]], ru_e=ru_e, rv_e=rv_e
+                    node=node, node_pos=self.rotor.nodes_pos[node], ru_e=ru_e, rv_e=rv_e
                 )
                 amp, phase = orbit.calculate_amplitude(angle=angle)
                 phase_values.append(phase)
 
             try:
-                probe_tag = p[2]
-            except IndexError:
-                probe_tag = f"Probe {i+1} - Node {p[0]}"
+                probe_tag = p.tag
+            except AttributeError:
+                try:
+                    probe_tag = p[2]
+                except IndexError:
+                    probe_tag = f"Probe {i+1} - Node {p[0]}"
 
             data[probe_tag] = Q_(phase_values, "rad").to(phase_units).m
 
@@ -2267,17 +2280,8 @@ class ForcedResponseResults(Results):
         Parameters
         ----------
         probe : list
-            List with tuples (node, orientation angle, tag).
-
-            node : int -> Indicate the node where the probe is located.
-
-            orientation : float -> Probe orientation angle about the shaft.
-            The 0 refers to +X direction.
-            The strings 'major' and 'minor' can also be used to reference the major
-            and minor axis.
-
-            tag : str, optional -> Probe tag to be add a DataFrame column title.
-        probe_units : str, option
+            List with rs.Probe objects.
+        probe_units : str, optional
             Units for probe orientation.
             Default is "rad".
         frequency_units : str, optional
@@ -2352,17 +2356,8 @@ class ForcedResponseResults(Results):
         Parameters
         ----------
         probe : list
-            List with tuples (node, orientation angle, tag).
-
-            node : int -> Indicate the node where the probe is located.
-
-            orientation : float -> Probe orientation angle about the shaft.
-            The 0 refers to +X direction.
-            The strings 'major' and 'minor' can also be used to reference the major
-            and minor axis.
-
-            tag : str, optional -> Probe tag to be add a DataFrame column title.
-        probe_units : str, option
+            List with rs.Probe objects.
+        probe_units : str, optional
             Units for probe orientation.
             Default is "rad".
         frequency_units : str, optional
@@ -2440,17 +2435,8 @@ class ForcedResponseResults(Results):
         Parameters
         ----------
         probe : list
-            List with tuples (node, orientation angle, tag).
-
-            node : int -> Indicate the node where the probe is located.
-
-            orientation : float -> Probe orientation angle about the shaft.
-            The 0 refers to +X direction.
-            The strings 'major' and 'minor' can also be used to reference the major
-            and minor axis.
-
-            tag : str, optional -> Probe tag to be add a DataFrame column title.
-        probe_units : str, option
+            List with rs.Probe objects.
+        probe_units : str, optional
             Units for probe orientation.
             Default is "rad".
         frequency_units : str, optional
@@ -2548,17 +2534,8 @@ class ForcedResponseResults(Results):
         Parameters
         ----------
         probe : list
-            List with tuples (node, orientation angle, tag).
-
-            node : int -> Indicate the node where the probe is located.
-
-            orientation : float -> Probe orientation angle about the shaft.
-            The 0 refers to +X direction.
-            The strings 'major' and 'minor' can also be used to reference the major
-            and minor axis.
-
-            tag : str, optional -> Probe tag to be add a DataFrame column title.
-        probe_units : str, option
+            List with rs.Probe objects.
+        probe_units : str, optional
             Units for probe orientation.
             Default is "rad".
         frequency_units : str, optional
@@ -4183,16 +4160,7 @@ class TimeResponseResults(Results):
         Parameters
         ----------
         probe : list
-            List with tuples (node, orientation angle, tag).
-
-            node : int -> Indicate the node where the probe is located.
-
-            orientation : float -> Probe orientation angle about the shaft.
-            The 0 refers to +X direction.
-            The strings 'major' and 'minor' can also be used to reference the major
-            and minor axis.
-
-            tag : str, optional -> Probe tag to be add a DataFrame column title.
+            List with rs.Probe objects.
         probe_units : str, option
             Units for probe orientation.
             Default is "rad".


Test patch:
diff --git a/ross/tests/test_probe.py b/ross/tests/test_probe.py
new file mode 100644
index 0000000..f8028c2
--- /dev/null
+++ b/ross/tests/test_probe.py
@@ -0,0 +1,21 @@
+import pytest
+from numpy.testing import assert_allclose
+
+from ross.probe import Probe
+from ross.units import Q_
+
+
+@pytest.fixture
+def probe():
+    return Probe(10, Q_(45, "degree"), "V1")
+
+
+def test_parameters(probe):
+    assert_allclose(probe.node, 10)
+    assert_allclose(probe.angle, 0.7853981633974483)
+
+
+def test_info(probe):
+    node, angle = probe.info
+    assert_allclose(node, 10)
+    assert_allclose(angle, 0.7853981633974483)
diff --git a/ross/tests/test_rotor_assembly.py b/ross/tests/test_rotor_assembly.py
index 893c803..5796410 100644
--- a/ross/tests/test_rotor_assembly.py
+++ b/ross/tests/test_rotor_assembly.py
@@ -10,6 +10,7 @@ from ross.bearing_seal_element import *
 from ross.disk_element import *
 from ross.materials import Material, steel
 from ross.point_mass import *
+from ross.probe import Probe
 from ross.rotor_assembly import *
 from ross.shaft_element import *
 from ross.units import Q_
@@ -1434,17 +1435,33 @@ def test_plot_mode(rotor7):
     assert_allclose(fig.data[-3]["z"][:5], expected_z, rtol=1e-5)
 
 
-def test_unbalance(rotor7):
-    unb = rotor7.run_unbalance_response(
+def test_unbalance(rotor3):
+    unb = rotor3.run_unbalance_response(
         node=0, unbalance_magnitude=1, unbalance_phase=0, frequency=[50, 100]
     )
-    amplitude_expected = np.array([0.00274, 0.003526])
+    amplitude_expected = np.array([0.003065, 0.004169])
     data = unb.data_magnitude(probe=[(0, 45)], probe_units="deg")
     assert_allclose(data["Probe 1 - Node 0"], amplitude_expected, rtol=1e-4)
+    data = unb.data_magnitude(probe=[Probe(0, Q_(45, "deg"), tag="Probe 1 - Node 0")])
+    assert_allclose(data["Probe 1 - Node 0"], amplitude_expected, rtol=1e-4)
 
-    phase_expected = np.array([0.730209, 0.545276])
+    phase_expected = np.array([0.785398, 0.785398])
     data = unb.data_phase(probe=[(0, 45)], probe_units="deg")
     assert_allclose(data["Probe 1 - Node 0"], phase_expected, rtol=1e-4)
+    data = unb.data_phase(probe=[Probe(0, Q_(45, "deg"), tag="Probe 1 - Node 0")])
+    assert_allclose(data["Probe 1 - Node 0"], phase_expected, rtol=1e-4)
+
+    amplitude_expected = np.array([0.003526, 0.005518])
+    data = unb.data_magnitude(probe=[(0, "major")])
+    assert_allclose(data["Probe 1 - Node 0"], amplitude_expected, rtol=1e-4)
+    data = unb.data_magnitude(probe=[Probe(0, "major", tag="Probe 1 - Node 0")])
+    assert_allclose(data["Probe 1 - Node 0"], amplitude_expected, rtol=1e-4)
+
+    phase_expected = np.array([1.5742, 1.573571])
+    data = unb.data_phase(probe=[(0, "major")], probe_units="deg")
+    assert_allclose(data["Probe 1 - Node 0"], phase_expected, rtol=1e-4)
+    data = unb.data_phase(probe=[Probe(0, "major", tag="Probe 1 - Node 0")])
+    assert_allclose(data["Probe 1 - Node 0"], phase_expected, rtol=1e-4)
 
 
 def test_deflected_shape(rotor7):