Make docstring even better

Author: cphyc

yt/data_objects/selection_objects/data_selection_objects.py

@@ -525,7 +525,7 @@ def piter(
             indices and the values will be whatever is assigned to the *result*
             attribute on the storage during iteration.
         barrier : bool
-            Should a barier be placed at the end of iteration?
+            Should a barrier be placed at the end of iteration?
         dynamic : bool
             This governs whether or not dynamic load balancing will be enabled.
             This requires one dedicated processor; if this is enabled with a set of
@@ -535,16 +535,28 @@ def piter(
             This specifies the reduction operation to be applied to the results
             from each processor.
             - None: no reduction will be applied and the storage object will
-              contain one result per chunk in the container.
+                contain one result per chunk in the container.
             - concat: the storage object will contain a flattened list of
-              each results.
+                each results.
             - cat_on_root: same as concat, but only the root processor will
-              contain anything.
+                contain anything.
             - sum, min, max: the storage object will contain the result
-              of applying the operation until getting a single value.
+                of applying the operation until getting a single value.
+
+        Important limitation
+        --------------------
+        When using `storage`, the result *must* be a dictionary. See the
+        examples below.
 
         Example
         -------
+
+        Here is an example of how to gather all data on root, reading in
+        parallel. Other MPI tasks will have nothing in `my_storage`.
+
+        >>> import yt
+        >>> ds = yt.load("output_00080")
+        ... ad = ds.all_data()
         >>> my_storage = {}
         ... for sto, chunk in ad.piter(storage=my_storage, reduction="cat_on_root"):
         ...     sto.result = {
@@ -557,18 +569,34 @@ def piter(
         ...     # Contains *all* the gas temperatures
         ...     my_storage["gas", "temperature"]
 
+        Here is an example of how to sum the total mass of all gas cells in
+        the dataset, storing the result in `my_storage` on all processors.
+
         >>> my_storage = {}
         ... for sto, chunk in ad.piter(storage=my_storage, reduction="sum"):
         ...     sto.result = {
         ...         "total_mass": chunk["gas", "cell_mass"].sum(),
         ...     }
         ... print("Total mass: ", my_storage["total_mass"])
 
-        Note
-        ----
-        `sto.result` has to be a dictionary.
+        Here is an example of how to read all data in parallel and
+        have the results available on all processors.
 
+        >>> my_storage = {}
+        ... for sto, chunk in ad.piter(storage=my_storage, reduction="cat"):
+        ...     sto.result = {("gas", "density"): chunk["gas", "density"]}
+        ... print(my_storage["gas", "density"])
+
+        This is equivalent (but faster, since reading is parallelized) to the
+        following
+
+        >>> my_storage = {("gas", "density"): ad["gas", "density"]}
         """
+        if reduction is not None and storage is None:
+            raise ValueError(
+                "If reduction is specified, you must pass in a storage dictionary."
+            )
+
         yield from parallel_objects(
             self.chunks([], "io"),
             njobs=njobs,

yt/utilities/parallel_tools/parallel_analysis_interface.py

@@ -506,15 +506,23 @@ def parallel_objects(
         indices and the values will be whatever is assigned to the *result*
         attribute on the storage during iteration.
     barrier : bool
-        Should a barier be placed at the end of iteration?
+        Should a barrier be placed at the end of iteration?
     dynamic : bool
         This governs whether or not dynamic load balancing will be enabled.
         This requires one dedicated processor; if this is enabled with a set of
         128 processors available, only 127 will be available to iterate over
         objects as one will be load balancing the rest.
     reduction : Literal[None, "sum", "max", "min", "cat", "cat_on_root"]
         This specifies the reduction operation to be applied to the results
-        from each processor.  The default is None, which doesn't apply any.
+        from each processor.
+        - None: no reduction will be applied and the storage object will
+            contain one result per chunk in the container.
+        - concat: the storage object will contain a flattened list of
+            each results.
+        - cat_on_root: same as concat, but only the root processor will
+            contain anything.
+        - sum, min, max: the storage object will contain the result
+            of applying the operation until getting a single value.
 
     Examples
     --------