From e2cc41609162801f9eae8c3788d624d3c6ad01ae Mon Sep 17 00:00:00 2001
From: Derrick Chambers <d-chambers@users.noreply.github.com>
Date: Wed, 23 Aug 2023 18:05:08 -0600
Subject: [PATCH] add get_sample_count and tests (#228)

---
 dascore/core/coords.py         | 20 ++++++++++++++++++++
 tests/test_core/test_coords.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/dascore/core/coords.py b/dascore/core/coords.py
index 34c5aeb6..d120013a 100644
--- a/dascore/core/coords.py
+++ b/dascore/core/coords.py
@@ -528,6 +528,26 @@ def update(self, **kwargs):
             out = out.convert_units(units)
         return out
 
+    def get_sample_count(self, value) -> int:
+        """
+        Return the number of samples represented by a value.
+
+        This is calculated by dividing the value by dt and rounding up.
+        Therefore, the output will always be greater or equal to 1.
+
+        Parameters
+        ----------
+        value
+            The value (supports units).
+        """
+        if not self.evenly_sampled:
+            msg = "Coordinate is not evenly sampled, cant get sample count."
+            raise CoordError(msg)
+        compat_val = self._get_compatible_value(value, relative=True)
+        duration = compat_val - self.min()
+        samples = int(np.ceil(duration / self.step))
+        return samples
+
 
 class CoordRange(BaseCoord):
     """
diff --git a/tests/test_core/test_coords.py b/tests/test_core/test_coords.py
index 6cc6b9c8..1a9d800b 100644
--- a/tests/test_core/test_coords.py
+++ b/tests/test_core/test_coords.py
@@ -1081,6 +1081,34 @@ def test_time_delta(self, evenly_sampled_time_delta_coord):
         assert out.dtype == coord.dtype
 
 
+class TestGetSampleCount:
+    """Tests for getting the number of samples from a coordinate."""
+
+    def test_uneven_raises(self, monotonic_float_coord):
+        """Unevenly sampled counts should raise."""
+        msg = "Coordinate is not evenly sampled"
+        with pytest.raises(CoordError, match=msg):
+            monotonic_float_coord.get_sample_count(10)
+
+    def test_dt_lt_step(self, evenly_sampled_coord):
+        """Ensure sample count is 1 when value < dt"""
+        dt = evenly_sampled_coord.step * 0.1
+        out = evenly_sampled_coord.get_sample_count(dt)
+        assert out == 1
+
+    def test_datetime(self, evenly_sampled_date_coord):
+        """Ensure evenly sampled date coords work."""
+        dt = evenly_sampled_date_coord.step
+        out = evenly_sampled_date_coord.get_sample_count(12 * dt)
+        assert out == 12
+
+    def test_timedelta(self, evenly_sampled_time_delta_coord):
+        """Ensure timedelta works."""
+        dt = evenly_sampled_time_delta_coord.step
+        out = evenly_sampled_time_delta_coord.get_sample_count(12 * dt)
+        assert out == 12
+
+
 class TestIssues:
     """Tests for special issues related to coords."""