import numpy as np
from astropy.visualization import SqrtStretch
from matplotlib import pyplot as plt

fig, ax = plt.subplots(figsize=(5, 5))

x = np.linspace(0, 1, 100)
stretch = SqrtStretch()
ax.plot(x, stretch(x, clip=True))

ax.axis('equal')
ax.plot(x, x, ls='dotted', color='k', alpha=0.3)
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.set_xlabel('Input Value')
ax.set_ylabel('Output Value')
ax.set_title(stretch.__class__.__name__)