Provide 'side' to pick whether we want the closest, left, or right value

Author: cphyc

yt/data_objects/time_series.py

@@ -6,7 +6,7 @@
 import weakref
 from abc import ABC, abstractmethod
 from functools import wraps
-from typing import Optional, Union
+from typing import Literal, Optional, Union
 
 import numpy as np
 from more_itertools import always_iterable
@@ -445,6 +445,7 @@ def _get_by_attribute(
         attribute: str,
         value: Union[unyt_quantity, tuple[float, str]],
         tolerance: Union[None, unyt_quantity, tuple[float, str]] = None,
+        side: Union[Literal["nearest"], Literal["left"], Literal["right"]] = "nearest",
     ) -> "Dataset":
         r"""
         Get a dataset at or near to a given value.
@@ -462,8 +463,14 @@ def _get_by_attribute(
             within the tolerance value. If None, simply return the
             nearest dataset.
             Default: None.
+        side : str
+            The side of the value to return. Can be 'nearest', 'left' or 'right'.
+            Default: 'nearest'.
         """
 
+        if side not in ("nearest", "left", "right"):
+            raise ValueError(f"side must be 'nearest', 'left' or 'right', got {side}.")
+
         # Use a binary search to find the closest value
         iL = 0
         iH = len(self._pre_outputs) - 1
@@ -518,7 +525,13 @@ def _get_by_attribute(
                 dsL = dsH = dsM
                 break
 
-        if abs(value - getattr(dsL, attribute)) < abs(value - getattr(dsH, attribute)):
+        if side == "left":
+            ds_best = dsL
+        elif side == "right":
+            ds_best = dsH
+        elif abs(value - getattr(dsL, attribute)) < abs(
+            value - getattr(dsH, attribute)
+        ):
             ds_best = dsL
         else:
             ds_best = dsH
@@ -534,6 +547,7 @@ def get_by_time(
         self,
         time: Union[unyt_quantity, tuple],
         tolerance: Union[None, unyt_quantity, tuple] = None,
+        side: Union[Literal["nearest"], Literal["left"], Literal["right"]] = "nearest",
     ):
         """
         Get a dataset at or near to a given time.
@@ -547,16 +561,26 @@ def get_by_time(
             within the tolerance value. If None, simply return the
             nearest dataset.
             Default: None.
+        side : str
+            The side of the value to return. Can be 'nearest', 'left' or 'right'.
+            Default: 'nearest'.
 
         Examples
         --------
         >>> ds = ts.get_by_time((12, "Gyr"))
         >>> t = ts[0].quan(12, "Gyr")
         ... ds = ts.get_by_time(t, tolerance=(100, "Myr"))
         """
-        return self._get_by_attribute("current_time", time, tolerance=tolerance)
+        return self._get_by_attribute(
+            "current_time", time, tolerance=tolerance, side=side
+        )
 
-    def get_by_redshift(self, redshift: float, tolerance: Optional[float] = None):
+    def get_by_redshift(
+        self,
+        redshift: float,
+        tolerance: Optional[float] = None,
+        side: Union[Literal["nearest"], Literal["left"], Literal["right"]] = "nearest",
+    ):
         """
         Get a dataset at or near to a given time.
 
@@ -569,6 +593,9 @@ def get_by_redshift(self, redshift: float, tolerance: Optional[float] = None):
             within the tolerance value. If None, simply return the
             nearest dataset.
             Default: None.
+        side : str
+            The side of the value to return. Can be 'nearest', 'left' or 'right'.
+            Default: 'nearest'.
 
         Examples
         --------