holobench 1.41.0__py3-none-any.whl → 1.43.0__py3-none-any.whl

bencher/utils.py

@@ -41,16 +41,22 @@ def make_namedtuple(class_name: str, **fields) -> namedtuple:
     return namedtuple(class_name, fields)(*fields.values())
 
 
-def get_nearest_coords(dataset: xr.Dataset, collapse_list=False, **kwargs) -> dict:
-    """Given an xarray dataset and kwargs of key value pairs of coordinate values, return a dictionary of the nearest coordinate name value pair that was found in the dataset
+def get_nearest_coords(dataset: xr.Dataset, collapse_list: bool = False, **kwargs) -> dict:
+    """Find the nearest coordinates in an xarray dataset based on provided coordinate values.
+
+    Given an xarray dataset and kwargs of key-value pairs of coordinate values, return a dictionary
+    of the nearest coordinate name-value pair that was found in the dataset.
 
     Args:
-        ds (xr.Dataset): dataset
+        dataset (xr.Dataset): The xarray dataset to search in
+        collapse_list (bool, optional): If True, when a coordinate value is a list, only the first
+                                       item is returned. Defaults to False.
+        **kwargs: Key-value pairs where keys are coordinate names and values are points to find
+                 the nearest match for
 
     Returns:
-        dict: nearest coordinate name value pair that matches the input coordinate name value pairs.
+        dict: Dictionary of coordinate name-value pairs with the nearest values found in the dataset
     """
-
     selection = dataset.sel(method="nearest", **kwargs)
     cd = selection.coords.to_dataset().to_dict()["coords"]
     cd2 = {}
@@ -61,7 +67,19 @@ def get_nearest_coords(dataset: xr.Dataset, collapse_list=False, **kwargs) -> di
     return cd2
 
 
-def get_nearest_coords1D(val: Any, coords) -> Any:
+def get_nearest_coords1D(val: Any, coords: List[Any]) -> Any:
+    """Find the closest coordinate to a given value in a list of coordinates.
+
+    For numeric values, finds the value in coords that is closest to val.
+    For non-numeric values, returns the exact match if found, otherwise returns val.
+
+    Args:
+        val (Any): The value to find the closest coordinate for
+        coords (List[Any]): The list of coordinates to search in
+
+    Returns:
+        Any: The closest coordinate value from the list
+    """
     if isinstance(val, (int, float)):
         return min(coords, key=lambda x_: abs(x_ - val))
     for i in coords:
@@ -70,19 +88,28 @@ def get_nearest_coords1D(val: Any, coords) -> Any:
     return val
 
 
-def hash_sha1(var: any) -> str:
-    """A hash function that avoids the PYTHONHASHSEED 'feature' which returns a different hash value each time the program is run"""
+def hash_sha1(var: Any) -> str:
+    """A hash function that avoids the PYTHONHASHSEED 'feature' which returns a different hash value each time the program is run.
+
+    Converts input to a consistent SHA1 hash string.
+
+    Args:
+        var (Any): The variable to hash
+
+    Returns:
+        str: A hexadecimal SHA1 hash of the string representation of the variable
+    """
     return hashlib.sha1(str(var).encode("ASCII")).hexdigest()
 
 
-def capitalise_words(message: str):
-    """Given a string of lowercase words, capitalise them
+def capitalise_words(message: str) -> str:
+    """Given a string of lowercase words, capitalise them.
 
     Args:
         message (str): lower case string
 
     Returns:
-        _type_: capitalised string
+        str: capitalised string where each word starts with an uppercase letter
     """
     capitalized_message = " ".join([word.capitalize() for word in message.split(" ")])
     return capitalized_message
@@ -101,7 +128,16 @@ def un_camel(camel: str) -> str:
     return capitalise_words(re.sub("([a-z])([A-Z])", r"\g<1> \g<2>", camel.replace("_", " ")))
 
 
-def mult_tuple(inp: Tuple[float], val: float) -> Tuple[float]:
+def mult_tuple(inp: Tuple[float, ...], val: float) -> Tuple[float, ...]:
+    """Multiply each element in a tuple by a scalar value.
+
+    Args:
+        inp (Tuple[float, ...]): The input tuple of floats to multiply
+        val (float): The scalar value to multiply each element by
+
+    Returns:
+        Tuple[float, ...]: A new tuple with each element multiplied by val
+    """
     return tuple(np.array(inp) * val)
 
 
@@ -118,15 +154,17 @@ def tabs_in_markdown(regular_str: str, spaces: int = 2) -> str:
     return regular_str.replace("\t", "".join(["&nbsp;"] * spaces))
 
 
-def int_to_col(int_val, sat=0.5, val=0.95, alpha=-1) -> tuple[float, float, float]:
+def int_to_col(
+    int_val: int, sat: float = 0.5, val: float = 0.95, alpha: float = -1
+) -> tuple[float, float, float] | tuple[float, float, float, float]:
     """Uses the golden angle to generate colors programmatically with minimum overlap between colors.
     https://martin.ankerl.com/2009/12/09/how-to-create-random-colors-programmatically/
 
     Args:
-        int_val (_type_): index of an object you want to color, this is mapped to hue in HSV
+        int_val (int): index of an object you want to color, this is mapped to hue in HSV
         sat (float, optional): saturation in HSV. Defaults to 0.5.
         val (float, optional): value in HSV. Defaults to 0.95.
-        alpha (int, optional): transparency.  If -1 then only RGB is returned, if 0 or greater, RGBA is returned. Defaults to -1.
+        alpha (float, optional): transparency.  If -1 then only RGB is returned, if 0 or greater, RGBA is returned. Defaults to -1.
 
     Returns:
         tuple[float, float, float] | tuple[float, float, float, float]: either RGB or RGBA vector
@@ -138,7 +176,23 @@ def int_to_col(int_val, sat=0.5, val=0.95, alpha=-1) -> tuple[float, float, floa
     return rgb
 
 
-def lerp(value, input_low: float, input_high: float, output_low: float, output_high: float):
+def lerp(
+    value: float, input_low: float, input_high: float, output_low: float, output_high: float
+) -> float:
+    """Linear interpolation between two ranges.
+
+    Maps a value from one range [input_low, input_high] to another range [output_low, output_high].
+
+    Args:
+        value (float): The input value to interpolate
+        input_low (float): The lower bound of the input range
+        input_high (float): The upper bound of the input range
+        output_low (float): The lower bound of the output range
+        output_high (float): The upper bound of the output range
+
+    Returns:
+        float: The interpolated value in the output range
+    """
     input_low = float(input_low)
     return output_low + ((float(value) - input_low) / (float(input_high) - input_low)) * (
         float(output_high) - output_low
@@ -146,10 +200,26 @@ def lerp(value, input_low: float, input_high: float, output_low: float, output_h
 
 
 def color_tuple_to_css(color: tuple[float, float, float]) -> str:
+    """Convert a RGB color tuple to CSS rgb format string.
+
+    Args:
+        color (tuple[float, float, float]): RGB color tuple with values in range [0.0, 1.0]
+
+    Returns:
+        str: CSS color string in format 'rgb(r, g, b)' with values in range [0, 255]
+    """
     return f"rgb{(color[0] * 255, color[1] * 255, color[2] * 255)}"
 
 
-def color_tuple_to_255(color: tuple[float, float, float]) -> tuple[float, float, float]:
+def color_tuple_to_255(color: tuple[float, float, float]) -> tuple[int, int, int]:
+    """Convert a RGB color tuple with values in range [0.0, 1.0] to values in range [0, 255].
+
+    Args:
+        color (tuple[float, float, float]): RGB color tuple with values in range [0.0, 1.0]
+
+    Returns:
+        tuple[int, int, int]: RGB color tuple with values clamped to range [0, 255]
+    """
     return (
         min(int(color[0] * 255), 255),
         min(int(color[1] * 255), 255),
@@ -157,25 +227,76 @@ def color_tuple_to_255(color: tuple[float, float, float]) -> tuple[float, float,
     )
 
 
-def gen_path(filename, folder="generic", suffix=".dat"):
+def gen_path(filename: str, folder: str = "generic", suffix: str = ".dat") -> str:
+    """Generate a unique path for a file in the cache directory.
+
+    Creates a directory structure in the 'cachedir' folder and returns a path
+    with a UUID to ensure uniqueness.
+
+    Args:
+        filename (str): Base name for the file
+        folder (str, optional): Subfolder within cachedir. Defaults to "generic".
+        suffix (str, optional): File extension. Defaults to ".dat".
+
+    Returns:
+        str: Absolute path to a unique file location
+    """
     path = Path(f"cachedir/{folder}/{filename}/")
     path.mkdir(parents=True, exist_ok=True)
     return f"{path.absolute().as_posix()}/{filename}_{uuid4()}{suffix}"
 
 
 def gen_video_path(video_name: str = "vid", extension: str = ".mp4") -> str:
+    """Generate a unique path for a video file in the cache directory.
+
+    Args:
+        video_name (str, optional): Base name for the video file. Defaults to "vid".
+        extension (str, optional): Video file extension. Defaults to ".mp4".
+
+    Returns:
+        str: Absolute path to a unique video file location
+    """
     return gen_path(video_name, "vid", extension)
 
 
-def gen_image_path(image_name: str = "img", filetype=".png") -> str:
+def gen_image_path(image_name: str = "img", filetype: str = ".png") -> str:
+    """Generate a unique path for an image file in the cache directory.
+
+    Args:
+        image_name (str, optional): Base name for the image file. Defaults to "img".
+        filetype (str, optional): Image file extension. Defaults to ".png".
+
+    Returns:
+        str: Absolute path to a unique image file location
+    """
     return gen_path(image_name, "img", filetype)
 
 
-def gen_rerun_data_path(rrd_name: str = "rrd", filetype=".rrd") -> str:
+def gen_rerun_data_path(rrd_name: str = "rrd", filetype: str = ".rrd") -> str:
+    """Generate a unique path for a rerun data file in the cache directory.
+
+    Args:
+        rrd_name (str, optional): Base name for the rerun data file. Defaults to "rrd".
+        filetype (str, optional): File extension. Defaults to ".rrd".
+
+    Returns:
+        str: Absolute path to a unique rerun data file location
+    """
     return gen_path(rrd_name, "rrd", filetype)
 
 
 def callable_name(any_callable: Callable[..., Any]) -> str:
+    """Extract the name of a callable object, handling various callable types.
+
+    This function attempts to extract the name of a callable object, including
+    regular functions, partial functions, and other callables.
+
+    Args:
+        any_callable (Callable[..., Any]): The callable object to get the name from
+
+    Returns:
+        str: The name of the callable
+    """
     if isinstance(any_callable, partial):
         return any_callable.func.__name__
     try:
@@ -184,8 +305,20 @@ def callable_name(any_callable: Callable[..., Any]) -> str:
         return str(any_callable)
 
 
-def listify(obj) -> list:
-    """Take an object and turn it into a list if its not already a list.  However if the object is none, don't turn it into a list"""
+def listify(obj: Any) -> List[Any] | None:
+    """Convert an object to a list if it's not already a list.
+
+    This function handles conversion of various object types to lists, with special
+    handling for None values and existing list/tuple types.
+
+    Args:
+        obj (Any): The object to convert to a list
+
+    Returns:
+        List[Any] | None: A list containing the object, the object itself if it was
+            already a list, a list from the tuple if it was a tuple, or None if the
+            input was None
+    """
     if obj is None:
         return None
     if isinstance(obj, list):
@@ -195,13 +328,29 @@ def listify(obj) -> list:
     return [obj]
 
 
-def get_name(var):
+def get_name(var: Any) -> str:
+    """Extract the name from a variable, handling param.Parameter objects.
+
+    Args:
+        var (Any): The variable to extract the name from
+
+    Returns:
+        str: The name of the variable
+    """
     if isinstance(var, param.Parameter):
         return var.name
     return var
 
 
-def params_to_str(param_list: List[param.Parameter]):
+def params_to_str(param_list: List[param.Parameter]) -> List[str]:
+    """Convert a list of param.Parameter objects to a list of their names.
+
+    Args:
+        param_list (List[param.Parameter]): List of parameter objects
+
+    Returns:
+        List[str]: List of parameter names
+    """
     return [get_name(i) for i in param_list]
 
 
@@ -212,8 +361,8 @@ def publish_file(filepath: str, remote: str, branch_name: str) -> str:  # pragma
 
         def publish_args(branch_name) -> Tuple[str, str]:
             return (
-                "https://github.com/dyson-ai/bencher.git",
-                f"https://github.com/dyson-ai/bencher/blob/{branch_name}")
+                "https://github.com/blooop/bencher.git",
+                f"https://github.com/blooop/bencher/blob/{branch_name}")
 
 
     Args:

bencher/variables/inputs.py

@@ -7,7 +7,15 @@ from bencher.variables.sweep_base import SweepBase, shared_slots
 
 
 class SweepSelector(Selector, SweepBase):
-    """A class to represent a parameter sweep of bools"""
+    """A class representing a parameter sweep for selectable options.
+
+    This class extends both Selector and SweepBase to provide parameter sweeping
+    capabilities for categorical variables that have a predefined set of options.
+
+    Attributes:
+        units (str): The units of measurement for the parameter
+        samples (int): The number of samples to take from the available options
+    """
 
     __slots__ = shared_slots
 
@@ -22,14 +30,26 @@ class SweepSelector(Selector, SweepBase):
             self.samples = samples
 
     def values(self) -> List[Any]:
-        """return all the values for a parameter sweep.  If debug is true return a reduced list"""
+        """Return all the values for the parameter sweep.
+
+        Returns:
+            List[Any]: A list of parameter values to sweep through
+        """
         return self.indices_to_samples(self.samples, self.objects)
 
 
 class BoolSweep(SweepSelector):
-    """A class to represent a parameter sweep of bools"""
+    """A class representing a parameter sweep for boolean values.
+
+    This class extends SweepSelector to provide parameter sweeping capabilities
+    specifically for boolean values (True and False).
 
-    def __init__(self, units: str = "ul", samples: int = None, default=True, **params):
+    Attributes:
+        units (str): The units of measurement for the parameter
+        samples (int): The number of samples to take (typically 2 for booleans)
+    """
+
+    def __init__(self, units: str = "ul", samples: int = None, default: bool = True, **params):
         SweepSelector.__init__(
             self,
             units=units,
@@ -41,7 +61,15 @@ class BoolSweep(SweepSelector):
 
 
 class StringSweep(SweepSelector):
-    """A class to represent a parameter sweep of strings"""
+    """A class representing a parameter sweep for string values.
+
+    This class extends SweepSelector to provide parameter sweeping capabilities
+    specifically for a list of string values.
+
+    Attributes:
+        units (str): The units of measurement for the parameter
+        samples (int): The number of samples to take from the available strings
+    """
 
     def __init__(
         self,
@@ -61,11 +89,21 @@ class StringSweep(SweepSelector):
 
 
 class EnumSweep(SweepSelector):
-    """A class to represent a parameter sweep of enums"""
+    """A class representing a parameter sweep for enum values.
+
+    This class extends SweepSelector to provide parameter sweeping capabilities
+    specifically for enumeration types, supporting both enum types and lists of enum values.
+
+    Attributes:
+        units (str): The units of measurement for the parameter
+        samples (int): The number of samples to take from the available enum values
+    """
 
     __slots__ = shared_slots
 
-    def __init__(self, enum_type: Enum | List[Enum], units="ul", samples=None, **params):
+    def __init__(
+        self, enum_type: Enum | List[Enum], units: str = "ul", samples: int = None, **params
+    ):
         # The enum can either be an Enum type or a list of enums
         list_of_enums = isinstance(enum_type, list)
         selector_list = enum_type if list_of_enums else list(enum_type)
@@ -82,11 +120,23 @@ class EnumSweep(SweepSelector):
 
 
 class IntSweep(Integer, SweepBase):
-    """A class to represent a parameter sweep of ints"""
+    """A class representing a parameter sweep for integer values.
+
+    This class extends both Integer and SweepBase to provide parameter sweeping capabilities
+    specifically for integer values within specified bounds or with custom sample values.
+
+    Attributes:
+        units (str): The units of measurement for the parameter
+        samples (int): The number of samples to take from the range
+        sample_values (List[int], optional): Specific integer values to use as samples instead of
+            generating them from bounds. If provided, overrides the samples parameter.
+    """
 
     __slots__ = shared_slots + ["sample_values"]
 
-    def __init__(self, units="ul", samples=None, sample_values=None, **params):
+    def __init__(
+        self, units: str = "ul", samples: int = None, sample_values: List[int] = None, **params
+    ):
         SweepBase.__init__(self)
         Integer.__init__(self, **params)
 
@@ -107,7 +157,14 @@ class IntSweep(Integer, SweepBase):
                 self.default = sample_values[0]
 
     def values(self) -> List[int]:
-        """return all the values for a parameter sweep.  If debug is true return the  list"""
+        """Return all the values for the parameter sweep.
+
+        If sample_values is provided, returns those values. Otherwise generates values
+        within the specified bounds.
+
+        Returns:
+            List[int]: A list of integer values to sweep through
+        """
         sample_values = (
             self.sample_values
             if self.sample_values is not None
@@ -136,11 +193,29 @@ class IntSweep(Integer, SweepBase):
 
 
 class FloatSweep(Number, SweepBase):
-    """A class to represent a parameter sweep of floats"""
+    """A class representing a parameter sweep for floating point values.
+
+    This class extends both Number and SweepBase to provide parameter sweeping capabilities
+    specifically for floating point values within specified bounds or with custom sample values.
+
+    Attributes:
+        units (str): The units of measurement for the parameter
+        samples (int): The number of samples to take from the range
+        sample_values (List[float], optional): Specific float values to use as samples instead of
+            generating them from bounds. If provided, overrides the samples parameter.
+        step (float, optional): Step size between samples when generating values from bounds
+    """
 
     __slots__ = shared_slots + ["sample_values"]
 
-    def __init__(self, units="ul", samples=10, sample_values=None, step=None, **params):
+    def __init__(
+        self,
+        units: str = "ul",
+        samples: int = 10,
+        sample_values: List[float] = None,
+        step: float = None,
+        **params,
+    ):
         SweepBase.__init__(self)
         Number.__init__(self, step=step, **params)
 
@@ -156,7 +231,14 @@ class FloatSweep(Number, SweepBase):
                 self.default = sample_values[0]
 
     def values(self) -> List[float]:
-        """return all the values for a parameter sweep.  If debug is true return a reduced list"""
+        """Return all the values for the parameter sweep.
+
+        If sample_values is provided, returns those values. Otherwise generates values
+        within the specified bounds, either using linspace (when step is None) or arange.
+
+        Returns:
+            List[float]: A list of float values to sweep through
+        """
         samps = self.samples
         if self.sample_values is None:
             if self.step is None:
@@ -166,7 +248,20 @@ class FloatSweep(Number, SweepBase):
         return self.sample_values
 
 
-def box(name, center, width):
+def box(name: str, center: float, width: float) -> FloatSweep:
+    """Create a FloatSweep parameter centered around a value with a given width.
+
+    This is a convenience function to create a bounded FloatSweep parameter with
+    bounds centered on a specific value, extending by the width in both directions.
+
+    Args:
+        name (str): The name of the parameter
+        center (float): The center value of the parameter
+        width (float): The distance from the center to the bounds in both directions
+
+    Returns:
+        FloatSweep: A FloatSweep parameter with the specified name, default, and bounds
+    """
     var = FloatSweep(default=center, bounds=(center - width, center + width))
     var.name = name
     return var
@@ -195,6 +290,18 @@ def p(
     return {"name": name, "values": values, "max_level": max_level, "samples": samples}
 
 
-def with_level(arr: list, level) -> list:
+def with_level(arr: list, level: int) -> list:
+    """Apply level-based sampling to a list of values.
+
+    This function uses an IntSweep with the provided values and applies level-based
+    sampling to it, returning the resulting values.
+
+    Args:
+        arr (list): List of values to sample from
+        level (int): The sampling level to apply (higher levels provide more samples)
+
+    Returns:
+        list: The level-sampled values
+    """
     return IntSweep(sample_values=arr).with_level(level).values()
     # return tmp.with_sample_values(arr).with_level(level).values()

bencher/video_writer.py

@@ -88,7 +88,8 @@ class VideoWriter:
             output_path = Path(output_path)
 
         with moviepy.video.io.VideoFileClip.VideoFileClip(video_path) as video:
-            frame_time = time if time is not None else video.duration
+            frame_time = time if time is not None else video.duration - 2.0 / video.fps
+            frame_time = max(frame_time, 0)
             frame = video.get_frame(frame_time)
             Image.fromarray(frame).save(output_path)
 

bencher/worker_job.py

@@ -6,8 +6,30 @@ from bencher.utils import hmap_canonical_input
 
 @dataclass
 class WorkerJob:
-    function_input_vars: List
-    index_tuple: Tuple[int]
+    """Represents a benchmark worker job with input variables and caching information.
+
+    This class encapsulates the information needed to execute a benchmark function,
+    including input variables, dimension information, and caching metadata. It handles
+    the preparation of function inputs and calculation of hash signatures for caching.
+
+    Attributes:
+        function_input_vars (List): The values of the input variables to pass to the function
+        index_tuple (Tuple[int]): The indices of these values in the N-dimensional result array
+        dims_name (List[str]): The names of the input dimensions
+        constant_inputs (dict): Dictionary of any constant input values
+        bench_cfg_sample_hash (str): Hash of the benchmark configuration without repeats
+        tag (str): Tag for grouping related jobs
+        function_input (dict): Complete input as a dictionary with dimension names as keys
+        canonical_input (Tuple[Any]): Canonical representation of inputs for caching
+        fn_inputs_sorted (List[Tuple[str, Any]]): Sorted representation of function inputs
+        function_input_signature_pure (str): Hash of the function inputs and tag
+        function_input_signature_benchmark_context (str): Comprehensive hash including benchmark context
+        found_in_cache (bool): Whether this job result was found in cache
+        msgs (List[str]): Messages related to this job's execution
+    """
+
+    function_input_vars: List[Any]
+    index_tuple: Tuple[int, ...]
     dims_name: List[str]
     constant_inputs: dict
     bench_cfg_sample_hash: str
@@ -15,13 +37,19 @@ class WorkerJob:
 
     function_input: dict = None
     canonical_input: Tuple[Any] = None
-    fn_inputs_sorted: List[str] = None
+    fn_inputs_sorted: List[Tuple[str, Any]] = None
     function_input_signature_pure: str = None
     function_input_signature_benchmark_context: str = None
     found_in_cache: bool = False
     msgs: List[str] = field(default_factory=list)
 
     def setup_hashes(self) -> None:
+        """Set up the function inputs and calculate hash signatures for caching.
+
+        This method prepares the function inputs by combining function input variables
+        with dimensions and constant inputs. It also calculates hash signatures used
+        for caching results and tracking job execution.
+        """
         self.function_input = dict(zip(self.dims_name, self.function_input_vars))
 
         self.canonical_input = hmap_canonical_input(self.function_input)