# stardard library
import sys, os

# external libraries
import numpy as np
from ray import imio, evaluate, morpho

def add_anything(a, b):
    return a + b

def write_synapse_to_vtk(neurons, coords, fn, im=None, t=(2,0,1), s=(1,-1,1),
        margin=None):
    """Output neuron shapes around pre- and post-synapse coordinates.
    
    The coordinate array is a (n+1) x m array, where n is the number of 
    post-synaptic sites (fly neurons are polyadic) and m = neurons.ndim, the
    number of dimensions of the image.
    """
    neuron_ids = neurons[zip(*(coords[:,t]*s))]
    synapse_volume = reduce(add_anything, 
        [(i+1)*(neurons==j) for i, j in enumerate(neuron_ids)])
    imio.write_vtk(synapse_volume, fn)
    if im is not None:
        imio.write_vtk(im, 
            os.path.join(os.path.dirname(fn), 'image.' + os.path.basename(fn)))

def get_box(a, coords, margin):
    """Obtain a box of size 2*margin+1 around coords in array a.

    Boxes close to the boundary are trimmed accordingly.
    """
    coords = np.array(coords)[np.newaxis, :]
    origin = np.zeros(coords.shape, dtype=int)
    shape = np.array(a.shape)[np.newaxis, :]
    topleft = np.concatenate((coords-margin, origin), axis=0).max(axis=0)
    bottomright = np.concatenate((coords+margin+1, shape), axis=0).min(axis=0)
    box = [slice(top, bottom) for top, bottom in zip(topleft, bottomright)]
    return a[box].copy()

@@ -26,16 +26,28 @@
  sys, os
+%0Aimport json
 %0A%0A# exte
@@ -530,24 +530,52 @@
 ge.%0A    %22%22%22%0A
+    coords = coords%5B:, t%5D*s%0A
     neuron_i
@@ -596,25 +596,122 @@
 ip(*
-(
 coords
-%5B:,t%5D*s))%5D
+)%5D%0A    mean_coords = coords.mean(axis=0).astype(np.uint)%0A    neurons = get_box(neurons, mean_coords, margin)
 %0A   
@@ -876,16 +876,62 @@
 t None:%0A
+        im = get_box(im, mean_coords, margin)%0A
         
@@ -1186,32 +1186,72 @@
 dingly.%0A    %22%22%22%0A
+    if margin is None:%0A        return a%0A
     coords = np.
@@ -1600,16 +1600,16 @@
 right)%5D%0A
-
     retu
@@ -1625,8 +1625,1064 @@
 .copy()%0A
+%0A%0Adef raveler_synapse_annotations_to_coords(fn):%0A    %22%22%22Obtain pre- and post-synaptic coordinates from Raveler annotations.%22%22%22%0A    with open(fn, 'r') as f:%0A        syn = json.load(f)%0A    tbars = %5Bnp.array(s%5B'location'%5D) for s in syn%5D%0A    posts = %5Bnp.array(%5Bp%5B'location'%5D for p in s%5B'partners'%5D%5D) for s in syn%5D%0A    return %5Bnp.concatenate((t%5Bnp.newaxis, :%5D, p), axis=0)%0A                                        for t, p in zip(tbars, posts)%5D%0A%0Adef write_all_synapses_to_vtk(neurons, list_of_coords, fn, im, t=(2,0,1), %0A        s=(1,-1,1), margin=None, single_pairs=True):%0A    for i, coords in enumerate(list_of_coords):%0A        if single_pairs:%0A            pre = coords%5B0%5D%0A            for j, post in enumerate(coords%5B1:%5D):%0A                pair_coords = np.concatenate(%0A                    (pre%5Bnp.newaxis, :%5D, post%5Bnp.newaxis, :%5D), axis=0)%0A                fn = fn%25(i, j)%0A                write_synapse_to_vtk(neurons, pair_coords, fn, im, t, s, margin)%0A        else:%0A            fn = fn%25i%0A            write_synapse_to_vtk(neurons, coords, fn, im, t, s, margin)%0A

Add functions to read in Raveler json