Merge pull request #214 from boutproject/animate_list-per-var-kwargs

Allow lists to be passed for kwargs in BoutDataset.animate_list()

Author: johnomotani

xbout/boutdataset.py

@@ -921,7 +921,8 @@ def animate_list(
             If a dict is passed, the dict entries are passed as arguments to
             tight_layout()
         **kwargs : dict, optional
-            Additional keyword arguments are passed on to each animation function
+            Additional keyword arguments are passed on to each animation function, per
+            variable if a sequence is given.
 
         Returns
         -------
@@ -976,6 +977,8 @@ def _expand_list_arg(arg, arg_name):
         aspect = _expand_list_arg(aspect, "aspect")
         extend = _expand_list_arg(extend, "extend")
         axis_coords = _expand_list_arg(axis_coords, "axis_coords")
+        for k in kwargs:
+            kwargs[k] = _expand_list_arg(kwargs[k], k)
 
         blocks = []
 
@@ -986,17 +989,19 @@ def is_list(variable):
                 or isinstance(variable, set)
             )
 
-        for subplot_args in zip(
-            variables,
-            axes,
-            poloidal_plot,
-            axis_coords,
-            vmin,
-            vmax,
-            logscale,
-            titles,
-            aspect,
-            extend,
+        for i, subplot_args in enumerate(
+            zip(
+                variables,
+                axes,
+                poloidal_plot,
+                axis_coords,
+                vmin,
+                vmax,
+                logscale,
+                titles,
+                aspect,
+                extend,
+            )
         ):
 
             (
@@ -1012,6 +1017,8 @@ def is_list(variable):
                 this_extend,
             ) = subplot_args
 
+            this_kwargs = {k: v[i] for k, v in kwargs.items()}
+
             divider = make_axes_locatable(ax)
             cax = divider.append_axes("right", size="5%", pad=0.1)
 
@@ -1051,7 +1058,7 @@ def is_list(variable):
                             animate=False,
                             axis_coords=this_axis_coords,
                             aspect=this_aspect,
-                            **kwargs,
+                            **this_kwargs,
                         )
                     )
                 else:
@@ -1065,7 +1072,7 @@ def is_list(variable):
                                 axis_coords=this_axis_coords,
                                 aspect=this_aspect,
                                 label=w.name,
-                                **kwargs,
+                                **this_kwargs,
                             )
                         )
                     legend = ax.legend()
@@ -1086,7 +1093,7 @@ def is_list(variable):
                         logscale=this_logscale,
                         aspect=this_aspect,
                         extend=this_extend,
-                        **kwargs,
+                        **this_kwargs,
                     )
                     for block in var_blocks:
                         blocks.append(block)
@@ -1104,7 +1111,7 @@ def is_list(variable):
                             logscale=this_logscale,
                             aspect=this_aspect,
                             extend=this_extend,
-                            **kwargs,
+                            **this_kwargs,
                         )
                     )
             else: