NMODL matexp solver

This notebook describes the implementation of the matexp solver, which solves the systems of ODEs defined in KINETIC blocks when the ODEs are linear and coupled.

For a more general tutorial on using the NMODL python interface, please see the tutorial notebook.


Mathematics

ODEs are linear if they can be written in the form \(\frac{dx}{dt} = A \cdot x\) where \(x\) is the state vector of the model and \(A\) is the Jacobian matrix. The analytic solution to this initial value problem is \(x = e^{A \Delta t} \cdot x_0\). The matrix \(e^{A \Delta t}\) is known as the propagator matrix and it advances the state of the model by one time step. Since this is the exact solution, it is compatible with the Crank−Nicholson method which requires at least second-order accuracy.


Matrix Exponential Function

The exponential function is defined over matrices as well as scalars by its Taylor series expansion \(e^x = \sum_{n=0}^{\infty} \frac{x^n}{n!}\). However, computing the exponential function by evaluating its Taylor series is slow and numerically unstable. There are many faster and more accurate methods of calculating the matrix exponential function. The NEURON simulator uses the Eigen library for its advanced linear algebra functions, including for the matrix exponential function. For this function, the Eigen library implements the algorithm introduced in the publication:

“The scaling and squaring method for the matrix exponential revisited,”
By Nicholas J. Higham,
SIAM J. Matrix Anal. Applic., 26:1179–1193, 2005.

Implementation

The MatexpVisitor implements the solver method matexp.

The visitor does the following:

  • make lists of all SOLVE statements, STEADYSTATE statements, and KINETIC blocks in the program

  • for each SOLVE block STEADYSTATE matexp statement

    • find the corresponding KINETIC block and convert it into a MATEXP_BLOCK

    • replace the SOLVE ... STEADYSTATE statement with the MATEXP_BLOCK

  • for each SOLVE block METHOD matexp statement

    • find the corresponding KINETIC block and convert it into a MATEXP_BLOCK

    • append the MATEXP_BLOCK to the program

  • remove the solved KINETIC blocks

MATEXP_BLOCKs have the following fields:

  • steadystate (boolean flag)

  • kinetic block

  • list of conserve statements, which the visitor removed from the kinetic block

During the code generation phase, all MATEXP_BLOCKs that were appended to the program are moved into the state update function. Then MATEXP_BLOCKs generate code in-place either in the initialization function or in the state update function. MATEXP_BLOCKs generate the following code:

  • define and zero-init the Jacobian matrix using the Eigen library

  • insert the body of the kinetic block

    • replace reaction statements with equivalent assignments to the Jacobian matrix

      • For example ~ X <-> Y (A, B) would become:

        • J[X, X] -= A

        • J[Y, X] += A

        • J[Y, Y] -= B

        • J[X, Y] += B

  • initialize the state vector

  • define the solution vector as solution = jacobian.exp() * state

  • conserve the sum of the states, if applicable

    • multiply each state variable by the expected sum divided by the actual sum

  • assign the solution vector to the state variables


Tests Cases

The unit tests may be helpful to understand what these functions are doing

The usecase tests compare several forumations of the same Hodgkin-Huxley model

  • Located in test/nmodl/transpiler/usecases/matexp

  • hh_cnexp.mod is the standard HH model

    • solved by cnexp

  • hh_matexp.mod is identical to hh_cnexp.mod except each gate (“M”, “H”, and “N”) is a two state Markov model

    • solved by matexp

  • hhkin.mod is modified to enumerate each all 12 states of each ion channel

    • solved by sparse

  • kkkinmatexp.mod is identical to hhkin.mod except it is solved by matexp


Examples

[1]:
%%capture
! pip install neuron
[11]:
import neuron.nmodl.dsl as nmodl


def run_matexp_solver(mod_string):
    # Parse NMDOL file (supplied as a string) into AST
    driver = nmodl.NmodlDriver()
    AST = driver.parse_string(mod_string)

    # Run SymtabVisitor to generate Symbol Table
    nmodl.symtab.SymtabVisitor().visit_program(AST)

    # Run matexp solver
    nmodl.visitor.MatexpVisitor().visit_program(AST)

    # Return the solution as a string
    return nmodl.to_nmodl(AST)

Ex. 1

Linear kinetic model with conserved sum of states

[3]:
ex1 = """
STATE {
    A
    B
}
BREAKPOINT {
    SOLVE states METHOD matexp
}
KINETIC states {
    ~ A <-> B (0.123, 0.456)
    CONSERVE A + B = 0.789
}
"""
print(run_matexp_solver(ex1))
STATE {
    A
    B
}

BREAKPOINT {
}

MATEXP_SOLVE (0) {
    nmodl_eigen_j[0] = nmodl_eigen_j[0]-(0.123)*dt
    nmodl_eigen_j[1] = nmodl_eigen_j[1]+(0.123)*dt
    nmodl_eigen_j[3] = nmodl_eigen_j[3]-(0.456)*dt
    nmodl_eigen_j[2] = nmodl_eigen_j[2]+(0.456)*dt
} CONSERVE = 0.789

Ex. 2

Non-linear ODEs

The matexp solver checks for linear equations and raises exception if it finds any non-linearities. It asserts that each reaction has exactly one reactant and one product. This means that the following statements will cause errors:

KINETIC nonlinear_errors {
    ~ A <-> B + C (kf, kr) : Multiple products!
    ~ 2 A <-> C (kf, kr)   : Multiple reactants!
    ~ A << (kf)            : Zero reactants!
}

However the solver does not check that the kinetic rates are independent from state variables, and so the following non-linear equation will be accepted by the program and solved inaccurately.

KINETIC nonlinear_silent_bug {
    ~ A <-> B (func(A), kr)
}