File: spatialstateupdate.py_

package info (click to toggle)
brian 2.9.0-2
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 6,872 kB
  • sloc: python: 51,820; cpp: 2,033; makefile: 108; sh: 72
file content (164 lines) | stat: -rw-r--r-- 6,582 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
{# USES_VARIABLES { Cm, dt, v, N, Ic, Ri,
                  _ab_star0, _ab_star1, _ab_star2, _b_plus, _b_minus,
                  _v_star, _u_plus, _u_minus,
                  _v_previous,
                  _gtot_all, _I0_all, v
                  _c,
                  _P_diag, _P_parent, _P_children,
                  _B, _morph_parent_i, _starts, _ends,
                  _morph_children, _morph_children_num, _morph_idxchild,
                  _invr0, _invrn, _invr,
                  r_length_1, r_length_2, area} #}
{# ITERATE_ALL { _idx } #}
"""
Solves the cable equation (spatial diffusion of currents).
This is where most time-consuming time computations are done.
"""
{% extends 'common_group.py_' %}

{# Preparation of the data structures #}
{% block before_code %}
# Inverse axial resistance
{{_invr}}[1:] = 1.0/({{Ri}}*(1/{{r_length_2}}[:-1] + 1/{{r_length_1}}[1:]));
# Cut sections
for _first in {{_starts}}:
    {{_invr}}[_first] = 0

# Linear systems
# The particular solution
"""a[i,j]=ab[u+i-j,j]"""  # u is the number of upper diagonals = 1

{{_ab_star0}}[1:] = {{_invr}}[1:] / {{area}}[:-1]
{{_ab_star2}}[:-1] = {{_invr}}[1:] / {{area}}[1:]
{{_ab_star1}}[:] = (-({{Cm}} / {{dt}}) - {{_invr}} / {{area}})
{{_ab_star1}}[:-1] -= {{_invr}}[1:] / {{area}}[:-1]

# Set the boundary conditions
for _counter, (_first, _last) in enumerate(zip({{_starts}},
                                               {{_ends}})):
    _last = _last -1  # the compartment indices are in the interval [starts, ends[
    # Inverse axial resistances at the ends: r0 and rn
    {{_invr0}}[_counter] = _invr0 = {{r_length_1}}[_first]/{{Ri}}
    {{_invrn}}[_counter] = _invrn = {{r_length_2}}[_last]/{{Ri}}
    # Correction for boundary conditions
    {{_ab_star1}}[_first] -= (_invr0 / {{area}}[_first])
    {{_ab_star1}}[_last] -= (_invrn / {{area}}[_last])
    # RHS for homogeneous solutions
    {{_b_plus}}[_last] = -(_invrn / {{area}}[_last])
    {{_b_minus}}[_first] = -(_invr0 / {{area}}[_first])
{% endblock %}

{% block maincode %}
from numpy import pi
from numpy import zeros
try:
    from scipy.linalg import solve_banded
except ImportError:
    raise ImportError("Install 'scipy' to run multi-compartmental neurons with numpy")

# scalar code
_vectorisation_idx = 1
{{scalar_code|autoindent}}

# vector code
_vectorisation_idx = LazyArange(N)
{{vector_code|autoindent}}

{{_v_previous}}[:] = {{v}}

# Particular solution
_b=-({{Cm}}/{{dt}}*{{v}})-_I0
_ab = zeros((3,N))
_ab[0,:] = {{_ab_star0}}
_ab[1,:] = {{_ab_star1}} - _gtot
_ab[2,:] = {{_ab_star2}}
{{_v_star}}[:] = solve_banded((1,1),_ab,_b,overwrite_ab=True,overwrite_b=True)
# Homogeneous solutions
_b[:] = {{_b_plus}}
_ab[0,:] = {{_ab_star0}}
_ab[1,:] = {{_ab_star1}} - _gtot
_ab[2,:] = {{_ab_star2}}
{{_u_plus}}[:] = solve_banded((1,1),_ab,_b,overwrite_ab=True,overwrite_b=True)
_b[:] = {{_b_minus}}
_ab[0,:] = {{_ab_star0}}
_ab[1,:] = {{_ab_star1}} - _gtot
_ab[2,:] = {{_ab_star2}}
{{_u_minus}}[:] = solve_banded((1,1),_ab,_b,overwrite_ab=True,overwrite_b=True)

# indexing for _P_children which contains the elements above the diagonal of the coupling matrix _P
children_rowlength = len({{_morph_children}})//len({{_morph_children_num}})

# Construct the coupling system with matrix _P in sparse form. s.t.
# _P_diag contains the diagonal elements
# _P_children contains the super diagonal entries
# _P_parent contains the single sub diagonal entry for each row
# _B contains the right hand side

_P_children_2d = {{_P_children}}.reshape(-1, children_rowlength)

for _i, (_i_parent, _i_childind, _first, _last, _invr0, _invrn) in enumerate(zip({{_morph_parent_i}},
                                                                                 {{_morph_idxchild}},
                                                                                 {{_starts}},
                                                                                 {{_ends}},
                                                                                 {{_invr0}},
                                                                                 {{_invrn}})):
    _last = _last - 1 # the compartment indices are in the interval [starts, ends[
    # Towards parent
    if _i == 0: # first section, sealed end
        {{_P_diag}}[0] = {{_u_minus}}[_first] - 1
        _P_children_2d[0, 0] = {{_u_plus}}[_first]

        # RHS
        {{_B}}[0] = -{{_v_star}}[_first]
    else:
        {{_P_diag}}[_i_parent] += (1 - {{_u_minus}}[_first]) * _invr0
        _P_children_2d[_i_parent, _i_childind] = -{{_u_plus}}[_first] * _invr0

        # RHS
        {{_B}}[_i_parent] += {{_v_star}}[_first] * _invr0

    # Towards children
    {{_P_diag}}[_i+1] = (1 - {{_u_plus}}[_last]) * _invrn
    {{_P_parent}}[_i] = -{{_u_minus}}[_last] * _invrn

    # RHS
    {{_B}}[_i+1] = {{_v_star}}[_last] * _invrn

# Solve the linear system (the result will be stored in the former rhs _B in the end)
# use efficient O(n) solution of the sparse linear system (structure-specific Gaussian elemination)
_morph_children_2d = {{_morph_children}}.reshape(-1, children_rowlength)

# part 1: lower triangularization

for _i in range(len({{_B}})-1, -1, -1):
    _num_children = {{_morph_children_num}}[_i];

    for _k in range(_num_children):
        _j = _morph_children_2d[_i, _k]  # child index

        # subtracting _subfac times the j-th from the _i-th row
        _subfac = _P_children_2d[_i, _k] / {{_P_diag}}[_j]

        {{_P_diag}}[_i] = {{_P_diag}}[_i]  - _subfac * {{_P_parent}}[_j-1]
        {{_B}}[_i] = {{_B}}[_i] - _subfac * {{_B}}[_j]

# part 2: forwards substitution
{{_B}}[0] = {{_B}}[0] / {{_P_diag}}[0]  # the first section does not have a parent
for _i, j in enumerate({{_morph_parent_i}}):
    {{_B}}[_i+1] -= {{_P_parent}}[_i] * {{_B}}[j]
    {{_B}}[_i+1] /= {{_P_diag}}[_i+1]

# For each section compute the final solution by linear combination of the general solution
for _i, (_B_parent, _j_start, _j_end) in enumerate(zip({{_B}}[{{_morph_parent_i}}],
                                                       {{_starts}},
                                                       {{_ends}})):
    _B_current = {{_B}}[_i+1]
    if _j_start == _j_end:
        {{v}}[_j_start] = ({{_v_star}}[_j_start] + _B_parent * {{_u_minus}}[_j_start]
                     + _B_current * {{_u_plus}}[_j_start])
    else:
        {{v}}[_j_start:_j_end] = ({{_v_star}}[_j_start:_j_end] + _B_parent * {{_u_minus}}[_j_start:_j_end]
                     + _B_current * {{_u_plus}}[_j_start:_j_end])

{{Ic}}[:] = {{Cm}}*({{v}} - {{_v_previous}})/{{dt}}
{% endblock %}