Improved example diamonds.py