diff --git a/compressed_sensing.py b/compressed_sensing.py
index 33e0e9df404baa9eff0102d89a55a9fb0f699cec..360d74a671c9a79772ef42a1755212828f455f4b 100644
--- a/compressed_sensing.py
+++ b/compressed_sensing.py
@@ -14,17 +14,14 @@ if __name__ == "__main__":
     n_samples = 1000
     sample_times = (sample_period / n_samples) * np.arange(n_samples)
     sample_values = sample_two_sins(f1, f2, sample_times)
-    fig, ax = plt.subplots()
-    ax.plot(sample_times, sample_values)
+    plt.plot(sample_times, sample_values)
+    plt.savefig("fig_a.png")
+    plt.close()
 
     sample_period = 0.02
     n_samples = 1000
     sample_times = (sample_period / n_samples) * np.arange(n_samples)
     sample_values = sample_two_sins(f1, f2, sample_times)
-    fig2, ax2 = plt.subplots()
-    ax2.plot(sample_times, sample_values)
-
-    plt.figure(fig.number)
-    plt.savefig("test.png")
-    plt.figure(fig2.number)
-    plt.savefig("test2.png")
+    plt.plot(sample_times, sample_values)
+    plt.savefig("fig_b.png")
+    plt.close()