A PHYSICS-INFORMED NEURAL NETWORK FOR COUPLED CALCIUM DYNAMICS IN A CABLE NEURON

Authors: zachary m miksis
  gillian queisser
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-01-12 22:52:33

Abstract: Transcranial magnetic stimulation (TMS) is a noninvasive treatment for a variety of neurological and neuropsychiatric disorders by triggering a calcium response through magnetic stimulation. To understand the full effects of this treatment, researchers will often use numerical simulations to model and study the calcium response. These simulations are limited to short-time simulations of single neurons due to computational complexity, restricting their use in clinical settings. In this paper, we explore an application of physics-informed neural networks (PINNs) to accurately produce long-time simulations of neuronal responses, opening the possibility of utilizing these methods in clinical applications to directly benefit patients.

Main paper: 
 introduction: 
 Transcranial magnetic stimulation (TMS) is an invaluable tool for treating a variety of neurological and neuropsychiatric disorders in a noninvasive way by using a time-varying magnetic field passed through the brain to stimulate neurons (Barker et al., 1985; Hallett, 2007) . This stimulation triggers a response from intracellular calcium, which is vital to regulating the transfer of information from synaptic sites to the cell nucleus (Borole et al., 2023) . Ultimately, this allows for treatment of various forms of degenerative neurological disorders, and TMS is therefore used extensively in both research and clinical settings (Lefaucheur et al., 2014) .While TMS is very important, its full effects are still not entirely understood. To investigate these effects, computational tools and studies are exceptionally important to complement experimental studies (Grein et al., 2014) . To this end, different numerical computation tools have been developed (Borole, 2022; Shirinpour et al., 2021; Guan & Queisser, 2022) . While these tools are very good at capturing fine grain details and the quick responses to fast TMS frequencies, they are limited to short simulated time spans of single neurons due to computational complexity. In order to address these limitations, we look to neural networks (from here on, "neuron" will be used to refer to a biological neuron, and "neural network" will be used to refer to the machine learning paradigm).Neural networks have been used in various applications, many focusing on image recognition and reconstruction (e.g., see Abraham et al. (2023) ). More recently, knowledge of physical models have been introduced into neural networks to form physics-informed neural networks (PINNs) (Raissi et al., 2019; 2017a; b) . In this work, we utilize these PINNs to incorporate partial differential equations (PDEs) that model calcium dynamics in neurons into the larger neural network. With a network that satisfies physical laws, we can build simulations that provide long-time results, overcoming one of the limitations of traditional numerical methods. This leads to another limitation of PINNs, in that they have been shown to have difficulty accurately to simulate diffusion problems (Saadat et al., 2022) . This can be overcome with strategies of both relaxation of the loss function and periodic activation functions (Snyder et al., 2023; Sitzmann et al., 2020) , which we utilize in our application. This paper is structured as follows: In Section 2, we describe the physical model of calcium dynamics that is used (with further details in Appendix A and B). In Section 3, we describe the details on PINNs as we used them in our application. In Section 4, we present results of a simulated cable neuron, with additional plots in Appendix C. In Section 5, we draw conclusions from our results and describe the broader impacts of this work. 
calcium dynamics: 
 Within a neuron, ion dynamics may be modeled by general 1D reaction-diffusion equations of the formEQUATIONwhere u is an ion concentration, D is the diffusion coefficient, and R(u) is a reaction term. From Fick's first law, we haveEQUATION)where J is the diffusion flux of the ion. This leads to the Neumann boundary conditionEQUATIONIn the cytosol, the non-organelle interior of a neuron, calcium is transported via various mechanisms and is buffered by the molecule calbindin. We model this using a dimension-reduced system (Borole et al., 2023) , given byEQUATIONEQUATIONwhere the reaction termEQUATIONand J P M represents the net Ca 2+ ion flux across the plasma membrane (separating intra-and extracellular space). This term can be broken down into the different mechanisms that transport calcium across the plasma membrane,EQUATIONwhere J P is the flux from plasma membrane Ca 2+ -ATPase pumps (PMCA), J N is the flux from Na + /Ca 2+ exchangers (NCX), J SY N is the flux through the post-synaptic density (PSD), and J V DCC is the flux from voltage dependent calcium channels (VDCCs). Details of how each flux term is calculated are given in Appendix A. 
physics-informed neural networks: 
 To simulate calcium dynamcis, we incorporate the above diffusion model into a physics-informed neural network (PINN) (Raissi et al., 2019; 2017a; b) . Consider a general partial differential equation of the formu t + N [u] = 0, (8) with x ∈ Ω, t ∈ [0, T ], where u(x, t) is the latent solution, N [•] is a nonlinear differential operator,and Ω is the computational domain. Define f (x, t) as the left-hand-side of Eq. ( 8),EQUATIONBy approximating u(x, t) using a deep neural network, we build a physics-informed neural network f (x, t) by combining this network with Eq. ( 9). The shared terms of the networks are learned by minimizing the mean squared error lossEQUATIONwhereM SE u = 1 N u Nu i=1 |u(t i u , x i u ) -u i | 2 (11)represents the error of the initial condition,M SE f = 1 N f N f i=1 |f (t i f , x i f )| 2 (12)represents the error of the operator f inside the domain, andM SE b = 1 N b N b i=1 |u x (t i b , x i b ) + J/D| 2 (13)represents the error of the Neumann boundary conditions. Here,{t i u , x i u , u i } Nu i=1 represents the initial condition training data, {t i f , x i f } N fi=1 represents the collocation points within the domain (excluding the boundaries), and {t i b , x i b } N b i=1 represents the boundary collocation points. In our minimization, we apply a relaxation parameter α ∈ (0, 1]. This is selected a priori by trial and error, and may require tuning to fit a particular model (Snyder et al., 2023) .As demonstrated in Saadat et al. (2022) and Sitzmann et al. ( 2020), diffusion problems such as the calcium model are partiularly difficult to model accurately with neural networks. PINNs were initally introduced with tanh activation function, and were demonstrated to perform well for advectiondominated problems (Raissi et al., 2019) . The SIREN network introduced the use of periodic activation functions (Sitzmann et al., 2020) , which have been demonstrated to correct for problems in modeling diffusion-dominated problems (Saadat et al., 2022) . Therefore, based on the results of the comparison of sin and tanh activation functions for presented in (Saadat et al., 2022) , we utilize a sin activation function in our network to provide consistent long-time modeling.In this section, we demonstrate the effectiveness of PINNs in modeling calcium dynamics. All constants used are given in Appendix B. We consider the computational domain (x, t) : [-1, 1] × [0, 1], and use a numerically computed solution as the exact solution (see Borole (2022) for details on numerical methods). The network consists of 4 hidden layers of 50 computational neurons each. The network is initialized with a uniform Xavier scheme, and an Adam optimizer with learning rate of 0.001 is used to minimize the M SE. N = 150 collocation points are randomly generated on both the interior of the domain and the boundary (50 on each spatial boundary and 50 on the initial condition). These can be seen in Fig. 1 .The full set of data points used is the discretized (x, t) : [-1, 1] × [0, 1] domain consisting of 101 spatial gridpoints and 7,193 temporal steps (for a total of 726,493 spatiotemporal data points). Of these points, the training set only consists of the 50 known initial condition (t = 0) points (defined by Eq. ( 11)), the validation set consists of the 50 points on each spatiotemporal boundary and 150 points on the domain interior that must satisfy the PDE (for a total of 250 data points, defined by Eqs. ( 12)-( 13)), and the remaining data points makeup the test set.It is important to note that while this case seems like a fairly simple problem that may be solved without PINN, this is only because this test case has been chosen to be small in early stages of this work. In a realistic application of simulating a full network of physical neurons over a full TMS treatment regime, a traditional numerical method will require tens of millions of timesteps over hundreds of thousands of spatial nodes. Even though this is a one-dimensional PDE, this quickly becomes an intractable simulation to run in any reasonable amount of time for use in a clinical setting. A pretrained PINN may be able to provide simulated results without significant computing resources, and can move this work beyond well-funded research and into a clinical setting. For the biological model, we introduce a constant stimulus at the soma for the entire simulation time, and a constant voltage is applied. As can be seen in Appendix A, the gating function for VDCCs requires the solution of two ODEs. For the purposes of simulating TMS treatment protocols, predetermined voltage data may be provided to review its effects. Since the gating function only requires voltage and temporal values, these gating probabilities may be precomputed prior to training the network. We therefore avoid the requirement of modeling additional ODEs, and provide these gating probabilities explicitly to the network.We first look at the values of the initial calcium and calbindin concentrations. In the numerical simulation, the initial concentrations are given as constant values. These may not represent the true steady state, and they will not satisfy the diffusion model, which leads to error in the PINN if they are used. We therefore consider the initial condition as the computed values after a small number of numerical timesteps, as the concentrations are approaching values that satisfy the diffusion model. As the number of numerical timesteps is selected a priori, the values are still not guaranteed to be consistent with the diffusion model, but this significantly improves the performance of the PINN. In Fig. 2 and Fig. 3 , we show the results of the network prediction of the initial concentrations. While there is a small amount of error, we can see that the network fairly accurately predicts these values. In Fig. 4 and Fig. 5 , we see that this error is significantly reduced to nearly vanishing as the simulation moves forward in time and the concentration values become consistent with the diffusion model. In this study, we presented a physics-informed neural network (PINN) for modeling a coupled calcium model in a cable neuron. This is significant for the study of repetitive transcranial magnetic stimulation (rTMS), as it allows for long-time simulations with near instantaneous results. While traditional numerical methods provide accurate solutions, the requirement of small timesteps results in significant computation time and hampering the use for simulating a full rTMS treatment that could last several minutes with voltage pulses in millisecond bursts (Shirinpour et al., 2021) . Training a PINN on a short timespan allows for feeding data into the fully trained network and producing a long simulation.Our results show that PINNs are capable of accurately capturing the calcium concentration over a cable neuron, which can be extended over time. While this work is limited to a single cable neuron, using this to create branches, and therefore a full neuron model, is an important direction for future work. 
broader impact appendix a: membrane flux components: 
 The plasma membrane flux in this study can be separated asEQUATIONIn this section, we give a brief description of each component as described in Borole et al. (2023) .PMCA Pumps A second-order Hill equation is used to model the plasma membrane Ca 2+ current of PMCA pumps,EQUATIONwhere ρ P is the density of PMCA pumps on the plasma membrane, I P is the single channel Ca 2+ current, c c is the cytosolic Ca 2+ concentration, and K P is the measure of Ca 2+ affinity.NCX Exchangers A second-order Hill equation is used to model the plasma membrane Ca 2+ current of NCX exchangers,EQUATION)where ρ N is the density of NCX exchangers on the plasma membrane, I N is the single channel Ca 2+ current, c c is the cytosolic Ca 2+ concentration, and K N is the measure of Ca 2+ affinity.EQUATIONEQUATIONVDCCs For VDCCs, we use a Borg-Graham model. The Ca 2+ current is given byEQUATIONwhere G(V, t) ∈ [0, 1] is the gating function and F (V, ∆[Ca 2+ ]) is the flux function. The difference between cytoplasmic and extracellular ion concentration is given by∆[Ca 2+ ] = c c -c o ,and the flux function is given by the Goldman-Hodgkin-Katz equation,EQUATION)where R is the universal gas constant, F is Faraday's constant, T is temperature in Kelvin, p Ca 2+ is the permeability of Ca 2+ ions through the channels, and z is the valence of a Ca 2+ ion. The gating function for N-type VDCCs is given byEQUATIONwhere the gating functions k(•) and l(•) satisfy the two ODEs∂k ∂t = k ∞ -k τ k and ∂l ∂t = l ∞ -l τ l ,wherek ∞ = α k (V ) α k (V ) + β k (V ) ; l ∞ = α l (V ) α l (V ) + β l (V ),and τ k = 1 α k + β k + τ k,0 ; τ l = 1 α l + β l + τ l,0 .The rate functions are defined asα k (V ) = K k exp z k γ k (V -V 1/2,k )F RT β k (V ) = K k exp -z k (1 -γ k )(V -V 1/2,k )F RT α l (V ) = K l exp z l γ l (V -V 1/2,l )F RT β l (V ) = K l exp -z l (1 -γ l )(V -V 1/2,l )F RTTranscranial Magnetic Stimulation is an important noninvasive treatment for a wide array of neurological disorders. While the neuronal response to TMS is still poorly understood, simulation tools are providing a way of studying these effects and how they will benefit patients. A limiting factor in this research has been the difficulty in developing long-time simulations using traditional numerical methods. Neural networks provide the ability to provide long-time simulations with ease that may be applied in both a research and clinical setting, benefitting researchers and patients alike. Our work is an important step in that direction, and has the potential to fill this need. 
 
 Appendix:  


Decision: The paper proposes a PINN based framework for magnetic simulation of calcium dynamics. The reviewers pointed out a few important things or questions which would have been great that authors could address it. That makes me be at the reject/accept border about this work. Due to their effort on using AI for this problem and under the condition that they address the reviewers' questions and also clarify a few parts mentioned in the review, I will vote for accept (poster).

Reviews: review: title: title: Are PINN's really needed for the application?
main_review: main_review: The paper proposes the use of a PINN to model calcium dynamics in (biological neurons). This is motivated by the high computational costs of simulating this using conventional methods which prevents long-term simulations. However, since the proposed PINN has 4 x 50 = 200 parameters, I am quite surprised that this is infeasible to simulate conventionally (also, it's a 1D PDE and it's not clear why/when it is stiff). The numerical results show very good agreement between conventional simulation and PINN, but I again wondering why a (artificial) neural network is needed when high accuracy can be achieved with such a simple neural network.
title: title: An exploration on applying PINNs to predicting calcium dynamics
main_review: main_review: This work details a PINN for modeling calcium dynamics, with application to Transcranial Magnetic Stimulation (TMS). The work is exploratory in nature and mentions that the goal is to extend the work to full neuron model and to simulations that extend over several minutes.

I know little about TMS, but I can appreciate the effort it appears to have taken to construct this PINN. 

Some experimental details are unclear, particularly with respect to the training/val/test data splits. It looks like the PINN is trained and tested on the same temporal range `t=[0,1]`, and I would have expected to see some results where it was trained on, say, `t=[0,0.8)` and tested on `t=[0.80,1]`. The results in Figure 2 and Figure 3 look great! However, the current test split likely isn't challenging enough to inform about how well this model will generalize in practice. An ablation study to quantify the improvement gained by using the `sin` activation function instead of `tanh` is missing as well.