Repository: MideTechnology/endaq-python (language: python)

Issue:
Return the PVSS of enveloped half sine
Current function returns the pulse width and amplitude, but it would be nice to also give the option of returning the PVSS itself and the time series of the pulse.


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 3dcd127..1cc1e8e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -100,6 +100,7 @@ instance/
 
 # Sphinx documentation
 docs/_build/
+docs/_doctest/
 
 # PyBuilder
 .pybuilder/
diff --git a/endaq/calc/shock.py b/endaq/calc/shock.py
index c4841e0..5934a15 100644
--- a/endaq/calc/shock.py
+++ b/endaq/calc/shock.py
@@ -3,8 +3,9 @@
 from __future__ import annotations
 
 import typing
-from typing import Tuple
+from typing import Optional
 from collections import namedtuple
+from dataclasses import dataclass
 import functools
 import warnings
 
@@ -232,10 +233,117 @@ def shock_spectrum(
     )
 
 
+@dataclass
+class HalfSineWavePulse:
+    """
+    The output data type for ``enveloping_half_sine``.
+
+    The significant data members are `amplitude` and `duration`, which can
+    simply be unpacked as if from a plain tuple:
+
+    .. testsetup::
+
+        import pandas as pd
+        df_pvss = pd.DataFrame([1, 1], index=[200, 400])
+
+        from endaq.calc.shock import enveloping_half_sine
+
+    .. testcode::
+
+        ampl, T = enveloping_half_sine(df_pvss)
+
+    However, users can also elect to use the other methods of this class to
+    generate other kinds of outputs.
+
+    .. note:: This class is not intended to be instantiated manually.
+    """
+
+    amplitude: pd.Series
+    duration: pd.Series
+
+    def __iter__(self):
+        return iter((self.amplitude, self.duration))
+
+    def to_time_series(
+        self,
+        tstart: Optional[float] = None,
+        tstop: Optional[float] = None,
+        dt: Optional[float] = None,
+        tpulse: Optional[float] = None,
+    ) -> pd.DataFrame:
+        """
+        Generate a time-series of the half-sine pulse.
+
+        :param tstart: the starting time of the resulting waveform; if `None`
+            (default), the range starts at `tpulse`
+        :param tstop: the ending time of the resulting waveform; if `None`
+            (default), the range ends at `tpulse + duration`
+        :param dt: the sampling period of the resulting waveform; defaults to
+            1/20th of the pulse duration
+        :param tpulse: the starting time of the pulse within the resulting
+            waveform; if `None` (default), the pulse starts at either:
+
+            - ``tstart``, if provided
+            - ``tstop - self.duration.max())``, if `tstop` is provided
+            - ``0.0`` otherwise
+        :return: a time-series of the half-sine pulse
+        """
+        if dt is None:
+            dt = self.duration.min() / 20
+        if dt > self.duration.min() / 8:
+            warnings.warn(
+                f"the sampling period {dt} is large relative to the pulse duration"
+                f" {self.duration.min()}; the waveform may not accurately represent"
+                f" the half-sine pulse's shock intensity"
+            )
+
+        default_start = 0.0
+        if tstop is not None:
+            default_start = tstop - self.duration.max()
+
+        if tpulse is None and tstart is None:
+            tpulse = tstart = default_start
+        elif tpulse is None:
+            tpulse = tstart
+        elif tstart is None:
+            tstart = tpulse
+
+        if tstop is None:
+            tstop = tpulse + self.duration.max()
+
+        if not (tstart <= tpulse <= tstop - self.duration.max()):
+            warnings.warn(
+                "half-sine pulse extends beyond the bounds of the time series"
+            )
+
+        t = np.arange(tstart, tstop, dt)
+
+        data = np.zeros((len(t), len(self.amplitude)), dtype=float)
+        t_data, ampl_data, T_data = np.broadcast_arrays(
+            t[..., None], self.amplitude.to_numpy(), self.duration.to_numpy()
+        )
+        t_mask = np.nonzero((t_data >= tpulse) & (t_data < tpulse + T_data))
+        data[t_mask] = ampl_data[t_mask] * np.sin(
+            np.pi * t_data[t_mask] / T_data[t_mask]
+        )
+
+        return pd.DataFrame(
+            data,
+            index=pd.Series(t, name="timestamp"),
+            columns=self.amplitude.index,
+        )
+
+    # def widened_duration(self, new_duration: float):
+    #    pass
+
+    # def pseudo_velocity(self):
+    #    pass
+
+
 def enveloping_half_sine(
     pvss: pd.DataFrame,
     damp: float = 0.0,
-) -> Tuple[pd.Series, pd.Series]:
+) -> HalfSineWavePulse:
     """
     Characterize a half-sine pulse whose PVSS envelopes the input.
 
@@ -274,7 +382,7 @@ def amp_factor(damp):
     max_pvss = pvss.max()
     max_f_pvss = pvss.mul(pvss.index, axis=0).max()
 
-    return namedtuple("HalfSinePulseParameters", "amplitude, period")(
+    return HalfSineWavePulse(
         amplitude=2 * np.pi * max_f_pvss,
-        period=max_pvss / (4 * amp_factor(damp) * max_f_pvss),
+        duration=max_pvss / (4 * amp_factor(damp) * max_f_pvss),
     )


Test patch:
diff --git a/tests/calc/test_shock.py b/tests/calc/test_shock.py
index d4df33e..02833ab 100644
--- a/tests/calc/test_shock.py
+++ b/tests/calc/test_shock.py
@@ -1,5 +1,8 @@
 # -*- coding: utf-8 -*-
 
+from unittest import mock
+import warnings
+
 import pytest
 import hypothesis as hyp
 import hypothesis.strategies as hyp_st
@@ -246,25 +249,57 @@ def test_pseudo_velocity_zero_padding(
     ).map(lambda array: pd.DataFrame(array, index=np.arange(1, 41))),
     damp=hyp_st.floats(0, 0.2),
 )
-@hyp.settings(deadline=None)  # this test tends to timeout
 def test_enveloping_half_sine(df_pvss, damp):
-    ampl, T = shock.enveloping_half_sine(df_pvss, damp=damp)
-    hyp.note(f"pulse amplitude: {ampl}")
-    hyp.note(f"pulse duration: {T}")
-
-    ampl = ampl[0]
-    T = T[0]
-
-    dt = min(
-        1 / (2 * df_pvss.index[-1]), T / 20
-    )  # guarantee sufficient # of samples to represent pulse
-    fs = 1 / dt
-    times = np.arange(int(fs * (T + 1 / df_pvss.index[0]))) / fs
-    pulse = np.zeros_like(times)
-    pulse[: int(T * fs)] = ampl * np.sin((np.pi / T) * times[: int(T * fs)])
+    env_half_sine = shock.enveloping_half_sine(df_pvss, damp=damp)
+    hyp.note(f"pulse amplitude: {env_half_sine.amplitude}")
+    hyp.note(f"pulse duration: {env_half_sine.duration}")
+
+    pulse = env_half_sine.to_time_series()
     pulse_pvss = shock.shock_spectrum(
-        pd.DataFrame(pulse, index=times), freqs=df_pvss.index, damp=damp, mode="pvss"
+        pulse, freqs=df_pvss.index, damp=damp, mode="pvss"
     )
 
     # This is an approximation -> give the result a fudge-factor for correctness
     assert (df_pvss / pulse_pvss).max().max() < 1.2
+
+
+class TestHalfSineWavePulse:
+    @pytest.mark.parametrize(
+        "tstart, tstop, dt, tpulse, warning_type",
+        [
+            # dt warnings
+            (None, None, 0.12, 0, None),  # dt < duration / 8 => OK
+            (None, None, 0.13, 0, UserWarning),  # dt > duration / 8 => WARNING
+            # trange warnings
+            (None, 0.5, None, 0, UserWarning),  # trange[1] < t0 + duration => WARNING
+            (0.5, None, None, 0, UserWarning),  # trange[0] > t0 => WARNING
+            (0.5, None, None, 1, None),  # OK
+        ],
+    )
+    def test_to_time_series_warnings(self, tstart, tstop, dt, tpulse, warning_type):
+        env_half_sine = shock.HalfSineWavePulse(
+            amplitude=pd.Series([1]),
+            duration=pd.Series([1]),
+        )
+
+        if warning_type is None:
+            with warnings.catch_warnings():
+                warnings.simplefilter("error")
+                env_half_sine.to_time_series(
+                    tstart=tstart, tstop=tstop, dt=dt, tpulse=tpulse
+                )
+        else:
+            with pytest.warns(warning_type):
+                env_half_sine.to_time_series(
+                    tstart=tstart, tstop=tstop, dt=dt, tpulse=tpulse
+                )
+
+    def test_tuple_like(self):
+        env_half_sine = shock.HalfSineWavePulse(
+            amplitude=mock.sentinel.amplitude,
+            duration=mock.sentinel.duration,
+        )
+
+        ampl, T = env_half_sine
+        assert ampl == mock.sentinel.amplitude
+        assert T == mock.sentinel.duration