update to code gen 2

examples/code-generation/Neuropiler/examples/example_LeakyIAF.py

@@ -1,4 +1,4 @@
-class LeakyIAF():
+class LeakyIAF(BaseAxonHillockModel):
     updates = ['spike_state', # (bool)
                'V' # Membrane Potential (mV)
               ]
@@ -21,4 +21,4 @@ def step():
             V = V + I * dt 
             if V > threshold:
                 V = reset_potential
-                spike = 1
+                spike_state = 1

examples/code-generation/Neuropiler/neuropile.py

@@ -2,7 +2,7 @@
 import ast
 import astunparse
 from python2c.python2c import *
-from python2c.translate import main_function
+from python2c.translate import main_function, prettyparseprint
 from collections import OrderedDict
 
 
@@ -15,7 +15,8 @@ def neuropile_main(file):
         if isinstance(node, ast.FunctionDef) or isinstance(node, ast.ClassDef):
             if isinstance(node, ast.ClassDef):
                 className = node.name
-                neuropiler_classes[className] = {'assignments': [], 'functions': {}}
+                model_base = node.bases[0].id
+                neuropiler_classes[className] = {'assignments': [], 'functions': {},'model_base': model_base}
                 for i in node.body:
                     if isinstance(i, ast.Assign):
                         neuropiler_classes[className]['assignments'].append(astunparse.unparse(i).replace('\n',''))
@@ -27,15 +28,28 @@ def neuropile_main(file):
 
 
 if __name__ == "__main__":
+    import neurokernel
+    import inspect
+
     parser = ArgumentParser(description="Translate Python code to C.")
     parser.add_argument("file", help=".py file to translate.")
     arguments = parser.parse_args()
     results = neuropile_main(arguments.file)
+
     for model_name in results.keys():
+        model_base = results[model_name]['model_base']
+        path = inspect.getfile(neurokernel).replace('__init__.py','')+'LPU/NDComponents/'+model_base[4:]+'s/'
         if 'step' in results[model_name]['functions']:
-            with open(model_name+".py", "w") as python_code:
+        	# write step function
+            with open(model_name+"_temp.py", "w") as python_code:
                 python_code.write(results[model_name]['functions']['step'])
+            # evaluate the assignments
             main_func_block = eval('main_function(' + (','.join(results[model_name]['assignments'])) + ')')
-            translated_code = translate.translate(model_name+".py", indent_size=4, main_func = main_func_block)
+            # translate the iterative update
+            translated_code = translate.translate(model_name+"_temp.py", indent_size=4, main_func = main_func_block)
+            # write the whole template
             with open(model_name+"_template.c", "w") as cuda_code:
-                cuda_code.write(translated_code)
+                cuda_code.write(translated_code)
+            wrapped_code = translate.wrapper(model_name+"_template.c", indent_size=4, model=model_name, model_base = model_base, assign = results[model_name]['assignments'])
+            with open(path+model_name+"_gen.py", "w") as python_code:
+                python_code.write(wrapped_code)

examples/code-generation/Neuropiler/python2c/translate.py

@@ -180,9 +180,15 @@ def get_op(op, arg1, arg2=None):
     elif isinstance(op, ast.BitAnd):
         return "(({}) & ({}))".format(arg1, arg2)
     elif isinstance(op, ast.USub):
-        return "(-{})".format(arg1.id)
+        if isinstance(arg1,ast.BinOp):
+            return "(-{})".format(handle_op_node(arg1))
+        else:
+            return "(-{})".format(arg1.id)
     elif isinstance(op, ast.UAdd):
-        return "(+{})".format(arg1.id)
+        if isinstance(arg1,ast.BinOp):
+            return "(+{})".format(handle_op_node(arg1))
+        else:
+            return "(+{})".format(arg1.id)
     raise Exception("Could not identify operator " + str(op))
 
 
@@ -197,7 +203,7 @@ def handle_op_node(node):
         if isinstance(node.left, ast.UnaryOp):
             node_left = handle_op_node(node.left)
         if isinstance(node.left, ast.Call):
-            arguments = '(' + ','.join([handle_op_node(i) for i in node.left.args]) + ')'
+            arguments = '(' + ','.join([str(handle_op_node(i)) for i in node.left.args]) + ')'
             node_left = node.left.func.id+'%(fletter)s' + arguments
         else:
             print(node.left)
@@ -212,7 +218,7 @@ def handle_op_node(node):
         if isinstance(node.right, ast.UnaryOp):
             node_right = handle_op_node(node.right)
         if isinstance(node.right, ast.Call):
-            arguments = '(' + ','.join([handle_op_node(i) for i in node.right.args]) + ')'
+            arguments = '(' + ','.join([str(handle_op_node(i)) for i in node.right.args]) + ')'
             node_right = node.right.func.id+'%(fletter)s' + arguments
         else:
             print(node.right)
@@ -223,10 +229,10 @@ def handle_op_node(node):
     if isinstance(node, ast.UnaryOp):
         return get_op(node.op, node.operand)
     elif isinstance(node, ast.Call):
-        arguments = '(' + ','.join([handle_op_node(i) for i in node.args]) + ')'
+        arguments = '(' + ','.join([str(handle_op_node(i)) for i in node.args]) + ')'
         return node.func.id +'%(fletter)s'+ arguments
     elif isinstance(node, ast.Compare):
-        return handle_op_node(node.left) + ''.join([handle_op_node(i) for i in node.ops]) + ''.join([handle_op_node(i) for i in node.comparators])
+        return handle_op_node(node.left) + ''.join([str(handle_op_node(i)) for i in node.ops]) + ''.join([str(handle_op_node(i)) for i in node.comparators])
     elif isinstance(node, ast.Num):
         return node.n
     elif isinstance(node, ast.Str):
@@ -258,7 +264,7 @@ def handle_op_node(node):
     elif isinstance(node, ast.Or):
         return '||'
     elif isinstance(node, ast.BoolOp):
-        return (handle_op_node(node.op)).join(['(' + handle_op_node(i)+ ')' for i in node.values])
+        return (handle_op_node(node.op)).join(['(' + str(handle_op_node(i))+ ')' for i in node.values])
     # print(node.left, node.right)
     raise Exception("Could not identify node op")
 
@@ -393,7 +399,6 @@ def translate(file_, indent_size=4, main_func = None):
     # Setup
     with open(file_, "r") as f:
         text = f.read()
-        # replace occurences of dt with ddt
         text = text.replace("dt","ddt")
         nodes = ast.parse(text).body
         code = text.splitlines()
@@ -418,7 +423,6 @@ def translate(file_, indent_size=4, main_func = None):
 
     return str(top)
 
-
 def import_block(top=None,model_base=None):
     if model_base=='BaseAxonHillockModel':
         top.append_block(blocks.StringBlock('from neurokernel.LPU.NDComponents.AxonHillockModels.BaseAxonHillockModel import *'))
@@ -427,7 +431,6 @@ def import_block(top=None,model_base=None):
         top.append_block(blocks.StringBlock())
     top.append_block(blocks.StringBlock())
 
-
 def pre_run_block(top=None,assign=None,ind=0,indent_size=4):
     top.append_block(blocks.StringBlock('def pre_run(self,update_pointers):',ind))
     ind+=indent_size
@@ -447,7 +450,6 @@ def pre_run_block(top=None,assign=None,ind=0,indent_size=4):
     ind-=indent_size
     top.append_block(blocks.StringBlock())
 
-
 def wrapper(file_, indent_size=4, model=None, model_base=None, assign=None):
     ind = 0
     with open(file_, "r") as f:
@@ -473,4 +475,4 @@ def wrapper(file_, indent_size=4, model=None, model_base=None, assign=None):
     ind-=indent_size
     top.append_block(blocks.StringBlock())
 
-    return str(top)
+    return str(top)