xs = np.linspace(-1, sum(strata.Ls[1:]), 5000)
nx = strata.populateIndices(xs)
Ey, _, _ = strata.populateMode(beta, xs)
# confinement = strata.confinementy(beta, xs, Ey)

rIdxAxes = plt.gca()
modeAxes = rIdxAxes.twinx()
modeAxes.set_frame_on(False)
modeAxes.get_yaxis().set_visible(False)
rIdxAxes.set_xlabel('Position (μm)')
rIdxAxes.set_ylabel('Refractive index $n$')

lNReal = rIdxAxes.plot(xs, nx.real, 'k', label=r'$\mathrm{Re}[n]$')
lNImag = rIdxAxes.plot(xs, nx.imag, 'orange', label=r'$\mathrm{Im}[n]$')
lMode = modeAxes.plot(xs, np.abs(Ey)**2, color='C0', label=r'Mode')

lns = lNReal + lNImag + lMode
labs = [line.get_label() for line in lns]
plt.legend(lns, labs)
rIdxAxes.set_ylim(0, 5)
rIdxAxes.set_xlim(-0.2, 11)
modeAxes.set_ylim(bottom=0)

plt.show()