import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm
from numpy.random import seed, normal
from astropy.modeling.models import Gaussian1D
from astropy.modeling.fitting import DogBoxLSQFitter
plt.rc('font', size=13)
seed(5)

fw = 10
nd, ncd = 31, 13
xd, xcd = np.arange(nd), np.arange(ncd)

psf = norm(5.4, 1.5).pdf(xcd) + normal(0, 0.01, ncd)
fitter = DogBoxLSQFitter()
m = fitter(Gaussian1D(), xcd, psf)

fig, ax = plt.subplots(figsize=(fw, fw*(ncd/nd)), constrained_layout=True)
ax.step(xcd, psf, where='mid', c='k')
ax.axvline(xcd[np.argmax(psf)], label='max')
ax.axvline(np.average(xcd, weights=psf), ls='--', label='centroid')
ax.axvline(m.mean.value, ls=':', label='gaussian')
ax.plot(xcd, m(xcd), ls=':')
ax.legend()
plt.setp(ax, yticks=[], ylabel='Flux', xlabel='Cross-dispersion axis [pix]', xlim=(0, ncd-1))
fig.show()