nonlinz.cpp

Go to the documentation of this file.

#include <../../nrnconf.h>
#include <stdio.h>
#include <math.h>
#include <InterViews/resource.h>
#include "nonlinz.h"
#include "nrnoc2iv.h"
#include "nrnmpi.h"
#include "cspmatrix.h"
#include "membfunc.h"

extern "C" int structure_change_cnt;
extern void v_setup_vectors();
extern void nrn_rhs(NrnThread*);
extern int nrndae_extra_eqn_count();
extern Symlist *hoc_built_in_symlist;
extern void (*nrnthread_v_transfer_)(NrnThread*);
extern spREAL *spGetElement(char*, int ,int);

extern void pargap_jacobi_rhs(double*, double*);
extern void pargap_jacobi_setup(int mode);

class NonLinImpRep {
public:
    NonLinImpRep();
    virtual ~NonLinImpRep();
    void delta(double);
    void didv();
    void dids();
    void dsdv();
    void dsds();
    int gapsolve();

    char* m_;
    int scnt_; // structure_change
    int n_v_, n_ext_, n_lin_, n_ode_, neq_v_, neq_;
    double** pv_;
    double** pvdot_;
    int* v_index_;
    double* rv_;
    double* jv_;
    double** diag_;
    double* deltavec_; // just like cvode.atol*cvode.atolscale for ode's
    double delta_; // slightly more efficient and easier for v.
    void current(int, Memb_list*, int);
    void ode(int, Memb_list*);

    double omega_;
    int iloc_; // current injection site of last solve
    float* vsymtol_;
    int maxiter_;
};

NonLinImp::NonLinImp() {
    rep_ = NULL;
}
NonLinImp::~NonLinImp() {
    if (rep_) {
        delete rep_;
    }
}
double NonLinImp::transfer_amp(int curloc, int vloc) {
    if (nrnmpi_numprocs > 1 && nrnthread_v_transfer_ && curloc != rep_->iloc_) {
          hoc_execerror("current injection site change not allowed with both gap junctions and nhost > 1", 0);
    }
    if (curloc != rep_->iloc_) {
      solve(curloc);
    }
    double x = rep_->rv_[vloc];
    double y = rep_->jv_[vloc];
    return sqrt(x*x+y*y);
}
double NonLinImp::input_amp(int curloc) {
    if (nrnmpi_numprocs > 1 && nrnthread_v_transfer_) {
      hoc_execerror("not allowed with both gap junctions and nhost>1", 0);
    }
    if (curloc != rep_->iloc_) {
      solve(curloc);
    }
    if (curloc < 0) { return 0.0; }
    double x = rep_->rv_[curloc];
    double y = rep_->jv_[curloc];
    return sqrt(x*x+y*y);
}
double NonLinImp::transfer_phase(int curloc, int vloc) {
    if (nrnmpi_numprocs > 1 && nrnthread_v_transfer_ && curloc != rep_->iloc_) {
          hoc_execerror("current injection site change not allowed with both gap junctions and nhost > 1", 0);
    }
    if (curloc != rep_->iloc_) {
      solve(curloc);
    }
    double x = rep_->rv_[vloc];
    double y = rep_->jv_[vloc];
    return atan2(y, x);
}
double NonLinImp::input_phase(int curloc) {
    if (nrnmpi_numprocs > 1 && nrnthread_v_transfer_) {
      hoc_execerror("not allowed with both gap junctions and nhost>1", 0);
    }
    if (curloc != rep_->iloc_) {
      solve(curloc);
    }
    if (curloc < 0) { return 0.0; }
    double x = rep_->rv_[curloc];
    double y = rep_->jv_[curloc];
    return atan2(y,x);
}
double NonLinImp::ratio_amp(int clmploc, int vloc) {
    if (nrnmpi_numprocs > 1 && nrnthread_v_transfer_) {
      hoc_execerror("not allowed with both gap junctions and nhost>1", 0);
    }
    if (clmploc < 0) { return 0.0; }
    if (clmploc != rep_->iloc_) {
      solve(clmploc);
    }
    double ax,bx,cx, ay,by,cy,bb;
    ax = rep_->rv_[vloc];
    ay = rep_->jv_[vloc];
    bx = rep_->rv_[clmploc];
    by = rep_->jv_[clmploc];
    bb = bx*bx + by*by;
    cx = (ax*bx + ay*by)/bb;
    cy = (ay*bx - ax*by)/bb;
    return sqrt(cx*cx + cy*cy);
}
void NonLinImp::compute(double omega, double deltafac, int maxiter) {
    v_setup_vectors();
    nrn_rhs(nrn_threads);
    if (rep_ && rep_->scnt_ != structure_change_cnt) {
        delete rep_;
        rep_ = NULL;
    }
    if (!rep_) {
        rep_ = new NonLinImpRep();
    }
    rep_->maxiter_ = maxiter;
    if (rep_->neq_ == 0) { return; }
    if (nrndae_extra_eqn_count() > 0) {
        hoc_execerror("Impedance calculation with LinearMechanism not implemented", 0);
    }
    if (nrn_threads->_ecell_memb_list) {
        hoc_execerror("Impedance calculation with extracellular not implemented", 0);
    }

    rep_->omega_ = 1000.* omega;
    rep_->delta(deltafac);
    // fill matrix
    cmplx_spClear(rep_->m_);
    rep_->didv();
    rep_->dsds();
#if 1   // when 0 equivalent to standard method
    rep_->dids();
    rep_->dsdv();
#endif
    
//  cmplx_spPrint(rep_->m_, 0, 1, 1);
//  for (int i=0; i < rep_->neq_; ++i) {
//      printf("i=%d %g %g\n", i, rep_->diag_[i][0], rep_->diag_[i][1]);
//  }
    int e = cmplx_spFactor(rep_->m_);
    switch (e) {
    case spZERO_DIAG:
        hoc_execerror("cmplx_spFactor error:", "Zero Diagonal");
    case spNO_MEMORY:
        hoc_execerror("cmplx_spFactor error:", "No Memory");
    case spSINGULAR:
        hoc_execerror("cmplx_spFactor error:", "Singular");
    }  

    rep_->iloc_ = -2;
}

int NonLinImp::solve(int curloc) {
    int rval = 0;
    NrnThread* _nt = nrn_threads;
    if (!rep_) {
        hoc_execerror("Must call Impedance.compute first", 0);
    }
    if (rep_->iloc_ != curloc) {
        int i;
        rep_->iloc_ = curloc;
        for (i=0; i < rep_->neq_; ++i) {
            rep_->rv_[i] = 0;
            rep_->jv_[i] = 0;
        }
        if (curloc >= 0) {rep_->rv_[curloc] = 1.e2/NODEAREA(_nt->_v_node[curloc]);}
        if (nrnthread_v_transfer_) {
          rval = rep_->gapsolve();
        }else{
          assert(rep_->m_);
          cmplx_spSolve(rep_->m_, rep_->rv_-1, rep_->rv_-1, rep_->jv_-1, rep_->jv_-1);
        }
    }
    return rval;
}

// too bad it is not easy to reuse the cvode/daspk structures. Most of
// mapping is already done there.

NonLinImpRep::NonLinImpRep() {
    int err;
    int i, j, ieq, cnt;
    NrnThread* _nt = nrn_threads;
    maxiter_ = 500;
    m_ = NULL;

    vsymtol_ = NULL;
    Symbol* vsym = hoc_table_lookup("v", hoc_built_in_symlist);
    if (vsym->extra) {
        vsymtol_ = &vsym->extra->tolerance;
    }
    // the equation order is the same order as for the fixed step method
    // for current balance. Remaining ode equation order is
    // defined by the memb_list.

    
    // how many equations
    n_v_ = _nt->end;
    n_ext_ = 0;
    if (_nt->_ecell_memb_list) {
        n_ext_ = _nt->_ecell_memb_list->nodecount*nlayer;
    }
    n_lin_ = nrndae_extra_eqn_count();
    n_ode_ = 0;
    for (NrnThreadMembList* tml = _nt->tml; tml; tml = tml->next) {
        Memb_list* ml = tml->ml;
        i = tml->index;
        nrn_ode_count_t s = memb_func[i].ode_count;
        if (s) {
            cnt = (*s)(i);
            n_ode_ += cnt * ml->nodecount;
        }
    }
    neq_v_ = n_v_ + n_ext_ + n_lin_;
    neq_ = neq_v_ + n_ode_;
    if (neq_ == 0) { return; }
    m_ = cmplx_spCreate(neq_, 1, &err);
    assert(err == spOKAY);
    pv_ = new double*[neq_];
    pvdot_ = new double*[neq_];
    v_index_ = new int[n_v_];
    rv_ = new double[neq_+1];
    rv_ += 1;
    jv_ = new double[neq_+1];
    jv_ += 1;
    diag_ = new double*[neq_];
    deltavec_ = new double[neq_];

    for (i=0; i < n_v_; ++i) {
        // utilize nd->eqn_index in case of use_sparse13 later
        Node* nd = _nt->_v_node[i];
        pv_[i] = &NODEV(nd);
        pvdot_[i] = nd->_rhs;
        v_index_[i] = i+1;
    }
    for (i=0; i < n_v_; ++i) {
        diag_[i] = cmplx_spGetElement(m_, v_index_[i], v_index_[i]);
    }
    for (i=neq_v_; i < neq_; ++i) {
        diag_[i] = cmplx_spGetElement(m_, i+1, i+1);
    }
    scnt_ = structure_change_cnt;
}

NonLinImpRep::~NonLinImpRep() {
    if (!m_) { return; }
    cmplx_spDestroy(m_);
    delete [] pv_;
    delete [] pvdot_;
    delete [] v_index_;
    delete [] (rv_ - 1);
    delete [] (jv_ - 1);
    delete [] diag_;
    delete [] deltavec_;
}

void NonLinImpRep::delta(double deltafac){ // also defines pv_,pvdot_ map for ode
    int i, j, nc, cnt, ieq;
    NrnThread* nt = nrn_threads;
    for (i=0; i < neq_; ++i) {
        deltavec_[i] = deltafac; // all v's wasted but no matter.
    }
    ieq = neq_v_;
    for (NrnThreadMembList* tml = nt->tml; tml; tml = tml->next) {
        Memb_list* ml = tml->ml;
        i = tml->index;
        nc = ml->nodecount;
        nrn_ode_count_t s = memb_func[i].ode_count;
        if (s && (cnt = (*s)(i)) > 0) {
            nrn_ode_map_t m = memb_func[i].ode_map;
            for (j=0; j < nc; ++j) {
(*m)(ieq, pv_ + ieq, pvdot_ + ieq, ml->data[j], ml->pdata[j], deltavec_ + ieq, i);
                ieq += cnt;
            }
        }
    }
    delta_ = (vsymtol_ && (*vsymtol_ != 0.)) ? *vsymtol_ : 1.;
    delta_ *= deltafac;
    
}

void NonLinImpRep::didv() {
    int i, j, ip;
    Node* nd;
    NrnThread* _nt = nrn_threads;
    // d2v/dv2 terms
    for (i=_nt->ncell; i < n_v_; ++i) {
        nd = _nt->_v_node[i];
        ip = _nt->_v_parent[i]->v_node_index;
        double* a = cmplx_spGetElement(m_, v_index_[ip], v_index_[i]);
        double* b = cmplx_spGetElement(m_, v_index_[i], v_index_[ip]);
        *a += NODEA(nd);
        *b += NODEB(nd);
        *diag_[i] -= NODEB(nd);
        *diag_[ip] -= NODEA(nd);
    }
    // jwC term
    Memb_list* mlc = _nt->tml->ml;
    int n = mlc->nodecount;
    for (i=0; i < n; ++i) {
        double* cd = mlc->data[i];
        j = mlc->nodelist[i]->v_node_index;
        diag_[v_index_[j]-1][1] += .001 * cd[0] * omega_;
    }
    // di/dv terms
    // because there may be several point processes of the same type
    // at the same location, we have to be careful to neither increment that
    // nd->v multiple times nor count the rhs multiple times.
    // So we can't take advantage of vectorized point processes.
    // To avoid this we do each mechanism item separately.
    // We assume there is no interaction between
    // separate locations. Note that interactions such as gap junctions
    // would not be handled in any case without computing a full jacobian.
    // i.e. calling nrn_rhs varying every state one at a time (that would
    // give the d2v/dv2 terms as well), but the expense is unwarranted.
    for (NrnThreadMembList* tml = _nt->tml; tml; tml = tml->next) {
        i = tml->index;
        if (i == CAP) { continue; }
        if (!memb_func[i].current) { continue; }
        Memb_list* ml = tml->ml;
        for (j = 0; j < ml->nodecount; ++j) { Node* nd = ml->nodelist[j];
            double x1;
            double x2;
            // zero rhs
            // save v
            NODERHS(nd) = 0;
            x1 = NODEV(nd);
            // v+dv
            NODEV(nd) += delta_;
            current(i, ml, j);
            // save rhs
            // zero rhs
            // restore v
            x2 = NODERHS(nd);
            NODERHS(nd) = 0;
            NODEV(nd) = x1;
            current(i, ml, j);
            // conductance
            // add to matrix
            *diag_[v_index_[nd->v_node_index]-1] -= (x2 - NODERHS(nd))/delta_;
        }
    }
}   

void NonLinImpRep::dids() {
    // same strategy as didv but now the innermost loop is over
    // every state but just for that compartment
    // outer loop over ode mechanisms in same order as we created the pv_ map   // so we can eas
    int ieq, i, in, is, iis;
    NrnThread* nt = nrn_threads;
    ieq = neq_ - n_ode_;
    for (NrnThreadMembList* tml = nt->tml; tml; tml = tml->next) {
      Memb_list* ml = tml->ml;
      i = tml->index;
      if (memb_func[i].ode_count && ml->nodecount) {
        int nc = ml->nodecount;
        nrn_ode_count_t s = memb_func[i].ode_count;
        int cnt = (*s)(i);
        if (memb_func[i].current) {
        double* x1 = rv_; // use as temporary storage
        double* x2 = jv_;
        for (in = 0; in < ml->nodecount; ++in) { Node* nd = ml->nodelist[in];
            // zero rhs
            NODERHS(nd) = 0;
            // compute rhs
            current(i, ml, in);
            // save rhs
            x2[in] = NODERHS(nd);
            // each state incremented separately and restored
            for (iis = 0; iis < cnt; ++iis) {           
                is = ieq + in*cnt + iis;
                // save s
                x1[is] = *pv_[is];
                // increment s and zero rhs
                *pv_[is] += deltavec_[is];
                NODERHS(nd) = 0;
                current(i, ml, in);
                *pv_[is] = x1[is]; // restore s
                double g = (NODERHS(nd) - x2[in])/deltavec_[is];
                if (g != 0.) {
                    double* elm = cmplx_spGetElement(m_, v_index_[nd->v_node_index], is+1);
                    elm[0] = -g;
                }
            }
            // don't know if this is necessary but make sure last
            // call with respect to original states
            current(i, ml, in);
        }
        }
        ieq += cnt*nc;
      }
    }
}

void NonLinImpRep::dsdv() {
    int ieq, i, in, is, iis;
    NrnThread* nt = nrn_threads;
    ieq = neq_ - n_ode_;
    for (NrnThreadMembList* tml = nt->tml; tml; tml = tml->next) {
      Memb_list* ml = tml->ml;
      i = tml->index;
      if (memb_func[i].ode_count && ml->nodecount) {
        int nc = ml->nodecount;
        nrn_ode_count_t s = memb_func[i].ode_count;
        int cnt = (*s)(i);
        if (memb_func[i].current) {
        double* x1 = rv_; // use as temporary storage
        double* x2 = jv_;
        // zero rhs, save v
        for (in = 0; in < ml->nodecount; ++in) { Node* nd = ml->nodelist[in];
            for (is = ieq + in*cnt, iis = 0; iis < cnt; ++iis, ++is) {
                *pvdot_[is] = 0.;
            }
            x1[in] = NODEV(nd);
        }
        // increment v only once in case there are multiple
        // point processes at the same location
        for (in = 0; in < ml->nodecount; ++in) { Node* nd = ml->nodelist[in];
            if (x1[in] == NODEV(nd)) {
                NODEV(nd) += delta_;
            }
        }
        // compute rhs. this is the rhs(v+dv)
        ode(i, ml);
        // save rhs, restore v, and zero rhs
        for (in = 0; in < ml->nodecount; ++in) { Node* nd = ml->nodelist[in];
            for (is = ieq + in*cnt, iis = 0; iis < cnt; ++iis, ++is) {
                x2[is] = *pvdot_[is];
                *pvdot_[is] = 0;
            }
            NODEV(nd) = x1[in];
        }
        // compute the rhs(v)
        ode(i, ml);
        // fill the ds/dv elements
        for (in = 0; in < ml->nodecount; ++in) { Node* nd = ml->nodelist[in];
            for (is = ieq + in*cnt, iis = 0; iis < cnt; ++iis, ++is) {
                double ds = (x2[is] - *pvdot_[is])/delta_;
                if (ds != 0.) {
                    double* elm = cmplx_spGetElement(m_, is+1, v_index_[nd->v_node_index]);
                    elm[0] = -ds;
                }
            }
        }
        }
        ieq += cnt*nc;
      }
    }
}

void NonLinImpRep::dsds() {
    int ieq, i, in, is, iis, ks, kks;
    NrnThread* nt = nrn_threads;
    // jw term
    for (i = neq_v_; i < neq_; ++i) {
        diag_[i][1] += omega_;
    }
    ieq = neq_v_;
    for (NrnThreadMembList* tml = nt->tml; tml; tml = tml->next) {
      Memb_list* ml = tml->ml;
      i = tml->index;
      if (memb_func[i].ode_count && ml->nodecount) {
        int nc = ml->nodecount;
        nrn_ode_count_t s = memb_func[i].ode_count;
        int cnt = (*s)(i);
        double* x1 = rv_; // use as temporary storage
        double* x2 = jv_;
        // zero rhs, save s
        for (in = 0; in < ml->nodecount; ++in) {
            for (is = ieq + in*cnt, iis = 0; iis < cnt; ++iis, ++is) {
                *pvdot_[is] = 0.;
                x1[is] = *pv_[is];
            }
        }
        // compute rhs. this is the rhs(s)
        ode(i, ml);
        // save rhs
        for (in = 0; in < ml->nodecount; ++in) {
            for (is = ieq + in*cnt, iis = 0; iis < cnt; ++iis, ++is) {
                x2[is] = *pvdot_[is];
            }
        }
        // iterate over the states
        for (kks=0; kks < cnt; ++kks) {
        // zero rhs, increment s(ks)
        for (in = 0; in < ml->nodecount; ++in) {
            ks = ieq + in*cnt + kks;
            for (is = ieq + in*cnt, iis = 0; iis < cnt; ++iis, ++is) {
                *pvdot_[is] = 0.;
            }
            *pv_[ks] += deltavec_[ks];
        }
        ode(i, ml);
        // store element and restore s
        // fill the ds/dv elements
        for (in = 0; in < ml->nodecount; ++in) { Node* nd = ml->nodelist[in];
            ks = ieq + in*cnt + kks;
            for (is = ieq + in*cnt, iis = 0; iis < cnt; ++iis, ++is) {
                double ds = ( *pvdot_[is] - x2[is])/deltavec_[is];
                if (ds != 0.) {
                    double* elm = cmplx_spGetElement(m_, is+1, ks+1);
                    elm[0] = -ds;
                }
                *pv_[ks] = x1[ks];
            }
        }
        // perhaps not necessary but ensures the last computation is with
        // the base states.
        ode(i, ml);
        }
        ieq += cnt*nc;
      }
    }
}

void NonLinImpRep::current(int im, Memb_list* ml, int in) {  // assume there is in fact a current
                                                             // method
    Pvmi s = memb_func[im].current;
    // fake a 1 element memb_list
    Memb_list mfake;
#if CACHEVEC != 0
    mfake.nodeindices = ml->nodeindices + in;
#endif
    mfake.nodelist = ml->nodelist + in;
    mfake.data = ml->data + in;
    mfake.pdata = ml->pdata + in;
    mfake.prop = ml->prop ? ml->prop + in : nullptr;
    mfake.nodecount = 1;
    mfake._thread = ml->_thread;
    (*s)(nrn_threads, &mfake, im);
}

void NonLinImpRep::ode(int im, Memb_list* ml) { // assume there is in fact an ode method
    Pvmi s = memb_func[im].ode_spec;
    (*s)(nrn_threads, ml, im);
}

int NonLinImpRep::gapsolve() {
  // On entry, rv_ and jv_ contain the complex b for A*x = b.
  // On return rv_ and jv_ contain complex solution, x.
  // m_ is the factored matrix for the trees without gap junctions
  // Jacobi method (easy for parallel)
  // A = D + R
  // D*x_(k+1) = (b - R*x_(k))
  // D is m_ (and includes the gap junction contribution to the diagonal)
  // R is the off diagonal matrix of the gaps.

  // one and only one stimulus
#if NRNMPI
  if (nrnmpi_numprocs > 1 && nrnmpi_int_sum_reduce(iloc_ >= 0 ? 1 : 0) != 1) {
    if (nrnmpi_myid == 0) {
      hoc_execerror("there can be one and only one impedance stimulus", 0);
    }
  }
#endif

  pargap_jacobi_setup(0);

  double *rx, *jx, *rx1, *jx1, *rb, *jb;
  if (neq_) {
    rx = new double[neq_];
    jx = new double[neq_];
    rx1 = new double[neq_];
    jx1 = new double[neq_];
    rb = new double[neq_];
    jb = new double[neq_];
  }

  // initialize for first iteration
  for (int i=0; i < neq_; ++i) {
    rx[i] = jx[i] = 0.0;
    rb[i] = rv_[i];
    jb[i] = jv_[i];
  }

  // iterate till change in x is small
  double tol = 1e-9;
  double delta;
  
  int success = 0;
  int iter;

  for (iter = 1; iter <= maxiter_; ++iter) {
    if (neq_) {
      cmplx_spSolve(m_, rb-1, rx1-1, jb-1, jx1-1);
    }
    
    // if any change in x > tol, then do another iteration.
    success = 1;
    delta = 0.0;
    for (int i=0; i < neq_; ++i) {
      double err = fabs(rx1[i] - rx[i]) + fabs(jx1[i] - jx[i]);
      if (err > tol) {
        success = 0;
      }
      if (delta < err) {
        delta = err;
      }
    }
#if NRNMPI
    if (nrnmpi_numprocs > 1) {
      success = nrnmpi_int_sum_reduce(success) / nrnmpi_numprocs;
    }
#endif
    if (success) {
      for (int i=0; i < neq_; ++i) {
        rv_[i] = rx1[i];
        jv_[i] = jx1[i];
      }
      break;
    }

    // setup for next iteration
    for (int i=0; i < neq_; ++i) {
      rx[i] = rx1[i];
      jx[i] = jx1[i];
      rb[i] = rv_[i];
      jb[i] = jv_[i];
    }
    pargap_jacobi_rhs(rb, rx);
    pargap_jacobi_rhs(jb, jx);
  }

  pargap_jacobi_setup(1); // tear down

  if (neq_) {
    delete [] rx;
    delete [] jx;
    delete [] rx1;
    delete [] jx1;
    delete [] rb;
    delete [] jb;
  }

  if (!success) {
    char buf[256];
    sprintf(buf, "Impedance calculation did not converge in %d iterations. Max state change on last iteration was %g (Iterations stop at %g)\n",
      maxiter_, delta, tol);
    execerror(buf, 0);
  }
  return iter;
}