<commit_before># stardard library
import sys, os
import json

# 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, 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)]
    mean_coords = coords.mean(axis=0).astype(np.uint)
    neurons = get_box(neurons, mean_coords, margin)
    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:
        im = get_box(im, mean_coords, margin)
        imio.write_vtk(im, 
            os.path.join(os.path.dirname(fn), 'image.' + os.path.basename(fn)))

def all_postsynaptic_sites(synapses):
    tbars, posts = zip(*synapses)
    return list(it.chain(*posts))

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.
    """
    if margin is None:
        return a
    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()

def synapses_from_raveler_session_data(fn):
    with open(fn) as f:
        d = pck.load(f)
    annots = d['annotations']['point']
    tbars = [a for a in annots if annots[a]['kind'] == 'T-bar']
    posts = [annots[a] for a in tbars]
    posts = [eval(p['value'].replace('false', 'False').replace('true', 'True'))
             for p in posts]
    posts = [p['partners'] for p in posts]
    posts = [map(lambda x: x[0], p) for p in posts]
    return zip(tbars, posts)


def raveler_synapse_annotations_to_coords(fn, output_format='pairs',
                                    t=(2,0,1), s=(1,-1,1), transform=True):
    """Obtain pre- and post-synaptic coordinates from Raveler annotations."""
    if not transform:
        t = (0, 1, 2)
        s = (1, 1, 1)
    with open(fn, 'r') as f:
        syns = json.load(f)['data']
    tbars = [np.array(syn['T-bar']['location'])[:, t]*s for syn in syns]
    posts = [np.array([p['location'] for p in syn['partners']])[:, t]*s 
        for syn in syns]
    if output_format == 'pairs':
        return zip(tbars, posts)
    elif output_format == 'arrays':
        return [np.concatenate((t[np.newaxis, :], p), axis=0)
                                        for t, p in zip(tbars, posts)]

def raveler_to_numpy_coord_transform(coords, t=(2, 0, 1), s=(1, -1, 1)):
    coords = np.array(coords)
    return coords[:, t]*s

def write_all_synapses_to_vtk(neurons, list_of_coords, fn, im, margin=None,
                                                            single_pairs=True):
    for i, coords in enumerate(list_of_coords):
        if single_pairs:
            pre = coords[0]
            for j, post in enumerate(coords[1]):
                pair_coords = np.concatenate(
                    (pre[np.newaxis, :], post[np.newaxis, :]), axis=0)
                fn = fn%(i, j)
                write_synapse_to_vtk(neurons, pair_coords, fn, im, margin)
        else:
            fn = fn%i
            write_synapse_to_vtk(neurons, coords, fn, im, margin)
<commit_msg>Transform Raveler coords to Numpy by default<commit_after># stardard library
import sys, os
import json
import cPickle as pck

# 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, 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)]
    mean_coords = coords.mean(axis=0).astype(np.uint)
    neurons = get_box(neurons, mean_coords, margin)
    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:
        im = get_box(im, mean_coords, margin)
        imio.write_vtk(im, 
            os.path.join(os.path.dirname(fn), 'image.' + os.path.basename(fn)))

def all_postsynaptic_sites(synapses):
    tbars, posts = zip(*synapses)
    return list(it.chain(*posts))

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.
    """
    if margin is None:
        return a
    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()

def tbar_post_pairs_to_arrays(pairs):
    return [np.concatenate((t[np.newaxis, :], p), axis=0) for t, p in pairs]

def synapses_from_raveler_session_data(fn, output_format='pairs', 
                                    t=(2, 1, 0), s=(1, -1, 1), transform=True):
    if not transform:
        t = (0, 1, 2)
        s = (1, 1, 1)
    with open(fn) as f:
        d = pck.load(f)
    annots = d['annotations']['point']
    tbars = [a for a in annots if annots[a]['kind'] == 'T-bar']
    posts = [annots[a] for a in tbars]
    posts = [eval(p['value'].replace('false', 'False').replace('true', 'True'))
             for p in posts]
    posts = [p['partners'] for p in posts]
    posts = [map(lambda x: x[0], p) for p in posts]
    tbars = [np.array(tbar)[np.array(t)]*s for tbar in tbars]
    posts = [np.array(post)[:, t] * s for post in posts]
    pairs = zip(tbars, posts)
    if output_format == 'pairs':
        return pairs
    elif output_format == 'arrays':
        return tbar_post_pairs_to_arrays(pairs)

def raveler_synapse_annotations_to_coords(fn, output_format='pairs',
                                    t=(2, 1, 0), s=(1, -1, 1), transform=True):
    """Obtain pre- and post-synaptic coordinates from Raveler annotations."""
    if not transform:
        t = (0, 1, 2)
        s = (1, 1, 1)
    with open(fn, 'r') as f:
        syns = json.load(f)['data']
    tbars = [np.array(syn['T-bar']['location'])[:, t]*s for syn in syns]
    posts = [np.array([p['location'] for p in syn['partners']])[:, t]*s 
        for syn in syns]
    pairs = zip(tbars, posts)
    if output_format == 'pairs':
        return pairs
    elif output_format == 'arrays':
        return tbar_post_pairs_to_arrays(pairs)

def raveler_to_numpy_coord_transform(coords, t=(2, 1, 0), s=(1, -1, 1)):
    coords = np.array(coords)
    return coords[:, t]*s

def write_all_synapses_to_vtk(neurons, list_of_coords, fn, im, margin=None,
                                                            single_pairs=True):
    for i, coords in enumerate(list_of_coords):
        if single_pairs:
            pre = coords[0]
            for j, post in enumerate(coords[1]):
                pair_coords = np.concatenate(
                    (pre[np.newaxis, :], post[np.newaxis, :]), axis=0)
                fn = fn%(i, j)
                write_synapse_to_vtk(neurons, pair_coords, fn, im, margin)
        else:
            fn = fn%i
            write_synapse_to_vtk(neurons, coords, fn, im, margin)

{'fraction_non_alphanumeric': 0.13707440100882723, 'fraction_numerical': 0.009836065573770493, 'mean_word_length': 2.6905537459283386, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 9, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2300630', 'n_tokens_mistral': 2851, 'n_tokens_neox': 2749, 'n_words': 740}