How to edit this code to have the same width and colour map as in the following figure? The script is based on this question.
import numpy as np impor
