Merge pull request #44 from sot/chunks

Add a linspace class method

.pre-commit-config.yaml

@@ -1,7 +1,7 @@
 repos:
 - repo: https://github.com/astral-sh/ruff-pre-commit
   # Ruff version.
-  rev: v0.1.5
+  rev: v0.6.7
   hooks:
     # Run the linter.
     - id: ruff

cxotime/cxotime.py

@@ -4,6 +4,7 @@
 from copy import copy
 from typing import Union
 
+import astropy.units as u
 import erfa
 import numpy as np
 import numpy.typing as npt
@@ -174,6 +175,58 @@ def __init__(self, *args, **kwargs):
 
         super(CxoTime, self).__init__(*args, **kwargs)
 
+    @classmethod
+    def linspace(
+        cls,
+        start: CxoTimeLike,
+        stop: CxoTimeLike,
+        num: int | None = None,
+        step_max: u.Quantity | None = None,
+    ):
+        """
+        Get a uniform time series that covers the given time range.
+
+        Output times either divide the time range into ``num`` intervals or are
+        uniformly spaced by up to ``step_max``, and cover the time
+        range from ``start`` to ``stop``.
+
+        Parameters
+        ----------
+        start : CxoTimeLike
+            Start time of the time range.
+        stop : CxoTimeLike
+            Stop time of the time range.
+        num : int | None
+            Number of time bins.
+        step_max : u.Quantity (timelike)
+            Maximum time interval step..  Should be positive nonzero.
+
+        Returns
+        -------
+        CxoTime
+            CxoTime with time bin edges for each interval.
+        """
+        start = CxoTime(start)
+        stop = CxoTime(stop)
+
+        if (num is None) == (step_max is None):
+            raise ValueError("exactly one of num and step_max must be defined")
+
+        if step_max is not None:
+            # Require that step_max is positive nonzero
+            if step_max <= 0 * u.s:
+                raise ValueError("step_max must be positive nonzero")
+
+            # Calculate chunks to cover time range, handling edge case of start == stop
+            num = int(max(np.ceil(abs(float((stop - start) / step_max))), 1))
+
+        if num <= 0:
+            raise ValueError("num must be positive nonzero int")
+
+        times = np.linspace(start, stop, num + 1)
+
+        return times
+
     @classmethod
     def now(cls):
         return cls()

cxotime/tests/test_cxotime.py

@@ -3,6 +3,7 @@
 Simple test of CxoTime.  The base Time object is extremely well
 tested, so this simply confirms that the add-on in CxoTime works.
 """
+
 import io
 import time
 from dataclasses import dataclass

cxotime/tests/test_cxotime_linspace.py

@@ -0,0 +1,207 @@
+import astropy.units as u
+import numpy as np
+import pytest
+
+from cxotime import CxoTime
+
+
+def run_linspace_step_test(start, stop, dt_max, expected_len, expected_values):
+    """General test function for CxoTime.linspace with step_max."""
+
+    result = CxoTime.linspace(start, stop, step_max=dt_max)
+
+    # Confirm that the first interval duration matches the expected value
+    assert abs(result[1] - result[0]) <= min(
+        dt_max, abs(CxoTime(stop) - CxoTime(start))
+    )
+
+    # Confirm that all the intervals are the same duration
+    interval1 = result[1] - result[0]
+    assert all(np.isclose((result[1:] - result[:-1]).sec, interval1.sec))
+
+    # Confirm that the time range is covered
+    assert result[0] == start
+    assert result[-1] == stop
+
+    # And confirm that the result is as expected
+    assert len(result) == expected_len
+    assert np.allclose(CxoTime(result).secs, CxoTime(expected_values).secs)
+
+
+def test_linspace_step_with_zero_range():
+    """Test that the result is correct when start==stop."""
+    run_linspace_step_test(
+        "2000:001", "2000:001", 1 * u.day, 2, ["2000:001", "2000:001"]
+    )
+
+
+def test_linspace_step_with_zero_range_and_bigger_step():
+    """Test that the result is correct when the step is larger than the range."""
+    run_linspace_step_test(
+        "2000:001", "2000:001", 1.5 * u.day, 2, ["2000:001", "2000:001"]
+    )
+
+
+def test_linspace_step_with_float_range():
+    """Test that the result is correct when the step is smaller than the range and more float-like."""
+    run_linspace_step_test(
+        "2024:001",
+        "2023:364",
+        12.5 * u.hour,
+        5,
+        [
+            "2024:001:00:00:00.000",
+            "2023:365:12:00:00.000",
+            "2023:365:00:00:00.000",
+            "2023:364:12:00:00.000",
+            "2023:364:00:00:00.000",
+        ],
+    )
+
+
+def test_linspace_step_odd_minutes():
+    """Test that the result is correct when the step is just some weird float of minutes."""
+    run_linspace_step_test(
+        "2020:020:00:12:00.000",
+        "2020:020:00:13:00.000",
+        23.5 * u.min,
+        2,
+        ["2020:020:00:12:00.000", "2020:020:00:13:00.000"],
+    )
+
+
+def test_linspace_negative_range():
+    """Test that the result is correct when stop < start"""
+    start = CxoTime("2000:005")
+    stop = CxoTime("2000:001")
+    dt_max = 24 * u.hour
+    result = CxoTime.linspace(start, stop, step_max=dt_max)
+    assert len(result) == 5
+    expected_values = ["2000:005", "2000:004", "2000:003", "2000:002", "2000:001"]
+    assert np.all(result == CxoTime(expected_values))
+
+
+def test_linspace_big_step():
+    """Test that the result is correct when the step is larger than the range."""
+    start = CxoTime("2020:001")
+    stop = CxoTime("2020:005")
+    dt_max = 30 * u.day
+    result = CxoTime.linspace(start, stop, step_max=dt_max)
+    assert len(result) == 2
+    assert np.all(result == CxoTime(["2020:001", "2020:005"]))
+
+
+def test_linspace_zero_step():
+    """Test that an error is raised if step_max is zero."""
+    start = CxoTime("2020:001")
+    stop = CxoTime("2020:005")
+    dt_max = 0 * u.day
+    with pytest.raises(ValueError, match="step_max must be positive nonzero"):
+        CxoTime.linspace(start, stop, step_max=dt_max)
+
+
+def test_linspace_negative_step():
+    """Test that an error is raised if step_max is negative."""
+    start = CxoTime("2020:001")
+    stop = CxoTime("2020:005")
+    dt_max = -1 * u.day
+    with pytest.raises(ValueError, match="step_max must be positive nonzero"):
+        CxoTime.linspace(start, stop, step_max=dt_max)
+
+
+def test_linspace_num_0():
+    """Test that an error is raised if num is zero."""
+    start = CxoTime("2000:001")
+    stop = CxoTime("2000:005")
+    num = 0
+    with pytest.raises(ValueError, match="num must be positive nonzero"):
+        CxoTime.linspace(start, stop, num=num)
+
+
+def test_linspace_num_neg():
+    """Test that an error is raised if num is negative."""
+    start = CxoTime("2000:001")
+    stop = CxoTime("2000:005")
+    num = -1
+    with pytest.raises(ValueError, match="num must be positive nonzero"):
+        CxoTime.linspace(start, stop, num=num)
+
+
+def test_linspace_num_1():
+    start = CxoTime("2000:001")
+    stop = CxoTime("2000:005")
+    num = 2
+    result = CxoTime.linspace(start, stop, num=num)
+    assert len(result) == num + 1
+    expected = ["2000:001", "2000:003", "2000:005"]
+    assert np.all(result == CxoTime(expected))
+
+
+def test_linspace_num_2():
+    start = "2000:001"
+    stop = "2024:001"
+    num = 12
+    result = CxoTime.linspace(start, stop, num=num)
+    assert len(result) == num + 1
+    expected = [
+        "2000:001:00:00:00.000",
+        "2001:365:12:00:00.417",
+        "2004:001:00:00:00.833",
+        "2005:365:12:00:01.250",
+        "2008:001:00:00:00.667",
+        "2009:365:12:00:00.083",
+        "2012:001:00:00:00.500",
+        "2013:365:11:59:59.917",
+        "2015:365:23:59:59.333",
+        "2017:365:11:59:58.750",
+        "2019:365:23:59:59.167",
+        "2021:365:11:59:59.583",
+        "2024:001:00:00:00.000",
+    ]
+    # There are very small numerical differences between the expected and actual values
+    # so this test uses allclose instead of ==.
+    assert np.allclose(CxoTime(result).secs, CxoTime(expected).secs)
+
+
+def test_linspace_num_3():
+    start = "2010:001"
+    stop = "2011:001"
+    num = 12
+    result = CxoTime.linspace(start, stop, num=num)
+    assert len(result) == num + 1
+    expected = [
+        "2010:001:00:00:00.000",
+        "2010:031:10:00:00.000",
+        "2010:061:20:00:00.000",
+        "2010:092:06:00:00.000",
+        "2010:122:16:00:00.000",
+        "2010:153:02:00:00.000",
+        "2010:183:12:00:00.000",
+        "2010:213:22:00:00.000",
+        "2010:244:08:00:00.000",
+        "2010:274:18:00:00.000",
+        "2010:305:04:00:00.000",
+        "2010:335:14:00:00.000",
+        "2011:001:00:00:00.000",
+    ]
+    # There are very small numerical differences between the expected and actual values
+    # so this test uses allclose instead of ==.
+    assert np.allclose(CxoTime(result).secs, CxoTime(expected).secs)
+
+
+def test_missing_args():
+    start = "2015:001"
+    stop = "2015:002"
+    with pytest.raises(
+        ValueError, match="exactly one of num and step_max must be defined"
+    ):
+        CxoTime.linspace(start, stop)
+
+
+def test_too_many_args():
+    start = "2015:001"
+    stop = "2015:002"
+    with pytest.raises(
+        ValueError, match="exactly one of num and step_max must be defined"
+    ):
+        CxoTime.linspace(start, stop, 1, 2)

ruff-base.toml

@@ -0,0 +1,59 @@
+# Copied originally from pandas. This config requires ruff >= 0.2.
+target-version = "py311"
+
+# fix = true
+lint.unfixable = []
+
+lint.select = [
+  "I", # isort
+  "F", # pyflakes
+  "E", "W", # pycodestyle
+  "YTT", # flake8-2020
+  "B", # flake8-bugbear
+  "Q", # flake8-quotes
+  "T10", # flake8-debugger
+  "INT", # flake8-gettext
+  "PLC", "PLE", "PLR", "PLW", # pylint
+  "PIE", # misc lints
+  "PYI", # flake8-pyi
+  "TID", # tidy imports
+  "ISC", # implicit string concatenation
+  "TCH", # type-checking imports
+  "C4", # comprehensions
+  "PGH" # pygrep-hooks
+]
+
+# Some additional rules that are useful
+lint.extend-select = [
+"UP009",  # UTF-8 encoding declaration is unnecessary
+"SIM118",  # Use `key in dict` instead of `key in dict.keys()`
+"D205",  # One blank line required between summary line and description
+"ARG001",  # Unused function argument
+"RSE102",  # Unnecessary parentheses on raised exception
+"PERF401",  # Use a list comprehension to create a transformed list
+]
+
+lint.ignore = [
+  "ISC001", # Disable this for compatibility with ruff format
+  "E402", # module level import not at top of file
+  "E731", # do not assign a lambda expression, use a def
+  "PLR2004", # Magic number
+  "B028", # No explicit `stacklevel` keyword argument found
+  "PLR0913", # Too many arguments to function call
+  "PLR1730", # Checks for if statements that can be replaced with min() or max() calls
+]
+
+extend-exclude = [
+  "docs",
+]
+
+[lint.pycodestyle]
+max-line-length = 100 # E501 reports lines that exceed the length of 100.
+
+[lint.extend-per-file-ignores]
+"__init__.py" = ["E402", "F401", "F403"]
+# For tests:
+# - D205: Don't worry about test docstrings
+# - ARG001: Unused function argument false positives for some fixtures
+# - E501: Line-too-long
+"**/tests/test_*.py" = ["D205", "ARG001", "E501"]