简体   繁体   中英

Matplotlib - mark_inset with different edges for axes

I want to plot a time series of a damped random walk in one subplot and then zoom into it in a second subplot. I know mark_inset from matplotlib, which works fine. The code I have so far is:

from mpl_toolkits.axes_grid1.inset_locator import mark_inset
from astroML.time_series import generate_damped_RW

fig = plt.figure()
ax = fig.add_subplot(111)
ax0 = fig.add_subplot(211)
ax1 = fig.add_subplot(212)

ax.set_ylabel('Brightness[mag]')
ax.yaxis.labelpad=30
ax.spines['top'].set_color('none')
ax.spines['bottom'].set_color('none')
ax.spines['left'].set_color('none')
ax.spines['right'].set_color('none')
ax.tick_params(labelcolor='w', top='off', bottom='off', left='off',
               right='off')

t = np.linspace(0, 5000, 100000)
data = generate_damped_RW(t, tau=100, xmean=20, z=0, SFinf=0.3,
                          random_state=1)
ax0.scatter(t, data, s=0.5)
ax0.text(1, 1, r'$E(m) = %.2f, \sigma(m) = %.2f$'%(np.mean(data),
                                                   np.std(data)),
         verticalalignment='top', horizontalalignment='right',
         transform=ax0.transAxes, fontsize=23)

mask = (t > 370) & (t < 470)
ax1.set_xlabel('Time[years]')
ax1.scatter(t[mask], data[mask], s=0.5)

mark_inset(ax0, ax1, loc1=2, loc=1, fc='none')

which creates a plot like this: 在此处输入图片说明

Which is almost what I want, except that the lines connecting the 2 subplots start at the upper edges of the box in the first subplot. Is it possible to have those start at the lower two edges while they still end up at the upper two in the second subplot? What would I have to do to achieve this?

The mark_inset has two arguments loc1 and loc2 to set the locations of the two connectors. Those locations are then the same for the box and and the inset axes.

We may however add two new arguments to the mark_inset function to set different locations for the start and end of the connector.

import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.inset_locator import TransformedBbox, BboxPatch, BboxConnector 
import numpy as np

fig, (ax, axins) = plt.subplots(nrows=2)

x = np.linspace(0,6*np.pi)
y = np.sin(x)
ax.plot(x,y)
axins.plot(x,y)
axins.set_xlim((2*np.pi, 2.5*np.pi))
axins.set_ylim((0, 1))

# draw a bbox of the region of the inset axes in the parent axes and
# connecting lines between the bbox and the inset axes area
# loc1, loc2 : {1, 2, 3, 4} 
def mark_inset(parent_axes, inset_axes, loc1a=1, loc1b=1, loc2a=2, loc2b=2, **kwargs):
    rect = TransformedBbox(inset_axes.viewLim, parent_axes.transData)

    pp = BboxPatch(rect, fill=False, **kwargs)
    parent_axes.add_patch(pp)

    p1 = BboxConnector(inset_axes.bbox, rect, loc1=loc1a, loc2=loc1b, **kwargs)
    inset_axes.add_patch(p1)
    p1.set_clip_on(False)
    p2 = BboxConnector(inset_axes.bbox, rect, loc1=loc2a, loc2=loc2b, **kwargs)
    inset_axes.add_patch(p2)
    p2.set_clip_on(False)

    return pp, p1, p2

mark_inset(ax, axins, loc1a=1, loc1b=4, loc2a=2, loc2b=3, fc="none", ec="crimson") 

plt.draw()
plt.show()

在此处输入图片说明

Unfortunately, mark_inset always has to connect the same corners (ie bottom right always has to connect to bottom right, etc.).

We can make our own function that mimics the mark_inset function though, to connect the two bottom corners with the two top corners in the inset ( custom_mark_inset in the code below).

This makes use of a Rectangle patch to draw the box on the primary axes, and the ConnectionPatch instances to draw the connecting lines between axes.

from mpl_toolkits.axes_grid1.inset_locator import mark_inset
#from astroML.time_series import generate_damped_RW
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np

fig = plt.figure()
ax = fig.add_subplot(111)
ax0 = fig.add_subplot(211)
ax1 = fig.add_subplot(212)

ax.set_ylabel('Brightness[mag]')
ax.yaxis.labelpad=30
ax.spines['top'].set_color('none')
ax.spines['bottom'].set_color('none')
ax.spines['left'].set_color('none')
ax.spines['right'].set_color('none')
ax.tick_params(labelcolor='w', top='off', bottom='off', left='off',
               right='off')

t = np.linspace(0, 5000, 10000)
#data = generate_damped_RW(t, tau=100, xmean=20, z=0, SFinf=0.3,
#                          random_state=1)
## Fake some data
data = np.sin(t/800.) + 20.

ax0.scatter(t, data, s=0.5)
ax0.text(1, 1, r'$E(m) = %.2f, \sigma(m) = %.2f$'%(np.mean(data),
                                                   np.std(data)),
         verticalalignment='top', horizontalalignment='right',
         transform=ax0.transAxes, fontsize=23)

mask = (t > 370) & (t < 470)
ax1.set_xlabel('Time[years]')
ax1.scatter(t[mask], data[mask], s=0.5)

def custom_mark_inset(axA, axB, fc='None', ec='k'):
    xx = axB.get_xlim()
    yy = axB.get_ylim()

    xy = (xx[0], yy[0])
    width = xx[1] - xx[0]
    height = yy[1] - yy[0]

    pp = axA.add_patch(patches.Rectangle(xy, width, height, fc=fc, ec=ec))

    p1 = axA.add_patch(patches.ConnectionPatch(
        xyA=(xx[0], yy[0]), xyB=(xx[0], yy[1]),
        coordsA='data', coordsB='data',
        axesA=axA, axesB=axB))

    p2 = axA.add_patch(patches.ConnectionPatch(
        xyA=(xx[1], yy[0]), xyB=(xx[1], yy[1]),
        coordsA='data', coordsB='data',
        axesA=axA, axesB=axB))

    return pp, p1, p2

pp, p1, p2 = custom_mark_inset(ax0, ax1)

plt.show()

在此处输入图片说明

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM