diff --git a/demo.py b/demo.py
index 3105f51..c72d0a8 100644
--- a/demo.py
+++ b/demo.py
@@ -12,7 +12,7 @@
 '''
 
 from simulator import Simulator
-from turbStats import tilt_mat, corr_mat
+from turbStats import tilt_mat, corr_mat, get_r0
 import matplotlib.pyplot as plt
 import torch
 
@@ -39,11 +39,19 @@
     x = x.transpose((2,0,1))
 x = torch.tensor(x, device = device, dtype=torch.float32)
 
+D = 0.1
+L = 3000
+# r0 can also be calculated from Cn2 and L using the get_r0 function: 
+# r0 = get_r0(Cn2, L)
+r0 = 0.05
+L = 3000
+im_size = x.shape[1] # The current version works for square image only
+
 # Generate correlation matrix for tilt. Do this once for each different turbulence parameter. 
-tilt_mat(x.shape[1], 0.1, 0.05, 3000)
+tilt_mat(im_size, D, r0, L)
 
 # Simulate
-simulator = Simulator(2, 512).to(device, dtype=torch.float32)
+simulator = Simulator(D/r0, im_size).to(device, dtype=torch.float32)
 
 out = simulator(x).detach().cpu().numpy()