Merge pull request #26 from the-virtual-brain/codejam9

Code generation without loopy

.gitignore

@@ -1,17 +1,2 @@
-*.swp
-build
-*.pyc
+env
 __pycache__
-*.swo
-*.o
-test_main
-*.S
-CMakeCache.txt
-CMakeFiles
-cmake_install.cmake
-compile_commands.json
-tests
-data
-.ipynb_checkpoints
-.idea
-tvb_hpc.egg-info

.gitlab-ci.yml

@@ -12,11 +12,17 @@ info:
     - whoami
     - which python
     - which pip
+    - docker info
+
+build_image:
+  stage: build
+  script:
+    - docker build -t tvb/hpc .
+  allow_failure: true
 
 build_env:
   stage: build
   script:
-    - conda install numpy scipy llvmlite numba scikit-learn
     - pip install -r requirements.txt
     - python setup.py install
 
@@ -39,7 +45,8 @@ unittest:
 example_hackathon:
   stage: test
   script:
-    - python examples/hackathon.py
+    - python -m tvb_hpc.examples.hackathon
+  allow_failure: true
 
 wheel:
   stage: package
@@ -56,6 +63,6 @@ upload_pypi:
 push_github:
   stage: package
   script:
-    - git push [email protected]:the-virtual-brain/tvb-hpc.git $(git rev-parse --abbrev-ref HEAD)
+    - git push [email protected]:the-virtual-brain/tvb-hpc.git
 
 # vim: sw=2 sts=2 et ai

.pre-commit

@@ -1,13 +1,3 @@
 #!/bin/bash
 
-if [[ -z "$(which flake8)" ]]
-then
-    echo "pre-commit: Can't find flake8; please pip -r requirements-dev.txt"
-    exit 1
-fi
-
-set -eu
-
-export PYTHONPATH=$(git rev-parse --show-toplevel)
-python3 -m unittest tvb_hpc.tests
 flake8 tvb_hpc

README.md

@@ -3,26 +3,23 @@
 This is a Python package for generating code for parameter sweeps and Bayesian
 inversion.
 
-Get the Docker image
-```
-docker pull maedoc/tvb-hpc
-```
+## Quickstart
 
-and run the tests
 ```
-docker run --rm -it -v ./:/root/hpc python -m unittest tvb_hpc.tests
+git clone https://gitlab.thevirtualbrain.org/tvb/hpc tvb-hpc
+cd tvb-hpc
+python -m venv env
+. env/bin/activate
+pip install -r requirements.txt
+python -m unittest tvb_hpc.tests
 ```
-(on Windows, replace `./` by `%CD%`.)
 
 Get hacking with the [`examples`](examples) or [tests](tvb_hpc/tests.py).
 
-## TODO
+## Targets
+
+By default the requirements.txt file will have Numba installed, which
+is an easy target to use, but other targets are straightforward to use
 
-- ensure par sweeps etc built as domains
-- rng
-- make high level usage easier (tavg, bold, gain, fcd, etc)
-- test on CUDA
-- parallel numba
-- simple SALib usage?
-- chunking of state & vectorization?
-- reach cuda hackathon performance numbers
+- `C` - ensure GCC or Clang are installed
+- `OpenCL` - `pip install pyopencl`

examples/progressive_knl.py

@@ -0,0 +1,39 @@
+import numpy as np
+import loopy as lp
+from tvb_hpc.numbacudatarget import NumbaCudaTarget
+target = NumbaCudaTarget()
+
+coupling_raw = """
+<> theta_i = state[i_nodie] {nosync=*}	
+<> sum = 0.0	
+<> wij = weights[j_node]
+if wij != 0.0
+    <int> dij = lengths[j_node] * rec_speed_dt		
+    <float32> theta_j = state[j_node] {nosync=*}
+    sum = sum + wij * sin(theta_j - theta_i)
+end
+"""
+model_raw = """
+theta_i = theta_i + dt * (omega + coupling_value * rec_n * sum) 	
+theta_i = theta_i + (sig * rand) 					
+theta_i = wrap_2_pi(theta_i) 					
+tavg[i_node] = tavg[i_node] + sin(theta_i) 	
+state[i_node] = theta_i {nosync=*}	
+"""
+node = 10
+#Raw coupling
+couplingknl = lp.make_kernel("{ [i_node, j_node]: 0<=i_node, j_node<n_node}",coupling_raw,target = target)
+couplingknl = lp.add_dtypes(couplingknl, {"lengths": np.float32, "state": np.float32, "weights": np.float32, "theta_i": np.float32, "rec_speed_dt": np.float32})
+couplingknl = lp.split_iname(couplingknl, "j_node", 1, outer_tag='l.0')
+#Raw model
+modelknl = lp.make_kernel("{ [i_node]: 0<=i_node<n_node}",model_raw,target = target)
+modelknl = lp.add_dtypes(modelknl, {"state": np.float32, "theta_i": np.float32, "tavg": np.float32})
+# Fuse
+knls = couplingknl, modelknl
+data_flow = [('state', 0, 1)]
+knl = lp.fuse_kernels(knls, data_flow = data_flow)
+print (knl)
+print ("****")
+knl = lp.split_iname(knl, "i_node", 128, outer_tag='g.0')
+print(lp.generate_code_v2(knl).all_code())
+

phase_plane_interactive/README.md

@@ -0,0 +1,138 @@
+# CodeJam 9 code generators
+
+Here we explore the code synthesis without the loopy framework. The basic idea is to stick with the simple structured model representation and use a combination of AST transformations and active code templates to target various scenarios and architectures. Ideally, the templates will be derived from existing hand-written implementations.
+
+
+
+Relevant files:
+* `models.py` : example of model structured definition
+
+* `dsl2python.py` : simple templated synthesis of a TVB-like Python class with a phase-plane interactive plot
+
+  *  `template.py` active template for Python model class
+  *  `ppi.py` : phase-plane interactive plot
+
+* `dsl2C.py` : example of AST transformations for C-like languages
+
+  *  `template.cu` active template for Cuda dfun function
+
+
+
+
+## Model description
+
+This was adapted from the loopy implementation. Note, that this is more a data-structure (string drifts), than a class with executable functions. Could be as well a JSON object... 
+
+While it would be better to have  a functional model class, it is difficult to create one without introducing redundancy in the notation -- especially in the drift function definitions. So let's just stick with this: 
+
+```Python
+class G2DO:
+    "Generic nonlinear 2-D (phase plane) oscillator."
+    state = 'W', 'V'
+    limit = (-5, 5), (-5, 5)
+    input = 'c_0'
+    param = 'a'
+    const = {'tau': 1.0, 'I': 0.0, 'a': -2.0, 'b': -10.0, 'c': 0.0, 'd': 0.02,
+             'e': 3.0, 'f': 1.0, 'g': 0.0, 'alpha': 1.0, 'beta': 1.0,
+             'gamma': 1.0}
+    drift = (
+        'd * tau * (alpha*W - f*V**3 + e*V**2 + g*V + gamma*I + gamma*c_0)',
+        'd * (a + b*V + c*V**2 - beta*W) / tau'
+    )
+    diffs = 1e-3, 1e-3
+    obsrv = 'W', 'V'
+
+```
+
+
+
+## Code synthesis
+
+For both C and Python we use the `mako` template engine to fill templates preparing the boilerplate for the `drift` expressions.
+
+The advantage of the `drift` having Python syntax is, that we can use the built-in `ast` library for parsing, and the `ctree` library for basic transformations and code synthesis. For the demo, the simple operator conversion had to be extended to cope with `**`. 
+
+In future, we can use the `ctree` framework for more complex transformations, such as array index arithmetic etc.
+
+
+
+A demo Python code synthesized for the phase-plane interactive plot:
+
+```Python
+class Generic2D:
+
+    def __init__(self,tau=1.0, I=0.0, a=-2.0, b=-10.0, c=0.0, d=0.02, e=3.0, f=1.0, g=0.0, alpha=1.0, beta=1.0, gamma=1.0):
+        self.limit = ((-5, 5), (-5, 5))
+        self.tau = tau
+        self.I = I
+        self.a = a
+        self.b = b
+        self.c = c
+        self.d = d
+        self.e = e
+        self.f = f
+        self.g = g
+        self.alpha = alpha
+        self.beta = beta
+        self.gamma = gamma
+
+    def dfun(self,state_variables, *args, **kwargs):
+
+        W, V = state_variables
+
+        tau = self.tau
+        I = self.I
+        a = self.a
+        b = self.b
+        c = self.c
+        d = self.d
+        e = self.e
+        f = self.f
+        g = self.g
+        alpha = self.alpha
+        beta = self.beta
+        gamma = self.gamma
+
+        c_0 = 0.0
+
+        dW = d * tau * (alpha*W - f*V**3 + e*V**2 + g*V + gamma*I + gamma*c_0)
+        dV = d * (a + b*V + c*V**2 - beta*W) / tau
+        return [W, V] 
+```
+
+
+
+A demo code synthesized to fit into an old code generator for TVB. 
+
+```C
+void model_dfun(
+        float * _dx, float *_x, float *mmpr, float*input
+    )   
+{       
+    float tau = mmpr[n_thr*0];
+    float I = mmpr[n_thr*1];
+    float a = mmpr[n_thr*2];
+    float b = mmpr[n_thr*3];
+    float c = mmpr[n_thr*4];
+    float d = mmpr[n_thr*5];
+    float e = mmpr[n_thr*6];
+    float f = mmpr[n_thr*7];
+    float g = mmpr[n_thr*8];
+    float alpha = mmpr[n_thr*9];
+    float beta = mmpr[n_thr*10];
+    float gamma = mmpr[n_thr*11];
+
+    float W = _x[n_thr*0];
+    float V = _x[n_thr*1];
+
+    float c_0 = input[0];
+
+    _dx[n_thr*0] = d * tau * (alpha * W - f * pow(V, 3) + e * pow(V, 2) + g * V + gamma * I + gamma * c_0);
+    _dx[n_thr*1] = d * tau * (alpha * W - f * pow(V, 3) + e * pow(V, 2) + g * V + gamma * I + gamma * c_0);
+
+}
+
+```
+
+
+

phase_plane_interactive/ast_tranfrom.py

@@ -0,0 +1,68 @@
+import ast 
+import astor
+
+
+class IndexArithmetic(ast.NodeTransformer):
+    def __init__(self, index_map):
+        self.index_map = index_map
+
+    def visit_Name(self, node):
+        if node.id not in self.index_map:
+            return node
+        index = self.index_map[node.id]
+
+        new_node = ast.Subscript(
+            value=ast.Name(id='X', ctx=ast.Load()),
+            slice=ast.Index(value=index),
+            ctx=node.ctx
+        )
+        return new_node
+
+def index_dfun(model):
+
+    n_svar = len(model.state)
+    svar_i = dict(zip(model.state,range(n_svar)))
+
+    dfun_asts = []
+    for drift in model.drift:
+        dfun_asts.append(ast.parse(drift).body[0].value)
+
+
+    svar_index = dict.fromkeys(model.state)
+
+    for i, var in enumerate(model.state):
+        index = ast.Tuple(
+                ctx=ast.Load(),
+                elts=[  ast.Num(n=i),
+                        ast.Name(id='t', ctx=ast.Load())
+                    ])
+        svar_index[var] = index
+
+    dfun_idx_asts = []
+    for i, var in enumerate(model.state):
+
+        index = svar_index[var]
+        dX = ast.Subscript(
+                value=ast.Name(id='dX', ctx=ast.Load()),
+                slice=ast.Index(value=index),
+                ctx=ast.Store() )
+        dfun_idx_ast = ast.Assign(
+                targets=[dX],
+                value=IndexArithmetic(svar_index).generic_visit(dfun_asts[i]))
+        dfun_idx_asts.append( dfun_idx_ast )
+
+    nowrap = lambda x:''.join(x)
+    dfun_strings = list(map(lambda x: astor.to_source(x, pretty_source=nowrap), 
+                            dfun_idx_asts ) )
+
+    return dfun_strings
+
+
+if __name__=="__main__":
+    from ppi import G2DO
+    from mako.template import Template
+
+    dfuns = index_dfun(G2DO)
+
+    template = Template(filename='numba_model.py')
+    print(template.render(dfuns=dfuns, const=G2DO.const, cvar=G2DO.input))