Merge pull request #15641 from anntzer/get_sample_data-npy

tacaswell · web-flow · commit a1c51a0a142e · 2020-05-22T08:56:47.000-04:00
API: Make get_sample_data autoload npy/npz files.
diff --git a/doc/api/api_changes_3.3/behaviour.rst b/doc/api/api_changes_3.3/behaviour.rst
@@ -294,3 +294,10 @@ FT_LOAD_DEFAULT, etc.  The old synonyms (respectively "either", "native",
 "auto", and "none") are still supported, but their use is discouraged.  To get
 normalized values, use `.backend_agg.get_hinting_flag`, which returns integer
 flag values.
+
+`.cbook.get_sample_data` auto-loads numpy arrays
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+When `.cbook.get_sample_data` is used to load a npy or npz file and the
+keyword-only parameter ``np_load`` is True, the file is automatically loaded
+using `numpy.load`.  ``np_load`` defaults to False for backwards compatibility,
+but will become True in a later release.
diff --git a/examples/axes_grid1/demo_axes_divider.py b/examples/axes_grid1/demo_axes_divider.py
@@ -6,14 +6,13 @@
 Axes divider to calculate location of axes and
 create a divider for them using existing axes instances.
 """
+
+from matplotlib import cbook
 import matplotlib.pyplot as plt
 
 
 def get_demo_image():
-    import numpy as np
-    from matplotlib.cbook import get_sample_data
-    f = get_sample_data("axes_grid/bivariate_normal.npy", asfileobj=False)
-    z = np.load(f)
+    z = cbook.get_sample_data("axes_grid/bivariate_normal.npy", np_load=True)
     # z is a numpy array of 15x15
     return z, (-3, 4, -4, 3)
 
diff --git a/examples/axes_grid1/demo_axes_grid.py b/examples/axes_grid1/demo_axes_grid.py
@@ -6,6 +6,7 @@
 Grid of 2x2 images with single or own colorbar.
 """
 
+from matplotlib import cbook
 import matplotlib.pyplot as plt
 from mpl_toolkits.axes_grid1 import ImageGrid
 
@@ -14,10 +15,7 @@
 
 
 def get_demo_image():
-    import numpy as np
-    from matplotlib.cbook import get_sample_data
-    f = get_sample_data("axes_grid/bivariate_normal.npy", asfileobj=False)
-    z = np.load(f)
+    z = cbook.get_sample_data("axes_grid/bivariate_normal.npy", np_load=True)
     # z is a numpy array of 15x15
     return z, (-3, 4, -4, 3)
 
diff --git a/examples/axes_grid1/demo_axes_grid2.py b/examples/axes_grid1/demo_axes_grid2.py
@@ -7,23 +7,15 @@
 """
 
 import numpy as np
-
+from matplotlib import cbook
+import matplotlib.colors
 import matplotlib.pyplot as plt
 from mpl_toolkits.axes_grid1 import ImageGrid
-import matplotlib.colors
 
 
 plt.rcParams["mpl_toolkits.legacy_colorbar"] = False
 
 
-def get_demo_image():
-    from matplotlib.cbook import get_sample_data
-    f = get_sample_data("axes_grid/bivariate_normal.npy", asfileobj=False)
-    z = np.load(f)
-    # z is a numpy array of 15x15
-    return z, (-3, 4, -4, 3)
-
-
 def add_inner_title(ax, title, loc, **kwargs):
     from matplotlib.offsetbox import AnchoredText
     from matplotlib.patheffects import withStroke
@@ -39,7 +31,8 @@ def add_inner_title(ax, title, loc, **kwargs):
 fig = plt.figure(figsize=(6, 6))
 
 # Prepare images
-Z, extent = get_demo_image()
+Z = cbook.get_sample_data("axes_grid/bivariate_normal.npy", np_load=True)
+extent = (-3, 4, -4, 3)
 ZS = [Z[i::3, :] for i in range(3)]
 extent = extent[0], extent[1]/3., extent[2], extent[3]
 
diff --git a/examples/axes_grid1/demo_axes_rgb.py b/examples/axes_grid1/demo_axes_rgb.py
@@ -13,9 +13,7 @@
 
 
 def get_rgb():
-    f = cbook.get_sample_data("axes_grid/bivariate_normal.npy")
-    Z = np.load(f)  # 15x15 numpy array.
-
+    Z = cbook.get_sample_data("axes_grid/bivariate_normal.npy", np_load=True)
     Z[Z < 0] = 0.
     Z = Z / Z.max()
 
diff --git a/examples/axes_grid1/demo_colorbar_of_inset_axes.py b/examples/axes_grid1/demo_colorbar_of_inset_axes.py
@@ -4,36 +4,25 @@
 ===========================
 
 """
-import matplotlib.pyplot as plt
 
+from matplotlib import cbook
+import matplotlib.pyplot as plt
 from mpl_toolkits.axes_grid1.inset_locator import inset_axes, zoomed_inset_axes
 
 
-def get_demo_image():
-    from matplotlib.cbook import get_sample_data
-    import numpy as np
-    f = get_sample_data("axes_grid/bivariate_normal.npy", asfileobj=False)
-    z = np.load(f)
-    # z is a numpy array of 15x15
-    return z, (-3, 4, -4, 3)
-
-
 fig, ax = plt.subplots(figsize=[5, 4])
 
-Z, extent = get_demo_image()
-
-ax.set(aspect=1,
-       xlim=(-15, 15),
-       ylim=(-20, 5))
+Z = cbook.get_sample_data("axes_grid/bivariate_normal.npy", np_load=True)
+extent = (-3, 4, -4, 3)
 
+ax.set(aspect=1, xlim=(-15, 15), ylim=(-20, 5))
 
 axins = zoomed_inset_axes(ax, zoom=2, loc='upper left')
 im = axins.imshow(Z, extent=extent, origin="lower")
 
 plt.xticks(visible=False)
 plt.yticks(visible=False)
 
-
 # colorbar
 cax = inset_axes(axins,
                  width="5%",  # width = 10% of parent_bbox width
@@ -43,7 +32,6 @@ def get_demo_image():
                  bbox_transform=axins.transAxes,
                  borderpad=0,
                  )
-
 fig.colorbar(im, cax=cax)
 
 plt.show()
diff --git a/examples/axes_grid1/demo_edge_colorbar.py b/examples/axes_grid1/demo_edge_colorbar.py
@@ -7,6 +7,7 @@
 of an image grid.
 """
 
+from matplotlib import cbook
 import matplotlib.pyplot as plt
 from mpl_toolkits.axes_grid1 import AxesGrid
 
@@ -15,10 +16,7 @@
 
 
 def get_demo_image():
-    import numpy as np
-    from matplotlib.cbook import get_sample_data
-    f = get_sample_data("axes_grid/bivariate_normal.npy", asfileobj=False)
-    z = np.load(f)
+    z = cbook.get_sample_data("axes_grid/bivariate_normal.npy", np_load=True)
     # z is a numpy array of 15x15
     return z, (-3, 4, -4, 3)
 
diff --git a/examples/axes_grid1/inset_locator_demo2.py b/examples/axes_grid1/inset_locator_demo2.py
@@ -9,19 +9,16 @@
 created via `~.mark_inset`.
 """
 
+from matplotlib import cbook
 import matplotlib.pyplot as plt
-
 from mpl_toolkits.axes_grid1.inset_locator import zoomed_inset_axes, mark_inset
 from mpl_toolkits.axes_grid1.anchored_artists import AnchoredSizeBar
 
 import numpy as np
 
 
 def get_demo_image():
-    from matplotlib.cbook import get_sample_data
-    import numpy as np
-    f = get_sample_data("axes_grid/bivariate_normal.npy", asfileobj=False)
-    z = np.load(f)
+    z = cbook.get_sample_data("axes_grid/bivariate_normal.npy", np_load=True)
     # z is a numpy array of 15x15
     return z, (-3, 4, -4, 3)
 
@@ -60,11 +57,9 @@ def add_sizebar(ax, size):
 ny, nx = Z.shape
 Z2[30:30+ny, 30:30+nx] = Z
 
-# extent = [-3, 4, -4, 3]
 ax2.imshow(Z2, extent=extent, origin="lower")
 
-
-axins2 = zoomed_inset_axes(ax2, 6, loc=1)  # zoom = 6
+axins2 = zoomed_inset_axes(ax2, zoom=6, loc=1)
 axins2.imshow(Z2, extent=extent, origin="lower")
 
 # sub region of the original image
diff --git a/examples/axes_grid1/simple_axesgrid2.py b/examples/axes_grid1/simple_axesgrid2.py
@@ -6,28 +6,21 @@
 Align multiple images of different sizes using
 `~mpl_toolkits.axes_grid1.axes_grid.ImageGrid`.
 """
+
+from matplotlib import cbook
 import matplotlib.pyplot as plt
 from mpl_toolkits.axes_grid1 import ImageGrid
 
 
-def get_demo_image():
-    import numpy as np
-    from matplotlib.cbook import get_sample_data
-    f = get_sample_data("axes_grid/bivariate_normal.npy", asfileobj=False)
-    z = np.load(f)
-    # z is a numpy array of 15x15
-    return z, (-3, 4, -4, 3)
-
-
 fig = plt.figure(figsize=(5.5, 3.5))
 grid = ImageGrid(fig, 111,  # similar to subplot(111)
                  nrows_ncols=(1, 3),
                  axes_pad=0.1,
                  label_mode="L",
                  )
 
-Z, extent = get_demo_image()  # demo image
-
+# demo image
+Z = cbook.get_sample_data("axes_grid/bivariate_normal.npy", np_load=True)
 im1 = Z
 im2 = Z[:, :10]
 im3 = Z[:, 10:]
diff --git a/examples/frontpage/3D.py b/examples/frontpage/3D.py
@@ -12,13 +12,12 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
-with cbook.get_sample_data('jacksboro_fault_dem.npz') as file, \
-     np.load(file) as dem:
-    z = dem['elevation']
-    nrows, ncols = z.shape
-    x = np.linspace(dem['xmin'], dem['xmax'], ncols)
-    y = np.linspace(dem['ymin'], dem['ymax'], nrows)
-    x, y = np.meshgrid(x, y)
+dem = cbook.get_sample_data('jacksboro_fault_dem.npz', np_load=True)
+z = dem['elevation']
+nrows, ncols = z.shape
+x = np.linspace(dem['xmin'], dem['xmax'], ncols)
+y = np.linspace(dem['ymin'], dem['ymax'], nrows)
+x, y = np.meshgrid(x, y)
 
 region = np.s_[5:50, 5:50]
 x, y, z = x[region], y[region], z[region]
diff --git a/examples/images_contours_and_fields/shading_example.py b/examples/images_contours_and_fields/shading_example.py
@@ -21,9 +21,8 @@ def main():
     x, y = np.mgrid[-5:5:0.05, -5:5:0.05]
     z = 5 * (np.sqrt(x**2 + y**2) + np.sin(x**2 + y**2))
 
-    with cbook.get_sample_data('jacksboro_fault_dem.npz') as file, \
-         np.load(file) as dem:
-        elev = dem['elevation']
+    dem = cbook.get_sample_data('jacksboro_fault_dem.npz', np_load=True)
+    elev = dem['elevation']
 
     fig = compare(z, plt.cm.copper)
     fig.suptitle('HSV Blending Looks Best with Smooth Surfaces', y=0.95)
diff --git a/examples/lines_bars_and_markers/scatter_demo2.py b/examples/lines_bars_and_markers/scatter_demo2.py
@@ -12,8 +12,8 @@
 # Load a numpy record array from yahoo csv data with fields date, open, close,
 # volume, adj_close from the mpl-data/example directory. The record array
 # stores the date as an np.datetime64 with a day unit ('D') in the date column.
-with cbook.get_sample_data('goog.npz') as datafile:
-    price_data = np.load(datafile)['price_data'].view(np.recarray)
+price_data = (cbook.get_sample_data('goog.npz', np_load=True)['price_data']
+              .view(np.recarray))
 price_data = price_data[-250:]  # get the most recent 250 trading days
 
 delta1 = np.diff(price_data.adj_close) / price_data.adj_close[:-1]
diff --git a/examples/mplot3d/custom_shaded_3d_surface.py b/examples/mplot3d/custom_shaded_3d_surface.py
@@ -13,13 +13,12 @@
 import numpy as np
 
 # Load and format data
-with cbook.get_sample_data('jacksboro_fault_dem.npz') as file, \
-     np.load(file) as dem:
-    z = dem['elevation']
-    nrows, ncols = z.shape
-    x = np.linspace(dem['xmin'], dem['xmax'], ncols)
-    y = np.linspace(dem['ymin'], dem['ymax'], nrows)
-    x, y = np.meshgrid(x, y)
+dem = cbook.get_sample_data('jacksboro_fault_dem.npz', np_load=True)
+z = dem['elevation']
+nrows, ncols = z.shape
+x = np.linspace(dem['xmin'], dem['xmax'], ncols)
+y = np.linspace(dem['ymin'], dem['ymax'], nrows)
+x, y = np.meshgrid(x, y)
 
 region = np.s_[5:50, 5:50]
 x, y, z = x[region], y[region], z[region]
diff --git a/examples/recipes/common_date_problems.py b/examples/recipes/common_date_problems.py
@@ -38,8 +38,8 @@
 import numpy as np
 import matplotlib.pyplot as plt
 
-with cbook.get_sample_data('goog.npz') as datafile:
-    r = np.load(datafile)['price_data'].view(np.recarray)
+r = (cbook.get_sample_data('goog.npz', np_load=True)['price_data']
+     .view(np.recarray))
 
 fig, ax = plt.subplots()
 ax.plot(r.date, r.close)
diff --git a/examples/recipes/fill_between_alpha.py b/examples/recipes/fill_between_alpha.py
@@ -21,8 +21,8 @@
 np.random.seed(19680801)
 
 # load up some sample financial data
-with cbook.get_sample_data('goog.npz') as datafile:
-    r = np.load(datafile)['price_data'].view(np.recarray)
+r = (cbook.get_sample_data('goog.npz', np_load=True)['price_data']
+     .view(np.recarray))
 # create two subplots with the shared x and y axes
 fig, (ax1, ax2) = plt.subplots(1, 2, sharex=True, sharey=True)
 
diff --git a/examples/specialty_plots/topographic_hillshading.py b/examples/specialty_plots/topographic_hillshading.py
@@ -22,23 +22,21 @@
 from matplotlib.colors import LightSource
 
 
-with np.load(get_sample_data('jacksboro_fault_dem.npz')) as dem:
-    z = dem['elevation']
-
-    #-- Optional dx and dy for accurate vertical exaggeration ----------------
-    # If you need topographically accurate vertical exaggeration, or you don't
-    # want to guess at what *vert_exag* should be, you'll need to specify the
-    # cellsize of the grid (i.e. the *dx* and *dy* parameters).  Otherwise, any
-    # *vert_exag* value you specify will be relative to the grid spacing of
-    # your input data (in other words, *dx* and *dy* default to 1.0, and
-    # *vert_exag* is calculated relative to those parameters).  Similarly, *dx*
-    # and *dy* are assumed to be in the same units as your input z-values.
-    # Therefore, we'll need to convert the given dx and dy from decimal degrees
-    # to meters.
-    dx, dy = dem['dx'], dem['dy']
-    dy = 111200 * dy
-    dx = 111200 * dx * np.cos(np.radians(dem['ymin']))
-    #-------------------------------------------------------------------------
+dem = get_sample_data('jacksboro_fault_dem.npz', np_load=True)
+z = dem['elevation']
+#-- Optional dx and dy for accurate vertical exaggeration ---------------------
+# If you need topographically accurate vertical exaggeration, or you don't want
+# to guess at what *vert_exag* should be, you'll need to specify the cellsize
+# of the grid (i.e. the *dx* and *dy* parameters).  Otherwise, any *vert_exag*
+# value you specify will be relative to the grid spacing of your input data
+# (in other words, *dx* and *dy* default to 1.0, and *vert_exag* is calculated
+# relative to those parameters).  Similarly, *dx* and *dy* are assumed to be in
+# the same units as your input z-values.  Therefore, we'll need to convert the
+# given dx and dy from decimal degrees to meters.
+dx, dy = dem['dx'], dem['dy']
+dy = 111200 * dy
+dx = 111200 * dx * np.cos(np.radians(dem['ymin']))
+#------------------------------------------------------------------------------
 
 # Shade from the northwest, with the sun 45 degrees from horizontal
 ls = LightSource(azdeg=315, altdeg=45)
diff --git a/examples/subplots_axes_and_figures/zoom_inset_axes.py b/examples/subplots_axes_and_figures/zoom_inset_axes.py
@@ -4,18 +4,15 @@
 ======================
 
 Example of an inset axes and a rectangle showing where the zoom is located.
-
 """
 
+from matplotlib import cbook
 import matplotlib.pyplot as plt
 import numpy as np
 
 
 def get_demo_image():
-    from matplotlib.cbook import get_sample_data
-    import numpy as np
-    f = get_sample_data("axes_grid/bivariate_normal.npy", asfileobj=False)
-    z = np.load(f)
+    z = cbook.get_sample_data("axes_grid/bivariate_normal.npy", np_load=True)
     # z is a numpy array of 15x15
     return z, (-3, 4, -4, 3)
 
diff --git a/examples/text_labels_and_annotations/date.py b/examples/text_labels_and_annotations/date.py
@@ -29,8 +29,7 @@
 # close, volume, adj_close from the mpl-data/example directory.  This array
 # stores the date as an np.datetime64 with a day unit ('D') in the 'date'
 # column.
-with cbook.get_sample_data('goog.npz') as datafile:
-    data = np.load(datafile)['price_data']
+data = cbook.get_sample_data('goog.npz', np_load=True)['price_data']
 
 fig, ax = plt.subplots()
 ax.plot('date', 'adj_close', data=data)
diff --git a/examples/text_labels_and_annotations/date_index_formatter.py b/examples/text_labels_and_annotations/date_index_formatter.py
@@ -14,8 +14,8 @@
 # Load a numpy record array from yahoo csv data with fields date, open, close,
 # volume, adj_close from the mpl-data/example directory. The record array
 # stores the date as an np.datetime64 with a day unit ('D') in the date column.
-with cbook.get_sample_data('goog.npz') as datafile:
-    r = np.load(datafile)['price_data'].view(np.recarray)
+r = (cbook.get_sample_data('goog.npz', np_load=True)['price_data']
+     .view(np.recarray))
 r = r[-30:]  # get the last 30 days
 
 # first we'll do it the default way, with gaps on weekends
diff --git a/examples/ticks_and_spines/centered_ticklabels.py b/examples/ticks_and_spines/centered_ticklabels.py
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
diff --git a/lib/mpl_toolkits/tests/test_axes_grid1.py b/lib/mpl_toolkits/tests/test_axes_grid1.py
diff --git a/tutorials/colors/colormapnorms.py b/tutorials/colors/colormapnorms.py
