improved the plotting slightly
switched to a PSD solver from np.linalg.solve