def get_coefficients(self): ar, cr, a, b, c, d = self.term.get_coefficients() # Real componenets crd = cr * self.delta coeffs = [2 * ar * (np.cosh(crd) - 1) / crd**2, cr] # Imaginary coefficients cd = c * self.delta dd = d * self.delta c2 = c**2 d2 = d**2 factor = 2.0 / (self.delta * (c2 + d2))**2 cos_term = np.cosh(cd) * np.cos(dd) - 1 sin_term = np.sinh(cd) * np.sin(dd) C1 = a * (c2 - d2) + 2 * b * c * d C2 = b * (c2 - d2) - 2 * a * c * d coeffs += [ factor * (C1 * cos_term - C2 * sin_term), factor * (C2 * cos_term + C1 * sin_term), c, d, ] return coeffs
def get_value(self, tau0): dt = self.delta ar, cr, a, b, c, d = self.term.get_coefficients() # Format the lags correctly tau0 = np.abs(np.atleast_1d(tau0)) tau = tau0[..., None] # Precompute some factors dpt = dt + tau dmt = dt - tau # Real parts: # tau > Delta crd = cr * dt cosh = np.cosh(crd) norm = 2 * ar / crd**2 K_large = np.sum(norm * (cosh - 1) * np.exp(-cr * tau), axis=-1) # tau < Delta crdmt = cr * dmt K_small = K_large + np.sum(norm * (crdmt - np.sinh(crdmt)), axis=-1) # Complex part cd = c * dt dd = d * dt c2 = c**2 d2 = d**2 c2pd2 = c2 + d2 C1 = a * (c2 - d2) + 2 * b * c * d C2 = b * (c2 - d2) - 2 * a * c * d norm = 1.0 / (dt * c2pd2)**2 k0 = np.exp(-c * tau) cdt = np.cos(d * tau) sdt = np.sin(d * tau) # For tau > Delta cos_term = 2 * (np.cosh(cd) * np.cos(dd) - 1) sin_term = 2 * (np.sinh(cd) * np.sin(dd)) factor = k0 * norm K_large += np.sum((C1 * cos_term - C2 * sin_term) * factor * cdt, axis=-1) K_large += np.sum((C2 * cos_term + C1 * sin_term) * factor * sdt, axis=-1) # tau < Delta edmt = np.exp(-c * dmt) edpt = np.exp(-c * dpt) cos_term = (edmt * np.cos(d * dmt) + edpt * np.cos(d * dpt) - 2 * k0 * cdt) sin_term = (edmt * np.sin(d * dmt) + edpt * np.sin(d * dpt) - 2 * k0 * sdt) K_small += np.sum(2 * (a * c + b * d) * c2pd2 * dmt * norm, axis=-1) K_small += np.sum((C1 * cos_term + C2 * sin_term) * norm, axis=-1) mask = tau0 >= dt return K_large * mask + K_small * (~mask)
def doublesoliton(x, t, c, x0): # assert c[0] > c[1], "c1 has to be bigger than c[2]" xi0 = (jnp.sqrt(c[0]) / 2 * (x - c[0] * t - x0[0]) ) # switch to moving coordinate frame xi1 = jnp.sqrt(c[1]) / 2 * (x - c[1] * t - x0[1]) part_1 = 2 * (c[0] - c[1]) numerator = c[0] * jnp.cosh(xi1)**2 + c[1] * jnp.sinh(xi0)**2 denominator_1 = (jnp.sqrt(c[0]) - jnp.sqrt(c[1])) * jnp.cosh(xi0 + xi1) denominator_2 = (jnp.sqrt(c[0]) + jnp.sqrt(c[1])) * jnp.cosh(xi0 - xi1) u = part_1 * numerator / (denominator_1 + denominator_2)**2 return u
def OVMadjsys(veh, lead, p, leadlen, vehstar, lam, curfol, regime, havefol, vehlen, lamp1, folp, relax, folrelax, shiftloss, *args): # add in the stuff for eqn 22 to work # veh = current position and speed of simulated traj # lead = current position and speed of leader # p = parameters # regime 1 = base model # regime 0 = shifted end # all other regimes correspond to unique pieces of the model # note lead and args[2] true measurements need to both be given for the specific time at which we compute the deriv # this intakes lambda and returns its derivative. out = jnp.zeros((1, 2)) if regime == 1: # if we are in the model then we get the first indicator function in eqn 22 out[0, 0] = 2 * (veh[0] - vehstar[0]) + p[0] * p[1] * p[3] * lam[1] * ( 1 / jnp.cosh(p[2] + p[4] - p[1] * (lead[0] - leadlen - veh[0] + relax)))**2 out[0, 1] = -lam[0] + p[3] * lam[1] elif regime == 0: # otherwise we are in shifted end, we just get the contribution from loss function, which is the same on the entire interval out[0, 0] = shiftloss else: out[0, 0] = 2 * (veh[0] - vehstar[0]) out[0, 1] = -lam[0] if havefol == 1: # if we have a follower in the current platoon we get another term (the second indicator functino term in eqn 22) out[0, 0] += -folp[0] * folp[1] * folp[3] * lamp1[1] * ( 1 / jnp.cosh(folp[2] + folp[4] - folp[1] * (-vehlen + veh[0] - curfol[0] + folrelax)))**2 return out
def evolution_pepo_imag_time(g: float, dt: float, bc: str, dtype: np.dtype, lx: Optional[int] = None, ly: Optional[int] = None) -> Operator: # PEPO for U(dt) ~ U_vert(dt/2) U_bond(dt) U_vert(dt/2) # # half bond operators: # # | | | | # U_bond = A -- A # | | | | # # expm(- H_bond dt) = expm(- (-XX) dt) = expm(dt XX) = cosh(dt) + sinh(dt) XX = A_0 A_0 + A_1 A_1 # with A_0 = (cosh(dt) ** 0.5) * 1 , A_1 = (sinh(dt) ** 0.5) * X # A & B legs: (p,p*,k) A = np.empty([2, 2, 2], dtype=dtype) A = index_update(A, index[:, :, 0], (np.cosh(dt)**0.5) * s0) A = index_update(A, index[:, :, 1], (np.sinh(dt)**0.5) * sx) # expm(- H_vert dt/2) = expm(- (-gZ) dt/2) = expm(g dt/2 Z) u_vert = np.asarray(expm(g * (dt / 2) * sz), dtype=dtype) return _build_evolution_pepo(u_vert, A, bc, lx, ly)
def get_celerite_matrices(self, x, diag, *, a=None, U=None, V=None, P=None): dt = self.delta ar, cr, a, b, c, d = self.term.get_coefficients() # Real part cd = cr * dt delta_diag = 2 * np.sum(ar * (cd - np.sinh(cd)) / cd**2) # Complex part cd = c * dt dd = d * dt c2 = c**2 d2 = d**2 c2pd2 = c2 + d2 C1 = a * (c2 - d2) + 2 * b * c * d C2 = b * (c2 - d2) - 2 * a * c * d norm = (dt * c2pd2)**2 sinh = np.sinh(cd) cosh = np.cosh(cd) delta_diag += 2 * np.sum( (C2 * cosh * np.sin(dd) - C1 * sinh * np.cos(dd) + (a * c + b * d) * dt * c2pd2) / norm) new_diag = diag + delta_diag return super().get_celerite_matrices(x, new_diag, a=a, U=U, V=V, P=P)
def logZ4(L, T): r = 2 * np.arange(0, L) cr = c(r, L, T) ch = jnp.cosh(2. / T)**2 - 1. / jnp.tanh(2. / T)**2 + 0.j return jnp.sum( jnp.log(jax.vmap(chebyUn, in_axes=(None, 0))(L // 2 - 1, cr))) + jnp.sum( jnp.log(cr[1:L // 2]**2 - 1)) + jnp.log(ch)
def log_pdf_hypsec(s): ''' Log-pdf for a **Hypersecant** distribution The term -np.log(2.0) Is just a normalization ''' return jnp.sum(-2 * jnp.log(jnp.cosh(s)) - jnp.log(2.0))
def log_abs_det_jacobian(self, x, y): # NB: because domain and codomain are two spaces with different dimensions, determinant of # Jacobian is not well-defined. Here we return `log_abs_det_jacobian` of `x` and the # flatten lower triangular part of `y`. # stick_breaking_logdet = log(y / r) = log(z_cumprod) (modulo right shifted) z1m_cumprod = 1 - cumsum(y * y) # by taking diagonal=-2, we don't need to shift z_cumprod to the right # NB: diagonal=-2 works fine for (2 x 2) matrix, where we get an empty array z1m_cumprod_tril = matrix_to_tril_vec(z1m_cumprod, diagonal=-2) stick_breaking_logdet = 0.5 * np.sum(np.log(z1m_cumprod_tril), axis=-1) tanh_logdet = -2 * np.sum(np.log(np.cosh(x)), axis=-1) return stick_breaking_logdet + tanh_logdet
def apply(self, s, numHidden=2, bias=False): """Restricted Boltzmann machine with complex parameters. Args: * ``s``: Computational basis configuration. * ``numHidden``: Number of hidden units. * ``bias``: ``Boolean`` indicating whether to use bias. """ layer = nn.Dense.shared(features=numHidden, name='rbm_layer', bias=bias, dtype=global_defs.tCpx, kernel_init=cplx_init, bias_init=partial(jax.nn.initializers.zeros, dtype=global_defs.tCpx)) return jnp.sum(jnp.log(jnp.cosh(layer(2 * s - 1))))
def c(l, L, T): return jnp.cosh(2. / T) / jnp.tanh(2. / T) - jnp.cos((jnp.pi * l) / L)
def cosh(x): if isinstance(x, JaxArray): x = x.value return JaxArray(jnp.cosh(x))
def d_tanh(x, scale=1.0): return jnp.diag(scale / (jnp.cosh(x)**2))
# extract result print('sizeA = ', Atemp.shape, ', RGstep = ', k + 1) lnz += np.log(ncon([Atemp], [[1, 2, 1, 2]])) / 2**(RGstep + 1) return lnz print('Calculate lnz') lnz = lnz_trg(betaval, chiM, log2L) Tval = 1 / betaval FreeEnergy = -Tval * lnz ##### Compare with exact results (thermodynamic limit) maglambda = 1 / (np.sinh(2 / Tval)**2) N = 1000000 x = np.linspace(0, np.pi, N + 1) y = np.log( np.cosh(2 * betaval) * np.cosh(2 * betaval) + (1 / maglambda) * np.sqrt(1 + maglambda**2 - 2 * maglambda * np.cos(2 * x))) FreeExact = -Tval * ((np.log(2) / 2) + 0.25 * sum(y[1:(N + 1)] + y[:N]) / N) RelFreeErr = abs((FreeEnergy - FreeExact) / FreeExact) print('FreeEnergy= %.5f' % FreeEnergy) print('Exact = ', FreeExact) print('RelFreeErr = ', RelFreeErr) print('Calculate dlnz') dlnz = grad(lnz_trg, argnums=0)(betaval, chiM, log2L) #dlnz = grad(lnz_trg)(betaval) print('dlnz: %.5f' % dlnz) #from jax import value_and_grad #lnz, dlnz = value_and_grad(lnz_trg, argnums=0)(betaval, chiM, log2L) #dlnZ2 = grad(grad(lnz_trg))(betaval) #print(dlnZ2)
def __call__(self, x): x = nn.Dense(features=alpha * x.shape[-1], dtype=jax.numpy.float32)(x) x = jnp.log(jnp.cosh(x)) return jnp.sum(x, axis=-1)
# Constrained expressions: myTfc = utfc(n,nC,m,x0=x0,xf=xf) t = myTfc.x H = myTfc.H g = lambda xi,t: np.dot(H(t),xi['xi']) xslow = lambda xi,t: g(xi,t)\ +(3.-t)/3.*(aEll*np.sin(xi['phi'])*np.cos(xi['th'])-g(xi,np.array([0.])))\ +t/3.*(3.+aHyp*np.sinh(xi['v'])*np.cos(xi['psi'])-g(xi,np.array([3.]))) yslow = lambda xi,t: g(xi,t)\ +(3.-t)/3.*(bEll*np.sin(xi['phi'])*np.sin(xi['th'])-g(xi,np.array([0.])))\ +t/3.*(bHyp*np.sinh(xi['v'])*np.sin(xi['psi'])-g(xi,np.array([3.]))) zslow = lambda xi,t: g(xi,t)\ +(3.-t)/3.*(cEll*np.cos(xi['phi'])-g(xi,np.array([0.])))\ +t/3.*((-1.)**step(xi['n'])*cHyp*np.cosh(xi['v'])-g(xi,np.array([3.]))) x = jit(xslow); y = jit(yslow); z = jit(zslow) # Plot conic sections th = np.linspace(0.,2.*np.pi,100) v = np.linspace(0.,3.0,100) matHyp = np.meshgrid(v,th) phi = np.linspace(0.,np.pi,100) matEll = np.meshgrid(phi,th) xHyp = lambda n,v,th: aHyp*np.sinh(np.abs(v))*np.cos(th)+3. yHyp = lambda n,v,th: bHyp*np.sinh(np.abs(v))*np.sin(th) zHyp = lambda n,v,th: (-1.)**step(n)*cHyp*np.cosh(np.abs(v))
def OVMadj(veh, lead, p, leadlen, lam, reg, relax, relaxadj, use_r, *args): # args[0] controls which column of lead contains position and speed info # args[1] has the length of the lead vehicle # args[2] holds the true measurements # args[3] holds lambda # this is what we integrate to compute the gradient of objective function after solving the adjoint system if use_r: # relaxation out = jnp.zeros((1, 6)) if reg == 1: s = lead[0] - leadlen - veh[0] + relax out[0, 0] = -p[3] * lam[1] * (jnp.tanh(p[2]) - jnp.tanh(p[2] + p[4] - p[1] * s)) out[0, 1] = p[0] * p[3] * lam[1] * (-s) * ( 1 / jnp.cosh(p[2] + p[4] - p[1] * s)**2) out[0, 2] = -p[0] * p[3] * lam[1] * ( 1 / jnp.cosh(p[2])**2 - 1 / jnp.cosh(p[2] + p[4] - p[1] * s)**2) out[0, 3] = -lam[1] * ( -veh[1] + p[0] * (jnp.tanh(p[2]) - jnp.tanh(p[2] + p[4] - p[1] * s))) out[0, 4] = p[0] * p[3] * lam[1] * ( 1 / jnp.cosh(p[2] + p[4] - p[1] * s)**2) out[0, 5] = -relaxadj * p[0] * p[1] * p[3] * lam[1] * ( 1 / jnp.cosh(p[2] + p[4] - p[1] * (s)))**2 # contribution from relaxation phenomenon else: out = jnp.zeros((1, 5)) if reg == 1: s = lead[0] - leadlen - veh[0] out[0, 0] = -p[3] * lam[1] * (jnp.tanh(p[2]) - jnp.tanh(p[2] + p[4] - p[1] * s)) out[0, 1] = p[0] * p[3] * lam[1] * (-s) * ( 1 / jnp.cosh(p[2] + p[4] - p[1] * s)**2) out[0, 2] = -p[0] * p[3] * lam[1] * ( 1 / jnp.cosh(p[2])**2 - 1 / jnp.cosh(p[2] + p[4] - p[1] * s)**2) out[0, 3] = -lam[1] * ( -veh[1] + p[0] * (jnp.tanh(p[2]) - jnp.tanh(p[2] + p[4] - p[1] * s))) out[0, 4] = p[0] * p[3] * lam[1] * ( 1 / jnp.cosh(p[2] + p[4] - p[1] * s)**2) return out
def logCosh(x): return jnp.log(jnp.cosh(x))
def CpIG(self, T): T = jnp.squeeze(T) return (self.CpIGA + self.CpIGB * (self.CpIGC / T / jnp.sinh(self.CpIGC / T))**2 + self.CpIGD * (self.CpIGE / T / jnp.cosh(self.CpIGE / T))**2)
tf.math.betainc, lambda a, b, x, name=None: scipy_special.betainc(a, b, x)) bincount = utils.copy_docstring(tf.math.bincount, _bincount) ceil = utils.copy_docstring(tf.math.ceil, lambda x, name=None: np.ceil(x)) # confusion_matrix = utils.copy_docstring( # tf.math.confusion_matrix, # lambda labels, predictions, num_classes=None, weights=None, # dtype=tf.int32, name=None: ...) conj = utils.copy_docstring(tf.math.conj, lambda x, name=None: np.conj(x)) cos = utils.copy_docstring(tf.math.cos, lambda x, name=None: np.cos(x)) cosh = utils.copy_docstring(tf.math.cosh, lambda x, name=None: np.cosh(x)) count_nonzero = utils.copy_docstring( tf.math.count_nonzero, lambda input, axis=None, keepdims=None, dtype=tf.int64, name=None: ( # pylint: disable=g-long-lambda utils.numpy_dtype(dtype)(np.count_nonzero(input, axis)))) cumprod = utils.copy_docstring(tf.math.cumprod, _cumprod) cumsum = utils.copy_docstring(tf.math.cumsum, _cumsum) digamma = utils.copy_docstring(tf.math.digamma, lambda x, name=None: scipy_special.digamma(x)) divide = utils.copy_docstring(tf.math.divide, lambda x, y, name=None: np.divide(x, y))