docs: add google style doc-strings and update auto-docs in relation to branch #358

Author: harryswift01

CodeEntropy/levels/execution/__init__.py

@@ -7,7 +7,18 @@ def chunk_frame_indices(
     frame_indices: list[int],
     chunk_size: int,
 ) -> list[tuple[int, ...]]:
-    """Split selected frame indices into deterministic fixed-size chunks."""
+    """Split frame indices into deterministic fixed-size chunks.
+
+    Args:
+        frame_indices: Ordered selected frame indices to split.
+        chunk_size: Maximum number of frame indices per chunk.
+
+    Returns:
+        A list of frame-index tuples preserving input order.
+
+    Raises:
+        ValueError: If ``chunk_size`` is less than one.
+    """
     if chunk_size < 1:
         raise ValueError("chunk_size must be >= 1")
 

CodeEntropy/levels/execution/chunks.py

@@ -22,7 +22,15 @@ class ExecutionPolicy:
     max_frame_in_flight_multiplier: int = 1
 
     def requested_worker_count(self, shared_data: dict[str, Any]) -> int:
-        """Return the worker-process count requested by the current run."""
+        """Return the requested worker-process count.
+
+        Args:
+            shared_data: Shared workflow data that may contain ``args`` with local Dask
+                or HPC worker settings.
+
+        Returns:
+            The requested worker count, clamped to at least one.
+        """
         args = shared_data.get("args")
 
         dask_workers = getattr(args, "dask_workers", None)
@@ -37,7 +45,15 @@ def requested_worker_count(self, shared_data: dict[str, Any]) -> int:
         return 1
 
     def frame_chunk_size(self, shared_data: dict[str, Any], *, n_frames: int) -> int:
-        """Choose a deterministic frame chunk size for the current run."""
+        """Choose a deterministic frame chunk size.
+
+        Args:
+            shared_data: Shared workflow data used to infer requested worker count.
+            n_frames: Number of selected frames for the current run.
+
+        Returns:
+            The frame chunk size clamped between the policy minimum and maximum.
+        """
         n_frames = max(1, int(n_frames))
         workers = self.requested_worker_count(shared_data)
         target_chunks = max(1, workers * int(self.target_frame_chunks_per_worker))
@@ -54,7 +70,15 @@ def max_frame_in_flight_tasks(
         *,
         n_chunks: int,
     ) -> int:
-        """Choose how many frame chunks may be active at once."""
+        """Choose the maximum number of active frame-chunk tasks.
+
+        Args:
+            shared_data: Shared workflow data used to infer requested worker count.
+            n_chunks: Number of frame chunks available for submission.
+
+        Returns:
+            The number of frame-chunk tasks allowed to be active at once.
+        """
         workers = self.requested_worker_count(shared_data)
         max_in_flight = workers * int(self.max_frame_in_flight_multiplier)
         return max(1, min(int(n_chunks), int(max_in_flight)))

CodeEntropy/levels/execution/policy.py

@@ -8,12 +8,30 @@
 
 
 def stable_keys(mapping: dict[Any, Any]) -> list[Any]:
-    """Return mapping keys in a deterministic order across Python processes."""
+    """Return mapping keys in deterministic order.
+
+    Args:
+        mapping: Mapping whose keys should be ordered independently of process hash
+            randomisation.
+
+    Returns:
+        A list of keys sorted by key type name and representation.
+    """
     return sorted(mapping.keys(), key=lambda key: (type(key).__name__, repr(key)))
 
 
 def merge_means(old_mean: Any, old_n: int, new_mean: Any, new_n: int) -> Any:
-    """Merge two means with their sample counts."""
+    """Merge two running means using their sample counts.
+
+    Args:
+        old_mean: Existing mean value, or ``None`` if no samples have been seen.
+        old_n: Number of samples represented by ``old_mean``.
+        new_mean: New mean value to merge.
+        new_n: Number of samples represented by ``new_mean``.
+
+    Returns:
+        The merged mean. If ``new_n`` is zero or negative, ``old_mean`` is returned.
+    """
     if new_n <= 0:
         return old_mean
     if old_mean is None or old_n <= 0:
@@ -23,7 +41,16 @@ def merge_means(old_mean: Any, old_n: int, new_mean: Any, new_n: int) -> Any:
 
 
 def incremental_mean(old: Any, new: Any, n: int) -> Any:
-    """Compute an incremental mean."""
+    """Update a running mean with one new sample.
+
+    Args:
+        old: Existing running mean, or ``None`` for the first sample.
+        new: New sample to incorporate.
+        n: One-based sample count after adding ``new``.
+
+    Returns:
+        The updated running mean.
+    """
     if old is None:
         return new.copy() if hasattr(new, "copy") else new
     return old + (new - old) / float(n)
@@ -34,7 +61,12 @@ class NeighborReducer:
 
     @staticmethod
     def initialise(shared_data: dict[str, Any]) -> None:
-        """Initialise parent-side neighbour reduction state."""
+        """Initialise parent-side neighbour accumulators.
+
+        Args:
+            shared_data: Shared workflow data containing ``groups``. The method writes
+                ``neighbor_totals`` and ``neighbor_samples``.
+        """
         shared_data["neighbor_totals"] = {
             group_id: 0 for group_id in shared_data["groups"].keys()
         }
@@ -47,7 +79,13 @@ def reduce_frame_output(
         shared_data: dict[str, Any],
         frame_neighbors: dict[int, tuple[int, int]] | None,
     ) -> None:
-        """Reduce one frame's neighbour-count payload."""
+        """Merge one frame's neighbour-count payload.
+
+        Args:
+            shared_data: Shared workflow data containing neighbour total/sample
+                accumulators.
+            frame_neighbors: Optional mapping of group id to ``(count, sample_count)``.
+        """
         if frame_neighbors is None:
             return
 
@@ -64,7 +102,13 @@ def merge_chunk_partial(
         neighbor_totals: dict[int, int],
         neighbor_samples: dict[int, int],
     ) -> None:
-        """Merge chunk-level neighbour totals/samples into parent accumulators."""
+        """Merge chunk-level neighbour totals and samples.
+
+        Args:
+            shared_data: Shared workflow data containing neighbour accumulators.
+            neighbor_totals: Mapping of group id to additive neighbour totals.
+            neighbor_samples: Mapping of group id to additive sample counts.
+        """
         totals = shared_data.get("neighbor_totals")
         samples = shared_data.get("neighbor_samples")
         if totals is None or samples is None:
@@ -79,7 +123,13 @@ def merge_chunk_partial(
 
     @staticmethod
     def finalise(shared_data: dict[str, Any]) -> None:
-        """Convert reduced neighbour totals into average neighbour counts."""
+        """Compute average neighbour counts from reduced totals.
+
+        Args:
+            shared_data: Shared workflow data containing ``groups``,
+                ``neighbor_totals``, and ``neighbor_samples``. The method writes
+                ``neighbors``.
+        """
         neighbors = {}
         for group_id in stable_keys(shared_data["groups"]):
             sample_count = shared_data["neighbor_samples"].get(group_id, 0)
@@ -100,7 +150,13 @@ def reduce_frame_output(
         shared_data: dict[str, Any],
         frame_out: dict[str, Any],
     ) -> None:
-        """Reduce one frame's covariance outputs into shared running means."""
+        """Reduce one frame covariance payload into parent accumulators.
+
+        Args:
+            shared_data: Shared workflow data containing covariance accumulators.
+            frame_out: Frame covariance payload with force, torque, and optional
+                force-torque sections.
+        """
         self._reduce_force_and_torque(shared_data, frame_out)
         self._reduce_forcetorque(shared_data, frame_out)
 
@@ -109,7 +165,12 @@ def merge_chunk_partial(
         shared_data: dict[str, Any],
         partial: CovarianceChunkPartial,
     ) -> None:
-        """Merge a compact worker covariance partial into parent accumulators."""
+        """Merge a worker covariance partial into parent accumulators.
+
+        Args:
+            shared_data: Shared workflow data containing covariance accumulators.
+            partial: Compact covariance partial returned by a worker frame chunk.
+        """
         self._merge_force_and_torque_partial(shared_data, partial)
         self._merge_forcetorque_partial(shared_data, partial)
 
@@ -118,7 +179,14 @@ def reduce_frame_map_output(
         shared_data: dict[str, Any],
         frame_out: dict[str, Any],
     ) -> None:
-        """Reduce one serial frame's complete MAP output into parent shared_data."""
+        """Reduce a complete serial MAP output.
+
+        Args:
+            shared_data: Shared workflow data containing covariance and neighbour
+                accumulators.
+            frame_out: MAP output containing optional ``covariance`` and ``neighbors``
+                entries.
+        """
         covariance = frame_out.get("covariance")
         if covariance is not None:
             self.reduce_frame_output(shared_data, covariance)
@@ -132,6 +200,13 @@ def _merge_force_and_torque_partial(
         shared_data: dict[str, Any],
         partial: CovarianceChunkPartial,
     ) -> None:
+        """Merge chunk force and torque means into parent accumulators.
+
+        Args:
+            shared_data: Shared workflow data containing force/torque accumulators,
+                frame counts, and ``group_id_to_index``.
+            partial: Worker covariance partial with force, torque, and count mappings.
+        """
         f_cov = shared_data["force_covariances"]
         t_cov = shared_data["torque_covariances"]
         counts = shared_data["frame_counts"]
@@ -183,6 +258,13 @@ def _merge_forcetorque_partial(
         shared_data: dict[str, Any],
         partial: CovarianceChunkPartial,
     ) -> None:
+        """Merge chunk force-torque block means into parent accumulators.
+
+        Args:
+            shared_data: Shared workflow data containing force-torque accumulators,
+                force-torque counts, and ``group_id_to_index``.
+            partial: Worker covariance partial with force-torque matrices and counts.
+        """
         ft_cov = shared_data["forcetorque_covariances"]
         ft_counts = shared_data["forcetorque_counts"]
         gid2i = shared_data["group_id_to_index"]
@@ -210,6 +292,13 @@ def _reduce_force_and_torque(
         shared_data: dict[str, Any],
         frame_out: dict[str, Any],
     ) -> None:
+        """Reduce frame force and torque matrices into running means.
+
+        Args:
+            shared_data: Shared workflow data containing force/torque accumulators,
+                frame counts, and ``group_id_to_index``.
+            frame_out: Frame covariance payload with ``force`` and ``torque`` sections.
+        """
         f_cov = shared_data["force_covariances"]
         t_cov = shared_data["torque_covariances"]
         counts = shared_data["frame_counts"]
@@ -266,6 +355,14 @@ def _reduce_forcetorque(
         shared_data: dict[str, Any],
         frame_out: dict[str, Any],
     ) -> None:
+        """Reduce frame force-torque matrices into running means.
+
+        Args:
+            shared_data: Shared workflow data containing force-torque accumulators,
+                force-torque counts, and ``group_id_to_index``.
+            frame_out: Frame covariance payload that may contain a ``forcetorque``
+                section.
+        """
         if "forcetorque" not in frame_out:
             return
 

CodeEntropy/levels/execution/reducers.py

@@ -35,6 +35,14 @@ def __init__(
         policy: ExecutionPolicy,
         universe_operations: Any | None = None,
     ) -> None:
+        """Initialise the frame scheduler.
+
+        Args:
+            frame_dag: Built or buildable frame-local DAG used for serial execution.
+            policy: Internal execution policy for chunking and in-flight limits.
+            universe_operations: Optional universe-operation adapter forwarded to worker
+                frame graphs.
+        """
         self._frame_dag = frame_dag
         self._policy = policy
         self._universe_operations = universe_operations
@@ -47,7 +55,17 @@ def execute(
         frame_indices: list[int],
         progress: _RichProgressSink | None = None,
     ) -> None:
-        """Execute frame-local MAP work and reduce it into ``shared_data``."""
+        """Execute frame-local MAP work and reduce outputs.
+
+        Args:
+            shared_data: Shared workflow data containing serial or Dask execution
+                inputs.
+            frame_indices: Ordered selected frame indices to execute.
+            progress: Optional progress sink used for reporting frame-stage progress.
+
+        Raises:
+            RuntimeError: If Dask execution is requested but unavailable or incomplete.
+        """
         task: TaskID | None = None
         if progress is not None:
             task = progress.add_task(
@@ -84,7 +102,14 @@ def _run_serial(
         progress: _RichProgressSink | None,
         task: TaskID | None,
     ) -> None:
-        """Execute frame-local MAP work serially and reduce immediately."""
+        """Execute frame MAP work serially.
+
+        Args:
+            shared_data: Shared workflow data mutated by parent-side reducers.
+            frame_indices: Ordered frame indices to process.
+            progress: Optional progress sink.
+            task: Optional progress task identifier.
+        """
         neighbor_helper = Neighbors()
 
         for frame_index in frame_indices:
@@ -111,7 +136,21 @@ def _run_dask(
         progress: _RichProgressSink | None,
         task: TaskID | None,
     ) -> None:
-        """Execute frame-chunk MAP tasks with bounded deterministic reduction."""
+        """Execute frame MAP work using bounded Dask futures.
+
+        Args:
+            shared_data: Shared workflow data mutated by parent-side reducers.
+            frame_indices: Ordered frame indices to process.
+            client: Dask distributed client-like object.
+            progress: Optional progress sink.
+            task: Optional progress task identifier.
+
+        Raises:
+            RuntimeError: If ``dask.distributed`` is unavailable or the number of
+                reduced frames does not match the selected frame count.
+            Exception: Propagates worker or Dask client errors after cancelling active
+                futures.
+        """
         try:
             from distributed import wait
         except ImportError as exc:
@@ -139,6 +178,12 @@ def _run_dask(
         pending_results: dict[int, FrameChunkResult] = {}
 
         def submit_next() -> bool:
+            """Submit the next frame-chunk task if one is available.
+
+            Returns:
+                ``True`` if a task was submitted, otherwise ``False`` when all tasks
+                    have already been submitted.
+            """
             nonlocal submitted
             try:
                 frame_task = next(frame_task_iter)
@@ -157,6 +202,11 @@ def submit_next() -> bool:
             return True
 
         def reduce_ready_results() -> None:
+            """Reduce completed frame chunks in deterministic chunk-index order.
+
+            Mutates enclosing scheduler state by consuming pending results, advancing
+            the next expected reduce index, and updating the completed-frame count.
+            """
             nonlocal completed, next_reduce_index
             while next_reduce_index in pending_results:
                 chunk_result = pending_results.pop(next_reduce_index)
@@ -232,7 +282,15 @@ def _make_frame_chunk_tasks(
         shared_data: dict[str, Any],
         frame_indices: list[int],
     ) -> list[FrameChunkTask]:
-        """Build explicit frame-chunk MAP tasks."""
+        """Build frame-chunk task descriptors.
+
+        Args:
+            shared_data: Shared workflow data used by the execution policy.
+            frame_indices: Ordered selected frame indices to split into chunks.
+
+        Returns:
+            A list of ``FrameChunkTask`` descriptors with deterministic chunk indices.
+        """
         chunk_size = self._policy.frame_chunk_size(
             shared_data,
             n_frames=len(frame_indices),