Documentation on providing time-dependent input to brainpy.dyn.DSRunner

[CloudyDory]

Hi, I hope to give a time-dependent input to a Hodgkin-Huxley model, but the tutorial does not provide a detailed guide on how to implement this. I found that brainpy.DSRunner can receive a function to argument inputs, so I try the following code:

import jax
import brainpy as bp
import brainpy.math as bm
from brainpy.channels import INa_HH1952, IL
import matplotlib.pyplot as plt

print(bp.__version__)
bm.set_platform('cpu')

class IK(bp.dyn.IonChannel):  # bp.Channel
    def __init__(self, size, E=-77., g_max=36., phi=1., method='rk4'):  # 'exp_auto'
        super(IK, self).__init__(size)
        self.g_max = g_max
        self.E = E
        self.phi = phi
      
        self.n = bm.Variable(bm.zeros(size))  # variables should be packed with bm.Variable
        
        self.integral = bp.odeint(self.dn, method=method)
    
    def dn(self, n, t, V):
        alpha_n = 0.01 * (V + 55) / (1 - bm.exp(-(V + 55) / 10))
        beta_n = 0.125 * bm.exp(-(V + 65) / 80)
        return self.phi * (alpha_n * (1. - n) - beta_n * n)
    
    def update(self, tdi, V):
        self.n.value = self.integral(self.n, tdi.t, V, dt=tdi.dt)
    
    def current(self, V):
        return self.g_max * self.n ** 4 * (self.E - V)

class HH(bp.dyn.CondNeuGroup):  # bp.CondNeuGroup
    def __init__(self, size):
        super().__init__(size, V_initializer=bp.init.Uniform(-70., -70.))
        self.IK = IK(size, E=-77., g_max=36.)
        self.INa = INa_HH1952(size, E=50., g_max=120.)
        self.IL = IL(size, E=-54.39, g_max=0.03)

def I_inject(shared):
    return jax.numpy.logical_and(0<=shared['t'], shared['t']<=300) * 6.0

neu = HH(size=1)

runner = bp.DSRunner(
    neu, 
    monitors = ['V'], 
    inputs = I_inject, 
    jit = False
)

runner.run(200)  # the running time is 200 ms

plt.figure()
plt.plot(runner.mon['ts'], runner.mon['V'])
plt.xlabel('t (ms)')
plt.ylabel('V (mV)')
plt.show()

However, the result is the same as setting the input current to zero. The function I_inject outputs the correct value, but it seems that the value is never used.

The documention on brainpy.DSRunner (https://brainpy.readthedocs.io/en/latest/apis/auto/generated/brainpy.DSRunner.html) also provides another piece of information:

This argument can be used to set the inputs to the Variable instances in the target. If you peruse to give time-dependent inputs, please use DSRunner.predict() or DSRunner.run() function.

But the documentation of DSRunner.predict is not very clear. It seems that we can only pre-generate the input data and use it when running the model. But what if the input can only be generated during runtime?

In summary, I have three questions:

1. Why doesn't providing a function to the inputs argument in brainpy.DSRunner work?
2. Should I provide the functions in Inputs Construction section of the documentation to brainpy.DSRunner?
3. How to provide input that can only be generated during runtime to brainpy.DSRunner?

Thanks!

[CloudyDory]

The above example works if we: 1. provide a pre-generated input current to DSRunner.run() function, and 2. Override the update method to use the given input current in the HH neuron.

However, providing a function to the inputs argument in brainpy.DSRunner still does not work.

[chaoming0625]

Dear @CloudyDory , thank you for your report. Actually, BrainPy has a very easy way to given any inputs.

For your model, you can give inputs using brainpy.DSRunner:

inputs = bp.inputs.section_input([0., 6.0, 0.], [100. , 200., 100.])

runner = bp.DSRunner(..., )
runner.run(inputs=inputs)

Or, you can run the mode using brainpy.math.for_loop:

inputs = bp.inputs.section_input([0., 6.0, 0.], [100. , 200., 100.])
indices = np.arange(inputs.size)


def run(i, x):
  neu.step_run(i, x)
  return neu.V.value


vs = bm.for_loop(run, (indices, inputs), progress_bar=True)

You can even run the model with your native Python syntax:

neu = HH(size=1)

vs = []
indices = np.arange(2000)  # 200 ms
for i in range(indices):
  neu.jit_step_run(i, 6. if 1000 < i < 3000 else 0.)
  vs.append(neu.V.value)
vs = bm.as_jax(vs)