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)

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity in the last 90 days. It will be closed in 14 days if no further activity occurs. If this issue is still relevant, please comment to keep it open. Thank you for your contributions!