figure source code

Dear Song,

Many thanks for sharing this awesome idea and the source code! I am wondering whether you can point me to the source code of the following figure in [Song2021](https://arxiv.org/pdf/2011.13456.pdf)?

![image](https://user-images.githubusercontent.com/8034156/176836324-05ce5049-f8f1-4c3c-9dce-c5cffbdda1ec.png)

Kind regards
Feng