Add fig2img

neurokit2/misc/__init__.py

@@ -5,7 +5,7 @@
 """
 
 from ._warnings import NeuroKitWarning
-from .random import check_random_state, check_random_state_children, spawn_random_state
+from .check_random_state import check_random_state, check_random_state_children, spawn_random_state
 from .check_type import check_type
 from .copyfunction import copyfunction
 from .expspace import expspace
@@ -21,6 +21,7 @@
 from .replace import replace
 from .type_converters import as_vector
 from .report import create_report
+from .fig2img import fig2img
 
 
 __all__ = [
@@ -43,4 +44,5 @@
     "check_random_state_children",
     "spawn_random_state",
     "create_report",
+    "fig2img",
 ]

neurokit2/misc/random.py → neurokit2/misc/check_random_state.py

@@ -25,6 +25,7 @@ def check_random_state(seed=None):
     rng: numpy.random.Generator or numpy.random.RandomState
         Random number generator.
     """
+
     # If seed is an integer, use the legacy RandomState class
     if isinstance(seed, numbers.Integral):
         return np.random.RandomState(seed)
@@ -77,7 +78,10 @@ def spawn_random_state(rng, n_children=1):
         if rng._bit_generator._seed_seq is not None:
             rng_class = type(rng)
             bit_generator_class = type(rng._bit_generator)
-            return [rng_class(bit_generator_class(seed=s)) for s in rng._bit_generator._seed_seq.spawn(n_children)]
+            return [
+                rng_class(bit_generator_class(seed=s))
+                for s in rng._bit_generator._seed_seq.spawn(n_children)
+            ]
     except TypeError:
         # The rng does not support spawning through SeedSequence, see below
         pass
@@ -94,7 +98,9 @@ def spawn_random_state(rng, n_children=1):
     return [np.random.RandomState(seed=s) for s in temp_rng.random_raw(n_children)]
 
 
-def check_random_state_children(random_state_parent, random_state_children, n_children=1):
+def check_random_state_children(
+    random_state_parent, random_state_children, n_children=1
+):
     """**Create new independent children random number generators to be used in sub-functions**
 
     Parameters

neurokit2/misc/fig2img.py

@@ -0,0 +1,46 @@
+import io
+
+
+def fig2img(fig):
+    """Matplotlib Figure to PIL Image
+
+    Convert a Matplotlib figure to a PIL Image
+
+    Parameters
+    ----------
+    fig : plt.figure
+        Matplotlib figure.
+
+    Returns
+    ----------
+    list
+        The rescaled values.
+
+
+    Examples
+    ----------
+    .. ipython:: python
+
+      import neurokit2 as nk
+      import matplotlib.pyplot as plt
+
+      plt.plot([1, 2, 3, 4, 5])  # Make plot
+      fig = plt.gcf()  # Get current figure
+      nk.fig2img(fig)  # Convert to PIL Image
+      plt.close(fig)  # Close figure
+
+    """
+
+    try:
+        import PIL.Image
+    except ImportError as e:
+        raise ImportError(
+            "fig2img(): the 'PIL' (Pillow) module is required for this function to run. ",
+            "Please install it first (`pip install pillow`).",
+        ) from e
+
+    buffer = io.BytesIO()
+    fig.savefig(buffer)
+    buffer.seek(0)
+    img = PIL.Image.open(buffer)
+    return img