Added plotting to optimization