Collector function to get `NeuronGroup` and return dictionary form of all its entities

dict_neurongroup.py

from brian2 import *
from brian2.utils.stringtools import get_identifiers

def collect_NeuronGroup(group):
    
    """
    Collector function to extract necessary information from NeuronGroup
    and represent them in dictionary format
    """
    
    dict = {}
    
    #get name
    dict['name'] = group.name

    # get size
    dict['N'] = group._N

    #get stateupdation method
    dict['method'] = group.method_choice
    
    #get equations
    dict['equations'] = collect_Equations(group.user_equations)

    #get identifiers from user equations, event codes
    dict['identifiers'] = collect_identifiers(group.user_equations.identifiers)
    
    if bool(group.events):
        dict['identifiers'].update(collect_identifiers(get_identifiers(group.events['spike']).union(get_identifiers(group.event_codes['spike']) - 
        group.user_equations.diff_eq_names)))

    #get spike event (user defined events are not considered)
    if bool(group.events):
        dict['events'] = collect_events(group)

    #get clock
    dict['clock'] = {'dt': str(group.clock.dt), 'name': group.clock.name}

    return dict

def collect_events(group):
    
    """
    Collector function to extract all spike event based information
    """
    
    dict = {}
    # add threshold
    dict['spike'] = {'threshold': group.events['spike']}
    
    #check reset defined
    if bool(group.event_codes):
        dict['spike'].update({'reset': group.event_codes['spike']})

    #check refractory
    if group._refractory:
        dict['spike'].update({'refractory': str(group._refractory)})
    
    return dict


def collect_identifiers(identifiers, level = 1):
    
    """
    Collector function to extract identifiers that come under NeuronGroup
    Each identifier have attributes: value and unit
    """
    
    dict = {}
    
    for identifier in identifiers:
        dict[identifier] = {'value': str(get_local_namespace(level + 1)[identifier])}
        # if value is not Quantity instance, then assume no Unit for that
        if isinstance(get_local_namespace(level + 1)[identifier], Quantity):
            dict[identifier].update({'unit': str(get_unit(get_local_namespace(level + 1)[identifier].dim))})

    return dict

def collect_Equations(equations):
    
    """
    Collector function to get governing Equations of the Neuron model.
    DIFFERENTIAL_EQUATIONS's variable attributes: expr, unit, type, dtype, flags
    SUBEXPRESSIONS's variable attributes: expr, unit, type, dtype, flags
    PARAMETERS's variable attributes: unit, type, dtype, flags
    """
    
    dict = {}
    
    # get DIFFERENTIAL_EQUATIONS
    for diff_names in equations.diff_eq_names:

        dict[diff_names] = {'expr': equations[diff_names].expr.code, 'unit': str(equations[diff_names].unit), 
                            'type': equations[diff_names].type, 'dtype': equations[diff_names].var_type}

        if len(equations[diff_names].flags) != 0:
            dict[diff_names]['flags'] = equations[diff_names].flags

    # get SUBEXPRESSIONS
    for subexp_names in equations.subexpr_names:
        
        dict[subexp_names] = {'expr': equations[subexp_names].expr.code, 'unit': str(equations[subexp_names].unit), 
                            'type': equations[subexp_names].type, 'dtype': equations[subexp_names].var_type}
        
        if len(equations[subexp_names].flags) != 0:
            dict[subexp_names]['flags'] = equations[subexp_names].flags

    # get PARAMETERS
    for param_names in equations.parameter_names:
        
        dict[param_names] = {'unit': str(equations[param_names].unit), 'type': equations[param_names].type, 
                            'dtype': equations[param_names].var_type}
        
        if len(equations[param_names].flags) != 0:
            dict[param_names]['flags'] = equations[param_names].flags
    
    return dict

Sample Testcase 1:

eqs = ''' dv/dt = (v_rest - v) / tau :volt '''
v_rest = -70 * mV
v_th = 0.8 * volt
tau = 10 * ms
grp = NeuronGroup(1, eqs, threshold = 'v > v_th', reset = 'v = v_rest', refractory = 2 * ms, method = 'euler')

output:

{
    "name": "neurongroup",
    "N": 1,
    "method": "euler",
    "equations": {
        "v": {
            "expr": "(v_rest - v) / tau",
            "unit": "V",
            "type": "differential equation",
            "dtype": "float"
        }
    },
    "identifiers": {
        "tau": {
            "value": "10. ms",
            "unit": "s"
        },
        "v_rest": {
            "value": "-70. mV",
            "unit": "V"
        },
        "v_th": {
            "value": "0.8 V",
            "unit": "V"
        }
    },
    "events": {
        "spike": {
            "threshold": "v > v_th",
            "reset": "v = v_rest",
            "refractory": "2. ms"
        }
    },
    "clock": {
        "dt": "100. us",
        "name": "defaultclock"
    }
}

Sample Input 2:

area = 100 * umetre ** 2
g_L = 1e-2 * siemens * cm ** -2 * area
E_L = 1000
Cm = 1 * ufarad * cm ** -2 * area
grp = NeuronGroup(10, '''dv/dt = I_leak / Cm : volt
                    I_leak = g_L*(E_L - v) : amp''')

output:

{
    "name": "neurongroup",
    "N": 10,
    "method": [
        "exact",
        "euler",
        "heun"
    ],
    "equations": {
        "v": {
            "expr": "I_leak / Cm",
            "unit": "V",
            "type": "differential equation",
            "dtype": "float"
        },
        "I_leak": {
            "expr": "g_L*(E_L - v)",
            "unit": "A",
            "type": "subexpression",
            "dtype": "float"
        }
    },
    "identifiers": {
        "g_L": {
            "value": "10. nS",
            "unit": "S"
        },
        "E_L": {
            "value": "1000"
        },
        "Cm": {
            "value": "1. pF",
            "unit": "F"
        }
    },
    "clock": {
        "dt": "100. us",
        "name": "defaultclock"
    }
}

Sample Input 3:

area = 20000 * umetre ** 2
Cm = 1 * ufarad * cm ** -2 * area
gl = 5e-5 * siemens * cm ** -2 * area
El = -65 * mV
EK = -90 * mV
ENa = 50 * mV
g_na = 100 * msiemens * cm ** -2 * area
g_kd = 30 * msiemens * cm ** -2 * area
VT = -63 * mV
I = 0.01*nA
eqs = Equations('''
dv/dt = (gl*(El-v) - g_na*(m*m*m)*h*(v-ENa) - g_kd*(n*n*n*n)*(v-EK) + I)/Cm : volt
dm/dt = 0.32*(mV**-1)*(13.*mV-v+VT)/
    (exp((13.*mV-v+VT)/(4.*mV))-1.)/ms*(1-m)-0.28*(mV**-1)*(v-VT-40.*mV)/
    (exp((v-VT-40.*mV)/(5.*mV))-1.)/ms*m : 1
dn/dt = 0.032*(mV**-1)*(15.*mV-v+VT)/
    (exp((15.*mV-v+VT)/(5.*mV))-1.)/ms*(1.-n)-.5*exp((10.*mV-v+VT)/(40.*mV))/ms*n : 1
dh/dt = 0.128*exp((17.*mV-v+VT)/(18.*mV))/ms*(1.-h)-4./(1+exp((40.*mV-v+VT)/(5.*mV)))/ms*h : 1
''')
grp = NeuronGroup(1, eqs, method='exponential_euler')

output:

{
    "name": "neurongroup",
    "N": 1,
    "method": "exponential_euler",
    "equations": {
        "m": {
            "expr": "0.32*(mV**-1)*(13.*mV-v+VT)/ (exp((13.*mV-v+VT)/(4.*mV))-1.)/ms*(1-m)-0.28*(mV**-1)*(v-VT-40.*mV)/ (exp((v-VT-40.*mV)/(5.*mV))-1.)/ms*m",
            "unit": "rad",
            "type": "differential equation",
            "dtype": "float"
        },
        "h": {
            "expr": "0.128*exp((17.*mV-v+VT)/(18.*mV))/ms*(1.-h)-4./(1+exp((40.*mV-v+VT)/(5.*mV)))/ms*h",
            "unit": "rad",
            "type": "differential equation",
            "dtype": "float"
        },
        "n": {
            "expr": "0.032*(mV**-1)*(15.*mV-v+VT)/ (exp((15.*mV-v+VT)/(5.*mV))-1.)/ms*(1.-n)-.5*exp((10.*mV-v+VT)/(40.*mV))/ms*n",
            "unit": "rad",
            "type": "differential equation",
            "dtype": "float"
        },
        "v": {
            "expr": "(gl*(El-v) - g_na*(m*m*m)*h*(v-ENa) - g_kd*(n*n*n*n)*(v-EK) + I)/Cm",
            "unit": "V",
            "type": "differential equation",
            "dtype": "float"
        }
    },
    "identifiers": {
        "exp": {
            "value": "<function wrap_function_dimensionless.<locals>.f at 0x7f7b6510d170>"
        },
        "gl": {
            "value": "10. nS",
            "unit": "S"
        },
        "mV": {
            "value": "mV",
            "unit": "V"
        },
        "VT": {
            "value": "-63. mV",
            "unit": "V"
        },
        "I": {
            "value": "10. pA",
            "unit": "A"
        },
        "g_na": {
            "value": "20. uS",
            "unit": "S"
        },
        "g_kd": {
            "value": "6. uS",
            "unit": "S"
        },
        "Cm": {
            "value": "200. pF",
            "unit": "F"
        },
        "ms": {
            "value": "ms",
            "unit": "s"
        },
        "EK": {
            "value": "-90. mV",
            "unit": "V"
        },
        "El": {
            "value": "-65. mV",
            "unit": "V"
        },
        "ENa": {
            "value": "50. mV",
            "unit": "V"
        }
    },
    "clock": {
        "dt": "100. us",
        "name": "defaultclock"
    }
}

note: exp is identified incorrectly as an identifier, so has to improve the way of collecting identifiers