diff --git a/binarycpython/utils/grid.py b/binarycpython/utils/grid.py
index f0937155683ddcc5da5d35c3927fdd9a2bf1b659..3a324eacd3665f5e1600140bfc4d9bfe18e16b4b 100644
--- a/binarycpython/utils/grid.py
+++ b/binarycpython/utils/grid.py
@@ -122,7 +122,7 @@ class Population(object):
 
         pass
 
-    def add_grid_variable(self, name, longname, valuerange, resolution, spacingfunc, precode, probdist, dphasevol, condition=None):
+    def add_grid_variable(self, name, longname, valuerange, resolution, spacingfunc, precode, probdist, dphasevol, parameter_name, condition=None):
         """
         Function to add grid variables to the grid_options.
 
@@ -170,6 +170,7 @@ class Population(object):
             "precode": precode,
             "probdist": probdist,
             "dphasevol": dphasevol,
+            "parameter_name": parameter_name,
             "condition": condition,
             "grid_variable_number": len(self.grid_options["grid_variables"]),
         }
@@ -454,33 +455,19 @@ class Population(object):
 
         # TODO: add imports
         # TODO: 
+        code_string += "from binarycpython.utils.probability_distributions import *\n"  
+        code_string += "import math\n"
+        code_string += "import numpy as np\n"
+        code_string += "\n\n"
 
         code_string += "starcount = 0\n"
-        code_string += "probsum = []\n"
-
+        code_string += "probabilities = {}\n"
+        code_string += "parameter_dict = {}\n"
         
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
+        # Prepare the probability
+        for el in sorted(self.grid_options["grid_variables"].items(), key=lambda x: x[1]['grid_variable_number']):
+            grid_variable = el[1]
+            code_string += 'probabilities["{}"] = 0\n'.format(grid_variable['parameter_name'])
 
         # Generate code
         print("Generating grid code")
@@ -489,29 +476,36 @@ class Population(object):
             grid_variable = el[1]
 
             # Adding for loop structure
-            code_string += indent * depth + 'for {} in range({}, {}):'.format(grid_variable['name'], grid_variable['valuerange'][0], grid_variable['valuerange'][1]) + '\n'
+            code_string += indent * depth + 'for {} in {}:'.format(grid_variable['name'], grid_variable['spacingfunc']) + '\n'
 
             # Add pre-code
             code_string += indent * (depth + 1) + '{}'.format(grid_variable['precode'].replace('\n', '\n'+indent * (depth + 1))) + '\n'
 
             # Calculate probability
-
+            code_string += indent * (depth + 1) + 'probabilities["{}"] = {}'.format(grid_variable['parameter_name'], grid_variable['probdist']) + '\n'
 
             # some testing line.
-            code_string += indent * (depth + 1) + 'print({})'.format(grid_variable['name']) + '\n'
+            # code_string += indent * (depth + 1) + 'print({})'.format(grid_variable['name']) + '\n'
+
+            # Add value to dict
+            code_string += indent * (depth + 1) + 'parameter_dict["{}"] = {}'.format(grid_variable['parameter_name'], grid_variable['parameter_name']) + '\n'
+
 
             # Add some space
             code_string += "\n"
 
-            # 
+            # increment depth
             depth += 1
 
+        # placeholder for calls to threading
+        code_string += indent * (depth) + 'print(parameter_dict)\n'
+        # starcount
+        code_string += indent * (depth) + 'starcount += 1\n'
+        code_string += indent * (depth) + 'print(probabilities)\n'
+        code_string += indent * (depth) + 'print("starcount: ", starcount)\n'
 
         # Write to file
         gridcode_filename = os.path.join(temp_custom_logging_dir(), 'example_grid.py')
         with open(gridcode_filename, 'w') as f:
             f.write(code_string)
-
-
-
 ################################################################################################
diff --git a/binarycpython/utils/probability_distributions.py b/binarycpython/utils/probability_distributions.py
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..57f8b6b32ffd2340c19b5ddc23cae035b4de1bd2 100644
--- a/binarycpython/utils/probability_distributions.py
+++ b/binarycpython/utils/probability_distributions.py
@@ -0,0 +1,2 @@
+def flat(parameter):
+    return 1
\ No newline at end of file
diff --git a/binarycpython/utils/spacingsfunctions.py b/binarycpython/utils/spacingsfunctions.py
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..53477b4d4bf25a8e84d1555a4d334f27ee9c5fa1 100644
--- a/binarycpython/utils/spacingsfunctions.py
+++ b/binarycpython/utils/spacingsfunctions.py
@@ -0,0 +1,2 @@
+def const():
+    return 
\ No newline at end of file
diff --git a/tests/population/grid_tests.py b/tests/population/grid_tests.py
index f78b7d7afabd76ae3aa009e9b4f5f24ed2577d25..e99807189837ddf31c35a8fe2d3dea2869212808 100644
--- a/tests/population/grid_tests.py
+++ b/tests/population/grid_tests.py
@@ -210,9 +210,9 @@ test_pop.add_grid_variable(
     longname='log primary mass',
     valuerange=[10, 20],
     resolution='10',
-    spacingfunc='const()',
-    precode='m2=m1**2',
-    probdist='flat()',
+    spacingfunc='np.linspace(0.213, 10.2, 10)',
+    precode='M_1=math.exp(lnm1)',
+    probdist='flat(M_1)',
     dphasevol='',
     parameter_name='M_1',
     condition='',
@@ -222,11 +222,11 @@ test_pop.add_grid_variable(
 test_pop.add_grid_variable(
     name='period',
     longname='period',
-    valuerange=['m1', 20],
+    valuerange=['M_1', 20],
     resolution='10',
-    spacingfunc='const()',
-    precode='',
-    probdist='flat()',
+    spacingfunc='np.linspace(1, 10, 10)',
+    precode='orbital_period = period**2',
+    probdist='flat(orbital_period)',
     parameter_name='orbital_period',
     dphasevol='',
     condition='',