How to set the number of compartments (d_lambda rule)

Jaxley allows you to change the number of compartments that each branch is made up of by using the set_ncomp() method:

cell.branch(0).set_ncomp(4)

To automatically choose the number of compartments, you can use the d_lambda rule, which can be implemented in Jaxley as follows:

import jaxley as jx

cell = jx.read_swc("path_to_swc.swc", ncomp=1)

# Reasonable default values for most models.
frequency = 100.0
d_lambda = 0.1  # Larger -> more coarse-grained.

for branch in cell.branches:
    diameter = 2 * branch.nodes["radius"].to_numpy()[0]
    c_m = branch.nodes["capacitance"].to_numpy()[0]
    r_a = branch.nodes["axial_resistivity"].to_numpy()[0]
    l = branch.nodes["length"].to_numpy()[0]

    lambda_f = 1e5 * np.sqrt(diameter / (4 * np.pi * frequency * c_m * r_a))
    ncomp = int((l / (d_lambda * lambda_f) + 0.9) / 2) * 2 + 1
    branch.set_ncomp(ncomp, initialize=False)
  
# After the loop, you have to run `cell.initialize()` because we passed
# `set_ncomp(..., initialize=False)` for speeding up the loop over branches.
cell.initialize()