Update test fit script to save figures