NAC as in https://arxiv.org/abs/1808.00508.
778666
def nac(x, depth, name=None, reuse=None):
  """"""
  with tf.variable_scope(name, default_name="nac", values=[x], reuse=reuse):
    x_shape = shape_list(x)
    w = tf.get_variable("w", [x_shape[-1], depth])
    m = tf.get_variable("m", [x_shape[-1], depth])
    w = tf.tanh(w) * tf.nn.sigmoid(m)
    x_flat = tf.reshape(x, [-1, x_shape[-1]])
    res_flat = tf.matmul(x_flat, w)
    return tf.reshape(res_flat, x_shape[:-1] + [depth])
1321970
def _dsinit(
       whichconst,
       cosim,  emsq,   argpo,   s1,     s2,
       s3,     s4,     s5,      sinim,  ss1,
       ss2,    ss3,    ss4,     ss5,    sz1,
       sz3,    sz11,   sz13,    sz21,   sz23,
       sz31,   sz33,   t,       tc,     gsto,
       mo,     mdot,   no,      nodeo,  nodedot,
       xpidot, z1,     z3,      z11,    z13,
       z21,    z23,    z31,     z33,    ecco,
       eccsq,  em,    argpm,  inclm, mm,
       nm,    nodem,
       irez,
       atime, d2201, d2211,  d3210, d3222,
       d4410, d4422, d5220,  d5232, d5421,
       d5433, dedt,  didt,   dmdt,
       dnodt, domdt, del1,   del2,  del3,
       xfact, xlamo, xli,    xni,
     ):

     q22    = 1.7891679e-6;
     q31    = 2.1460748e-6;
     q33    = 2.2123015e-7;
     root22 = 1.7891679e-6;
     root44 = 7.3636953e-9;
     root54 = 2.1765803e-9;
     rptim  = 4.37526908801129966e-3; # equates to 7.29211514668855e-5 rad/sec
     root32 = 3.7393792e-7;
     root52 = 1.1428639e-7;
     x2o3   = 2.0 / 3.0;
     znl    = 1.5835218e-4;
     zns    = 1.19459e-5;

     #  sgp4fix identify constants and allow alternate values
     xke = whichconst.xke

     #  -------------------- deep space initialization ------------
     irez = 0;
     if 0.0034906585 < nm < 0.0052359877:
         irez = 1;
     if 8.26e-3 <= nm <= 9.24e-3 and em >= 0.5:
         irez = 2;

     #  ------------------------ do solar terms -------------------
     ses  =  ss1 * zns * ss5;
     sis  =  ss2 * zns * (sz11 + sz13);
     sls  = -zns * ss3 * (sz1 + sz3 - 14.0 - 6.0 * emsq);
     sghs =  ss4 * zns * (sz31 + sz33 - 6.0);
     shs  = -zns * ss2 * (sz21 + sz23);
     #  sgp4fix for 180 deg incl
     if inclm < 5.2359877e-2 or inclm > pi - 5.2359877e-2:
       shs = 0.0;
     if sinim != 0.0:
       shs = shs / sinim;
     sgs  = sghs - cosim * shs;

     #  ------------------------- do lunar terms ------------------
     dedt = ses + s1 * znl * s5;
     didt = sis + s2 * znl * (z11 + z13);
     dmdt = sls - znl * s3 * (z1 + z3 - 14.0 - 6.0 * emsq);
     sghl = s4 * znl * (z31 + z33 - 6.0);
     shll = -znl * s2 * (z21 + z23);
     #  sgp4fix for 180 deg incl
     if inclm < 5.2359877e-2 or inclm > pi - 5.2359877e-2:
         shll = 0.0;
     domdt = sgs + sghl;
     dnodt = shs;
     if sinim != 0.0:

         domdt = domdt - cosim / sinim * shll;
         dnodt = dnodt + shll / sinim;


     #  ----------- calculate deep space resonance effects --------
     dndt   = 0.0;
     theta  = (gsto + tc * rptim) % twopi
     em     = em + dedt * t;
     inclm  = inclm + didt * t;
     argpm  = argpm + domdt * t;
     nodem  = nodem + dnodt * t;
     mm     = mm + dmdt * t;
     

     #  -------------- initialize the resonance terms -------------
     if irez != 0:

         aonv = pow(nm / xke, x2o3);

         #  ---------- geopotential resonance for 12 hour orbits ------
         if irez == 2:

             cosisq = cosim * cosim;
             emo    = em;
             em     = ecco;
             emsqo  = emsq;
             emsq   = eccsq;
             eoc    = em * emsq;
             g201   = -0.306 - (em - 0.64) * 0.440;

             if em <= 0.65:

                 g211 =    3.616  -  13.2470 * em +  16.2900 * emsq;
                 g310 =  -19.302  + 117.3900 * em - 228.4190 * emsq +  156.5910 * eoc;
                 g322 =  -18.9068 + 109.7927 * em - 214.6334 * emsq +  146.5816 * eoc;
                 g410 =  -41.122  + 242.6940 * em - 471.0940 * emsq +  313.9530 * eoc;
                 g422 = -146.407  + 841.8800 * em - 1629.014 * emsq + 1083.4350 * eoc;
                 g520 = -532.114  + 3017.977 * em - 5740.032 * emsq + 3708.2760 * eoc;

             else:

                 g211 =   -72.099 +   331.819 * em -   508.738 * emsq +   266.724 * eoc;
                 g310 =  -346.844 +  1582.851 * em -  2415.925 * emsq +  1246.113 * eoc;
                 g322 =  -342.585 +  1554.908 * em -  2366.899 * emsq +  1215.972 * eoc;
                 g410 = -1052.797 +  4758.686 * em -  7193.992 * emsq +  3651.957 * eoc;
                 g422 = -3581.690 + 16178.110 * em - 24462.770 * emsq + 12422.520 * eoc;
                 if em > 0.715:
                     g520 =-5149.66 + 29936.92 * em - 54087.36 * emsq + 31324.56 * eoc;
                 else:
                     g520 = 1464.74 -  4664.75 * em +  3763.64 * emsq;

             if em < 0.7:

                 g533 = -919.22770 + 4988.6100 * em - 9064.7700 * emsq + 5542.21  * eoc;
                 g521 = -822.71072 + 4568.6173 * em - 8491.4146 * emsq + 5337.524 * eoc;
                 g532 = -853.66600 + 4690.2500 * em - 8624.7700 * emsq + 5341.4  * eoc;

             else:

                 g533 =-37995.780 + 161616.52 * em - 229838.20 * emsq + 109377.94 * eoc;
                 g521 =-51752.104 + 218913.95 * em - 309468.16 * emsq + 146349.42 * eoc;
                 g532 =-40023.880 + 170470.89 * em - 242699.48 * emsq + 115605.82 * eoc;

             sini2=  sinim * sinim;
             f220 =  0.75 * (1.0 + 2.0 * cosim+cosisq);
             f221 =  1.5 * sini2;
             f321 =  1.875 * sinim  *  (1.0 - 2.0 * cosim - 3.0 * cosisq);
             f322 = -1.875 * sinim  *  (1.0 + 2.0 * cosim - 3.0 * cosisq);
             f441 = 35.0 * sini2 * f220;
             f442 = 39.3750 * sini2 * sini2;
             f522 =  9.84375 * sinim * (sini2 * (1.0 - 2.0 * cosim- 5.0 * cosisq) +
                     0.33333333 * (-2.0 + 4.0 * cosim + 6.0 * cosisq) );
             f523 = sinim * (4.92187512 * sini2 * (-2.0 - 4.0 * cosim +
                    10.0 * cosisq) + 6.56250012 * (1.0+2.0 * cosim - 3.0 * cosisq));
             f542 = 29.53125 * sinim * (2.0 - 8.0 * cosim+cosisq *
                    (-12.0 + 8.0 * cosim + 10.0 * cosisq));
             f543 = 29.53125 * sinim * (-2.0 - 8.0 * cosim+cosisq *
                    (12.0 + 8.0 * cosim - 10.0 * cosisq));
             xno2  =  nm * nm;
             ainv2 =  aonv * aonv;
             temp1 =  3.0 * xno2 * ainv2;
             temp  =  temp1 * root22;
             d2201 =  temp * f220 * g201;
             d2211 =  temp * f221 * g211;
             temp1 =  temp1 * aonv;
             temp  =  temp1 * root32;
             d3210 =  temp * f321 * g310;
             d3222 =  temp * f322 * g322;
             temp1 =  temp1 * aonv;
             temp  =  2.0 * temp1 * root44;
             d4410 =  temp * f441 * g410;
             d4422 =  temp * f442 * g422;
             temp1 =  temp1 * aonv;
             temp  =  temp1 * root52;
             d5220 =  temp * f522 * g520;
             d5232 =  temp * f523 * g532;
             temp  =  2.0 * temp1 * root54;
             d5421 =  temp * f542 * g521;
             d5433 =  temp * f543 * g533;
             xlamo =  (mo + nodeo + nodeo-theta - theta) % twopi
             xfact =  mdot + dmdt + 2.0 * (nodedot + dnodt - rptim) - no;
             em    = emo;
             emsq  = emsqo;

         #  ---------------- synchronous resonance terms --------------
         if irez == 1:

             g200  = 1.0 + emsq * (-2.5 + 0.8125 * emsq);
             g310  = 1.0 + 2.0 * emsq;
             g300  = 1.0 + emsq * (-6.0 + 6.60937 * emsq);
             f220  = 0.75 * (1.0 + cosim) * (1.0 + cosim);
             f311  = 0.9375 * sinim * sinim * (1.0 + 3.0 * cosim) - 0.75 * (1.0 + cosim);
             f330  = 1.0 + cosim;
             f330  = 1.875 * f330 * f330 * f330;
             del1  = 3.0 * nm * nm * aonv * aonv;
             del2  = 2.0 * del1 * f220 * g200 * q22;
             del3  = 3.0 * del1 * f330 * g300 * q33 * aonv;
             del1  = del1 * f311 * g310 * q31 * aonv;
             xlamo = (mo + nodeo + argpo - theta) % twopi
             xfact = mdot + xpidot - rptim + dmdt + domdt + dnodt - no;

         #  ------------ for sgp4, initialize the integrator ----------
         xli   = xlamo;
         xni   = no;
         atime = 0.0;
         nm    = no + dndt;

     return (
       em,    argpm,  inclm, mm,
       nm,    nodem,
       irez, atime,
       d2201, d2211,  d3210, d3222,
       d4410, d4422, d5220,  d5232,
       d5421, d5433, dedt,  didt,
       dmdt,  dndt, dnodt, domdt,
       del1,   del2,  del3, xfact,
       xlamo, xli,    xni,
       )
457325
def ccmod_xstep(k):
    

    YU = mp_D_Y - mp_D_U[k]
    b = mp_ZSf[k] + mp_drho * sl.rfftn(YU, None, mp_cri.axisN)
    Xf = sl.solvedbi_sm(mp_Zf[k], mp_drho, b, axis=mp_cri.axisM)
    mp_D_X[k] = sl.irfftn(Xf, mp_cri.Nv, mp_cri.axisN)
922644
def ac3(space):
    
    #determine arcs
    arcs = {}
    for name in space.variables:
        arcs[name] = set([])
    for const in space.constraints:
        for vname1,vname2 in product(const.vnames,const.vnames):
            if vname1 != vname2:
                #this is pessimistic, we assume that each constraint
                #pairwisely couples all variables it affects
                arcs[vname1].add(vname2)

    #enforce node consistency
    for vname in space.variables:
        for const in space.constraints:
            _unary(space,const,vname)

    #assemble work list
    worklist = set([])
    for v1 in space.variables:
        for v2 in space.variables:
            for const in space.constraints:
                if _binary(space,const,v1,v2):
                    for name in arcs[v1]:
                        worklist.add((v1,name))

    #work through work list
    while worklist:
        v1,v2 = worklist.pop()
        for const in space.constraints:
            if _binary(space,const,v1,v2):
                for vname in arcs[v1]:
                    worklist.add((v1,vname))
802461
def conv(attrs, inputs, proto_obj):
    """"""
    new_attrs = translation_utils._fix_attribute_names(attrs, {'kernel_shape' : 'kernel',
                                                               'strides' : 'stride',
                                                               'pads': 'pad',
                                                               'dilations': 'dilate',
                                                               'group': 'num_group'})
    new_attrs = translation_utils._add_extra_attributes(new_attrs, {'num_group' : 1})
    new_attrs = translation_utils._fix_bias('Convolution', new_attrs, len(inputs))

    new_attrs = translation_utils._fix_channels('Convolution', new_attrs, inputs, proto_obj)
    kernel = new_attrs['kernel']
    stride = new_attrs['stride'] if 'stride' in new_attrs else []
    padding = new_attrs['pad'] if 'pad' in new_attrs else []
    dilations = new_attrs['dilate'] if 'dilate' in new_attrs else []
    num_filter = new_attrs['num_filter']
    num_group = new_attrs['num_group']
    no_bias = new_attrs['no_bias'] if 'no_bias' in new_attrs else 0
    bias = None if no_bias is True else inputs[2]

    # Unlike ONNX, MXNet's convolution operator does not support asymmetric padding, so we first
    # use 'Pad' operator, which supports asymmetric padding. Then use the convolution operator.
    pad_width = (0, 0, 0, 0) + translation_utils._pad_sequence_fix(padding, kernel_dim=len(kernel))
    pad_op = symbol.pad(inputs[0], mode='constant', pad_width=pad_width)

    conv_op = symbol.Convolution(pad_op, inputs[1], bias,
                                 kernel=kernel, stride=stride, dilate=dilations,
                                 num_filter=num_filter, num_group=num_group, no_bias=no_bias)

    return conv_op, new_attrs, inputs
297166
def nmap(a, b, c, d, curvefn=None, normfn=None):
    
    if not curvefn:
        curvefn = lambda x: x

    def map(n):
        r = 1.0 * (n - a) / (b - a)
        out = curvefn(r) * (d - c) + c
        if not normfn:
            return out
        return normfn(out)
    return map
1293651
function nthRoot(node) {
  if (!Node.Type.isFunction(node, 'nthRoot')) {
    return Node.Status.noChange(node);
  }

  const radicandNode = getRadicandNode(node);
  if (Node.Type.isOperator(radicandNode)) {
    if (radicandNode.op === '^') {
      return nthRootExponent(node);
    }
    else if (radicandNode.op === '*') {
      return nthRootMultiplication(node);
    }
  }
  else if (Node.Type.isConstant(radicandNode)) {
    return nthRootConstant(node);
  }

  return Node.Status.noChange(node);
}
648614
def AcosTriAngles(v, f, normalize):
    

    if normalize:
        nm = lambda x : NormalizedNx3(x)
    else:
        nm = lambda x : x

    return Ch(lambda v :
        Sum3xN(NormalizedNx3(TriEdges(f, 1, 0, nm(v))) * NormalizedNx3(TriEdges(f, 2, 0, nm(v)))) &
        Sum3xN(NormalizedNx3(TriEdges(f, 2, 1, nm(v))) * NormalizedNx3(TriEdges(f, 0, 1, nm(v)))) &
        Sum3xN(NormalizedNx3(TriEdges(f, 0, 2, nm(v))) * NormalizedNx3(TriEdges(f, 1, 2, nm(v)))))
300467
def cas2mach(cas, h):
    
    tas = cas2tas(cas, h)
    M   = tas2mach(tas, h)
    return M
233897
def ecc(self, n):
        

        """
        ny, nx = self._profile.shape
        xmax, ymax = self._xymax
        xcm, ycm = self._cm

        # create (X, Y) grids relative to CM
        Y, X = np.mgrid[ymax:-ymax:1j*ny, -xmax:xmax:1j*nx]
        X -= xcm
        Y -= ycm

        # create grid of weights = profile * R^n
        Rsq = X*X + Y*Y
        if n == 1:
            W = np.sqrt(Rsq, out=Rsq)
        elif n == 2:
            W = Rsq
        else:
            if n & 1:  # odd n
                W = np.sqrt(Rsq)
            else:  # even n
                W = np.copy(Rsq)
            # multiply by R^2 until W = R^n
            for _ in range(int((n-1)/2)):
                W *= Rsq
        W *= self._profile

        # create grid of e^{i*n*phi} * W
        i_n_phi = np.zeros_like(X, dtype=complex)
        np.arctan2(Y, X, out=i_n_phi.imag)
        i_n_phi.imag *= n
        exp_phi = np.exp(i_n_phi, out=i_n_phi)
        exp_phi *= W

        return abs(exp_phi.sum()) / W.sum()
121462
def Axn(mt, x, n):
    
    return (mt.Mx[x] - mt.Mx[x + n]) / mt.Dx[x]
371693
private void _symmetricTql2 ()
  {
    // This is derived from the Algol procedures tql2, by
    // Bowdler, Martin, Reinsch, and Wilkinson, Handbook for
    // Auto. Comp., Vol.ii-Linear Algebra, and the corresponding
    // Fortran subroutine in EISPACK.
    for (int i = 1; i < m_nDim; i++)
      m_aEVe[i - 1] = m_aEVe[i];
    m_aEVe[m_nDim - 1] = 0.0;

    double f = 0.0;
    double tst1 = 0.0;
    for (int l = 0; l < m_nDim; l++)
    {
      // Find small subdiagonal element
      tst1 = Math.max (tst1, MathHelper.abs (m_aEVd[l]) + MathHelper.abs (m_aEVe[l]));
      int m = l;
      while (m < m_nDim)
      {
        if (MathHelper.abs (m_aEVe[m]) <= EPSILON * tst1)
          break;
        m++;
      }

      // If m == l, d[l] is an eigenvalue,
      // otherwise, iterate.

      if (m > l)
      {
        do
        {
          // Compute implicit shift
          double g = m_aEVd[l];
          double p = (m_aEVd[l + 1] - g) / (2.0 * m_aEVe[l]);
          double r = MathHelper.hypot (p, 1.0);
          if (p < 0)
          {
            r = -r;
          }
          m_aEVd[l] = m_aEVe[l] / (p + r);
          m_aEVd[l + 1] = m_aEVe[l] * (p + r);
          final double dl1 = m_aEVd[l + 1];
          double h = g - m_aEVd[l];
          for (int i = l + 2; i < m_nDim; i++)
          {
            m_aEVd[i] -= h;
          }
          f += h;

          // Implicit QL transformation.
          p = m_aEVd[m];
          double c = 1.0;
          double c2 = c;
          double c3 = c;
          final double el1 = m_aEVe[l + 1];
          double s = 0.0;
          double s2 = 0.0;
          for (int i = m - 1; i >= l; i--)
          {
            c3 = c2;
            c2 = c;
            s2 = s;
            g = c * m_aEVe[i];
            h = c * p;
            r = MathHelper.hypot (p, m_aEVe[i]);
            m_aEVe[i + 1] = s * r;
            s = m_aEVe[i] / r;
            c = p / r;
            p = c * m_aEVd[i] - s * g;
            m_aEVd[i + 1] = h + s * (c * g + s * m_aEVd[i]);

            // Accumulate transformation.

            for (int k = 0; k < m_nDim; k++)
            {
              h = m_aEigenVector[k][i + 1];
              m_aEigenVector[k][i + 1] = s * m_aEigenVector[k][i] + c * h;
              m_aEigenVector[k][i] = c * m_aEigenVector[k][i] - s * h;
            }
          }
          p = -s * s2 * c3 * el1 * m_aEVe[l] / dl1;
          m_aEVe[l] = s * p;
          m_aEVd[l] = c * p;

          // Check for convergence.

        } while (MathHelper.abs (m_aEVe[l]) > EPSILON * tst1);
      }
      m_aEVd[l] = m_aEVd[l] + f;
      m_aEVe[l] = 0.0;
    }

    // Sort eigenvalues and corresponding vectors.
    for (int i = 0; i < m_nDim - 1; i++)
    {
      int k = i;
      double p = m_aEVd[i];
      for (int j = i + 1; j < m_nDim; j++)
      {
        if (m_aEVd[j] < p)
        {
          k = j;
          p = m_aEVd[j];
        }
      }
      if (k != i)
      {
        m_aEVd[k] = m_aEVd[i];
        m_aEVd[i] = p;
        for (int j = 0; j < m_nDim; j++)
        {
          p = m_aEigenVector[j][i];
          m_aEigenVector[j][i] = m_aEigenVector[j][k];
          m_aEigenVector[j][k] = p;
        }
      }
    }
  }
929959
def cublasZsyrk(handle, uplo, trans, n, k, alpha, A, lda, beta, C, ldc):
    
    
    status = _libcublas.cublasZsyrk_v2(handle,
                                       _CUBLAS_FILL_MODE[uplo], 
                                       _CUBLAS_OP[trans], 
                                       n, k, ctypes.byref(cuda.cuDoubleComplex(alpha.real,    
                                                                               alpha.imag)),
                                       int(A), lda,
                                       ctypes.byref(cuda.cuDoubleComplex(beta.real,
                                                                         beta.imag)),
                                       int(C), ldc)
    cublasCheckStatus(status)
914207
def nalu(x, depth, epsilon=1e-30, name=None, reuse=None):
  """"""
  with tf.variable_scope(name, default_name="nalu", values=[x], reuse=reuse):
    x_shape = shape_list(x)
    x_flat = tf.reshape(x, [-1, x_shape[-1]])
    gw = tf.get_variable("w", [x_shape[-1], depth])
    g = tf.nn.sigmoid(tf.matmul(x_flat, gw))
    g = tf.reshape(g, x_shape[:-1] + [depth])
    a = nac(x, depth, name="nac_lin")
    log_x = tf.log(tf.abs(x) + epsilon)
    m = nac(log_x, depth, name="nac_log")
    return g * a + (1 - g) * tf.exp(m)
313524
def cublasDgemv(handle, trans, m, n, alpha, A, lda, x, incx, beta, y, incy):
    

    status = _libcublas.cublasDgemv_v2(handle,
                                       _CUBLAS_OP[trans], m, n,
                                       ctypes.byref(ctypes.c_double(alpha)),
                                       int(A), lda, int(x), incx,
                                       ctypes.byref(ctypes.c_double(beta)),
                                       int(y), incy)
    cublasCheckStatus(status)
805756
def _symm_current(C):
    """"""
    nans = np.isnan(C)
    C[nans] = np.einsum('klij', C)[nans]
    return C
1261167
def eigenvectors_from_displacements(disp,masses):
    
    nphonons,natoms,ndirections = disp.shape
    sqrt_masses = np.sqrt(masses)
    return np.einsum("nax,a->nax",disp,sqrt_masses)