Changing x-axis scale piecewise in matplotlib

I've implemented next helper function plot_xrescaled(...) in order to solve your tas (see visualizing image at the end of my answer).

As far as I know there is no such built-in function in matplotlib that does this piece-wise rescaling.

Function accepts x and y 1D arrays plus xscale which contains pairs of (x0, scale) where x0 is beginning of range to be scaled and scale is scaling factor (scale of 1 is unscaled x, same as input, scale of 0.5 compresses x twice). Each range of x is scaled starting from x0 position till next x0. E.g. if you have xscale = [(0., 1.), (2 * np.pi, 0.5), (4 * np.pi, 2.)] then x range (0;2pi) is scaled at factor of 1x (unscaled), then range (2pi;4pi) is scaled 0.5. then (4pi;till_end) is scaled 2x.

Also I output x ticks with all necessary scaling considered. In my function argument nxticks controls how many x ticks (points) are there on x axis showed, it is approximate, sometimes a bit less ticks is shown. Also beginning and end of each scaled range is shown as x tick starting with # symbol. Also there's argument xround which controls how many digits after decimal points are shown in x ticks e.g. if xround = 3 then x value 2.3456789 will be shown as 2.345 (i.e. rounded to 3 decimal digits).

In my code I test my function sin(x). See output plot below to understand how my code works for this function and 3 scaled ranges.

Try it online!

def plot_xrescaled(x, y, xscale = [], *, nxticks = 12, xround = 2):
    import matplotlib, matplotlib.pyplot as plt, numpy as np
    def check_sorted(a):
        assert np.array_equal(np.arange(a.size), np.argsort(a))
    if len(xscale) == 0:
        xscale = np.zeros((0, 2), dtype = np.float32)
    else:
        xscale = np.array(xscale)
    check_sorted(xscale[:, 0])
    check_sorted(x)
    if len(xscale) == 0 or xscale[0, 0] > x[0]:
        xscale = np.concatenate((np.array([(x[0], 1)]), xscale))
    iscales = np.searchsorted(xscale[:, 0], x, side = 'right') - 1
    assert iscales[0] >= 0
    px1, nx, xt, xtl = xscale[0, 0], np.zeros_like(x), [], []
    for i in range(xscale.shape[0]):
        (x0, sc), b = xscale[i], iscales == i
        x1 = xscale[i + 1, 0] if i + 1 < xscale.shape[0] else x[-1]
        nx[b] = px1 + (x[b] - x0) * sc
        npx1 = px1 + (x1 - x0) * sc
        if np.count_nonzero(b) > 0:
            xt += [px1, npx1]
            xtl += [x0, x1]
        px1 = npx1
    nxt, nxtl, xtb, xts = [], [], [], (xt[-1] - xt[0]) / nxticks
    for b, e, bl, el in zip(xt[0::2], xt[1::2], xtl[0::2], xtl[1::2]):
        c = b
        while c <= e:
            if not (10 ** -6 < e - c < 0.5 * xts):
                nxt += [c]
                nxtl += [str(round(bl + (el - bl) * (c - b) / (e - b), xround))]
                if c in (b, e):
                    nxtl[-1] = '#' + nxtl[-1]
            if c == e:
                break
            c = min(c + xts, e)
    xt, xtl = nxt, nxtl
    plt.plot(nx, y)
    plt.xticks(xt, xtl)
    plt.show()

# Test
import numpy as np
xscale = [(0., 1.), (2 * np.pi, 0.5), (4 * np.pi, 2.)]
x = np.linspace(0, 6 * np.pi, 200)
y = np.sin(x)
plot_xrescaled(x, y, xscale)

Output plot: