## FORM LOSS AND JACOBIAN ***********************************************************************************
L0 = lambda xi, IC: r(z, xi, IC)[0, :] - IC['R0']
Ld0 = lambda xi, IC: IC['c'] * v(z, xi, IC)[0, :] - IC['V0']
Lf = lambda xi, IC: r(z, xi, IC)[-1, :]
Ldf = lambda xi, IC: IC['c'] * v(z, xi, IC)[-1, :]
Ls = lambda xi, IC: IC['c']**2 * a(z, xi, IC) - IC['ag'] + lam(z, xi)
# Htf = lambda xi,IC: np.dot(lam(z,xi)[-1,:],(-1./2.*lam(z,xi)[-1,:] + IC['ag']))
# Updated because need to at lam_r * v term for spectral method
Htf = lambda xi,IC: np.dot(lam(z,xi)[-1,:],(-1./2.*lam(z,xi)[-1,:] + IC['ag'])) \
+ np.dot(-IC['c'] *lamr(z,xi)[-1,:], IC['c'] * v(z,xi,IC)[-1,:]) + IC['Gam']
L = jit(lambda xi,IC: np.hstack(( Ls(xi,IC)[1:-1,:].flatten(), \
L0(xi,IC).flatten(), \
Ld0(xi,IC).flatten(), \
Lf(xi,IC).flatten(), \
Ldf(xi,IC).flatten() )) )
## INITIALIZE VARIABLES *************************************************************************************
xis = onp.zeros((Hs(z).shape[1], 3))
xic = onp.zeros((Hc(z).shape[1], 3))
if W == False:
b = np.sqrt(2) * onp.ones(1)
else:
b = np.sqrt(10) * onp.ones(1)
xi = TFCDictRobust({'xis':xis,\
'xic':xic})
def apply_psi(params, inputs, signal):
inputs = jnp.hstack((inputs, signal))
for fun, param in zip(apply_funs1, params):
inputs = fun(param, inputs)
return inputs
def apply_g(params, inputs, signal, **kwargs):
inputs = jnp.hstack((inputs, signal))
for fun, param in zip(apply_funs2, params):
inputs = fun(param, inputs)
return inputs
def compute_mle(self, y, compute_ci=True):
"""Compute maximum likelihood estimates.
Parameter
---------
y: jnp.array or dict, (n_samples)
Response. if dict is
"""
if self.compute_ci is None:
self.compute_ci = compute_ci
if type(y) is dict:
y_train = y['train']
if len(y['train']) == 0:
raise ValueError('Training set is empty after burned in.')
y_dev = y.get('dev', None)
else:
y = {'train': y}
y_train = y['train']
n_samples = len(y_train) - self.burn_in
X = jnp.hstack([
jnp.hstack([
jnp.ones(n_samples)[:, jnp.newaxis], self.XS['train'][name]
]) if name in self.S else jnp.hstack(
[jnp.ones(n_samples)[:, jnp.newaxis], self.X['train'][name]])
for name in self.filter_names
])
X = jnp.hstack([jnp.ones(n_samples)[:, jnp.newaxis], X])
XtX = X.T @ X
Xty = X.T @ y_train[self.burn_in:]
mle = jnp.linalg.lstsq(XtX, Xty, rcond=None)[0]
self.b['mle'] = {}
self.w['mle'] = {}
self.intercept['mle'] = {}
# slicing the mle matrix into each filter
l = jnp.cumsum(
jnp.hstack(
[0, [self.n_features[name] + 1 for name in self.n_features]]))
idx = [jnp.array((l[i], l[i + 1])) for i in range(len(l) - 1)]
self.idx = idx
for i, name in enumerate(self.filter_names):
mle_params = mle[idx[i][0]:idx[i][1]][:, jnp.newaxis].astype(
self.dtype)
self.intercept['mle'][name] = mle_params[0]
if name in self.S:
self.b['mle'][name] = mle_params[1:]
self.w['mle'][name] = self.S[name] @ self.b['mle'][name]
else:
self.w['mle'][name] = mle_params[1:]
self.intercept['mle']['global'] = mle[0]
self.p['mle'] = {}
self.y_pred['mle'] = {}
for name in self.filter_names:
if name in self.S:
self.p['mle'].update({name: self.b['mle'][name]})
else:
self.p['mle'].update({name: self.w['mle'][name]})
self.p['mle']['intercept'] = self.intercept['mle']
self.y['train'] = y_train[self.burn_in:]
if len(self.y['train']) == 0:
raise ValueError('Training set is empty after burned in.')
self.y_pred['mle']['train'] = self.forwardpass(self.p['mle'],
kind='train')
# # get filter confidence interval
if self.compute_ci:
self._get_filter_variance(w_type='mle')
self._get_response_variance(w_type='mle', kind='train')
if type(y) is dict and 'dev' in y:
self.y['dev'] = y_dev[self.burn_in:]
if len(self.y['dev']) == 0:
raise ValueError('Dev set is empty after burned in.')
self.y_pred['mle']['dev'] = self.forwardpass(self.p['mle'],
kind='dev')
if self.compute_ci:
self._get_response_variance(w_type='mle', kind='dev')
self.mle_computed = True
def aug_b(t):
bval = b(t)
return np.hstack((bval, bval))
## FORM LOSS AND JACOBIAN ***********************************************************************************
L0 = lambda xi, IC: r(z, xi, IC)[0, :] - IC['R0']
Ld0 = lambda xi, IC: xi['b']**2 * v(z, xi, IC)[0, :] - IC['V0']
Lf = lambda xi, IC: r(z, xi, IC)[-1, :]
Ldf = lambda xi, IC: xi['b']**2 * v(z, xi, IC)[-1, :]
Ls = lambda xi, IC: xi['b']**4 * a(z, xi, IC) - IC['ag'] + lam(z, xi)
# Htf = lambda xi,IC: np.dot(lam(z,xi)[-1,:],(-1./2.*lam(z,xi)[-1,:] + IC['ag']))
# Updated because need to at lam_r * v term for spectral method
Htf = lambda xi,IC: np.dot(lam(z,xi)[-1,:],(-1./2.*lam(z,xi)[-1,:] + IC['ag'])) \
+ np.dot(-xi['b']**2 *lamr(z,xi)[-1,:], xi['b']**2 * v(z,xi,IC)[-1,:]) + IC['Gam']
L = jit(lambda xi,IC: np.hstack([Ls(xi,IC)[1:-1,:].flatten(), \
L0(xi,IC).flatten(), \
Ld0(xi,IC).flatten(), \
Lf(xi,IC).flatten(), \
Ldf(xi,IC).flatten(), \
Htf(xi,IC)] ))
## INITIALIZE VARIABLES *************************************************************************************
xis = onp.zeros((Hs(z).shape[1], 3))
xic = onp.zeros((Hc(z).shape[1], 3))
if W == False:
b = np.sqrt(2) * onp.ones(1)
else:
b = np.sqrt(10) * onp.ones(1)
xi = TFCDictRobust({'xis':xis,\
'xic':xic,\
egrad(u, 2), 2)(xi, *x)
L = lambda xi, *x: laplace(xi, *x) - np.exp(-x[0]) * (x[0] - 2. + x[1]
**3 + 6. * x[1])
# Calculate the A and B matrices
zXi = np.zeros((tfc.basisClass.numBasisFunc))
A = np.vstack([
jacfwd(L, 0)(zXi, *x),
H(x[0][x0ind], x[1][x0ind]),
H(x[0][xfind], x[1][xfind]),
H(x[0][y0ind], x[1][y0ind]),
H(x[0][yfind], x[1][yfind])
])
B = np.hstack([
-L(zXi, *x), x[1][x0ind]**3, (1. + x[1][xfind]**3) * np.exp(-1.),
x[0][y0ind] * np.exp(-x[0][y0ind]),
(x[0][yfind] + 1.) * np.exp(-x[0][yfind])
])
# Calculate the xi values
xi = np.dot(np.linalg.pinv(A), B)
# Calculate the error
dark = np.meshgrid(np.linspace(x0[0], xf[0], n),
np.linspace(x0[1], xf[1], n))
x = (dark[0].flatten(), dark[1].flatten())
ur = real(*x)
ue = u(xi, *x)
err = ur - ue
testErr[j, k] = np.max(np.abs(err))
def test_fun5(p, t, w1, w2):
X = jnp.hstack([p, t])
z1 = nonlin(jnp.matmul(X, w1))
z2 = jnp.matmul(z1, w2)
return jnp.reshape(z2, ())
def model_bnn(p, t, w1, b1, w2, b2, w3, b3):
X = jnp.hstack([p, t]) # jnp.concatenate((p, t), axis=1)
z1 = nonlin(jnp.matmul(X, w1) + jnp.transpose(b1))
z2 = nonlin(jnp.matmul(z1, w2) + jnp.transpose(b2))
z3 = jnp.matmul(z2, w3) + jnp.transpose(b3)
return jnp.reshape(z3, ()) # z3.squeeze()
def jacres_u(U, Uold, peq, neq, sepq, accq, zccq):
Mp = peq.M
Mn = neq.M
Ms = sepq.M
Ma = accq.M
Mz = zccq.M
Np = peq.N
Nn = neq.N
cmat_pe, cmat_ne, \
uvec_pe, uvec_sep, uvec_ne,\
Tvec_acc, Tvec_pe, Tvec_sep, Tvec_ne, Tvec_zcc, \
phie_pe, phie_sep, phie_ne, phis_pe, phis_ne, \
j_pe,j_ne,eta_pe,eta_ne = unpack(U, Mp, Np, Mn, Nn, Ms, Ma, Mz)
cmat_old_pe, cmat_old_ne,\
uvec_old_pe, uvec_old_sep, uvec_old_ne,\
Tvec_old_acc, Tvec_old_pe, Tvec_old_sep, Tvec_old_ne, Tvec_old_zcc,\
_, _, \
_, _,\
_,_,\
_,_,_= unpack(Uold,Mp, Np, Mn, Nn, Ms, Ma, Mz)
bc_u0p = peq.bc_zero_neumann(uvec_pe[0], uvec_pe[1])
res_up = vmap(peq.electrolyte_conc)(uvec_pe[0:Mp], uvec_pe[1:Mp + 1],
uvec_pe[2:Mp + 2], Tvec_pe[0:Mp],
Tvec_pe[1:Mp + 1], Tvec_pe[2:Mp + 2],
j_pe[0:Mp], uvec_old_pe[1:Mp + 1])
bc_uMp = peq.bc_u_sep_p(uvec_pe[Mp], uvec_pe[Mp + 1], Tvec_pe[Mp],
Tvec_pe[Mp + 1], uvec_sep[0], uvec_sep[1],
Tvec_sep[0], Tvec_sep[1])
bc_u0s = peq.bc_inter_cont(uvec_pe[Mp], uvec_pe[Mp + 1], uvec_sep[0],
uvec_sep[1])
res_us = vmap(sepq.electrolyte_conc)(uvec_sep[0:Ms], uvec_sep[1:Ms + 1],
uvec_sep[2:Ms + 2], Tvec_sep[0:Ms],
Tvec_sep[1:Ms + 1],
Tvec_sep[2:Ms + 2],
uvec_old_sep[1:Ms + 1])
bc_uMs = sepq.bc_u_sep_n(uvec_ne[0], uvec_ne[1], Tvec_ne[0], Tvec_ne[1],
uvec_sep[Ms], uvec_sep[Ms + 1], Tvec_sep[Ms],
Tvec_sep[Ms + 1])
bc_u0n = neq.bc_inter_cont(uvec_ne[0], uvec_ne[1], uvec_sep[Ms],
uvec_sep[Ms + 1])
res_un = vmap(neq.electrolyte_conc)(uvec_ne[0:Mn], uvec_ne[1:Mn + 1],
uvec_ne[2:Mn + 2], Tvec_ne[0:Mn],
Tvec_ne[1:Mn + 1], Tvec_ne[2:Mn + 2],
j_ne[0:Mn], uvec_old_ne[1:Mn + 1])
bc_uMn = neq.bc_zero_neumann(uvec_ne[Mn], uvec_ne[Mn + 1])
""" positive electrode"""
arg_up = [
uvec_pe[0:Mp], uvec_pe[1:Mp + 1], uvec_pe[2:Mp + 2], Tvec_pe[0:Mp],
Tvec_pe[1:Mp + 1], Tvec_pe[2:Mp + 2], j_pe[0:Mp], uvec_old_pe[1:Mp + 1]
]
arg_uMp = [
uvec_pe[Mp], uvec_pe[Mp + 1], Tvec_pe[Mp], Tvec_pe[Mp + 1],
uvec_sep[0], uvec_sep[1], Tvec_sep[0], Tvec_sep[1]
]
A_up = vmap(grad(peq.electrolyte_conc, range(0, len(arg_up) - 1)))(*arg_up)
bc_u0p_grad = np.array([[-1, 1]]).T
bc_uMp_grad = ((jax.grad(peq.bc_u_sep_p, range(0,
len(arg_uMp)))))(*arg_uMp)
#uu
bc_uup = {"right": bc_u0p_grad[0:2], "left": bc_uMp_grad[0:2]}
J_uupp = build_tridiag(Mp, A_up[0:3], **bc_uup)
# interface boundar8
J_uups = coo_matrix(
(np.ravel(np.asarray(bc_uMp_grad[4:6])), ([Mp + 1, Mp + 1], [0, 1])),
shape=(Mp + 2, Ms + 2 + Mn + 2))
J_uup = hstack([J_uupp, J_uups])
# uT
bc_uTp = {"left": bc_uMp_grad[2:4]}
J_uTpp = build_tridiag(Mp, A_up[3:6], **bc_uTp)
# interface
J_uTps = coo_matrix(
(np.ravel(np.asarray(bc_uMp_grad[6:8])), ([Mp + 1, Mp + 1], [0, 1])),
shape=(Mp + 2, Ms + 2 + Mn + 2))
J_uTp = hstack([J_uTpp, J_uTps])
# uj
J_ujp = build_diag(Mp, A_up[6], "long")
""" Separator """
arg_u0s = [uvec_pe[Mp], uvec_pe[Mp + 1], uvec_sep[0], uvec_sep[1]]
arg_us = [
uvec_sep[0:Ms], uvec_sep[1:Ms + 1], uvec_sep[2:Ms + 2], Tvec_sep[0:Ms],
Tvec_sep[1:Ms + 1], Tvec_sep[2:Ms + 2], uvec_old_sep[1:Ms + 1]
]
arg_uMs = [
uvec_ne[0], uvec_ne[1], Tvec_ne[0], Tvec_ne[1], uvec_sep[Ms],
uvec_sep[Ms + 1], Tvec_sep[Ms], Tvec_sep[Ms + 1]
]
A_us = vmap(grad(sepq.electrolyte_conc, range(0,
len(arg_us) - 1)))(*arg_us)
bc_u0s_grad = (jax.grad(peq.bc_inter_cont, range(0,
len(arg_u0s))))(*arg_u0s)
bc_uMs_grad = (jax.grad(sepq.bc_u_sep_n, range(0, len(arg_uMs))))(*arg_uMs)
#uu
bc_uus = {"right": bc_u0s_grad[2:4], "left": bc_uMs_grad[4:6]}
J_uuss = build_tridiag(Ms, A_us[0:3], **bc_uus)
# positive sep interface
J_uusp = coo_matrix(
(np.ravel(np.asarray(bc_u0s_grad[0:2])), ([0, 0], [Mp, Mp + 1])),
shape=(Ms + 2, Mp + 2))
#negative sep interface
J_uusn = coo_matrix(
(np.ravel(np.asarray(bc_uMs_grad[0:2])), ([Ms + 1, Ms + 1], [0, 1])),
shape=(Ms + 2, Mn + 2))
J_uus = hstack([J_uusp, J_uuss, J_uusn])
# uT
bc_uTs = {"left": bc_uMs_grad[6:8]}
J_uTss = build_tridiag(Ms, A_us[3:6], **bc_uTs)
# J_uTsp = coo_matrix((np.ravel(np.asarray(bc_u0s_grad[2:4])),([0,0],[Mp,Mp+1] )),shape=(Ms+2,Mp+2))
J_uTsp = empty_rec(Ms + 2, Mp + 2)
J_uTsn = coo_matrix(
(np.ravel(np.asarray(bc_uMs_grad[2:4])), ([Ms + 1, Ms + 1], [0, 1])),
shape=(Ms + 2, Mn + 2))
J_uTs = hstack([J_uTsp, J_uTss, J_uTsn])
""" negative electrode"""
arg_un = [
uvec_ne[0:Mn], uvec_ne[1:Mn + 1], uvec_ne[2:Mn + 2], Tvec_ne[0:Mn],
Tvec_ne[1:Mn + 1], Tvec_ne[2:Mn + 2], j_ne[0:Mn], uvec_old_ne[1:Mn + 1]
]
arg_u0n = [uvec_ne[0], uvec_ne[1], uvec_sep[Ms], uvec_sep[Ms + 1]]
A_un = vmap(grad(neq.electrolyte_conc, range(0, len(arg_un) - 1)))(*arg_un)
bc_u0n_grad = grad(neq.bc_inter_cont, range(0, len(arg_u0n)))(*arg_u0n)
bc_uMn_grad = np.array([[-1, 1]]).T
#uu
bc_uun = {"right": bc_u0n_grad[0:2], "left": bc_uMn_grad[0:2]}
J_uunn = build_tridiag(Mn, A_un[0:3], **bc_uun)
J_uuns = coo_matrix((np.ravel(np.asarray(bc_u0n_grad[2:4])),
([0, 0], [Mp + 2 + Ms, Mp + 2 + Ms + 1])),
shape=(Mn + 2, Ms + 2 + Mp + 2))
J_uun = hstack([J_uuns, J_uunn])
# uT
J_uTnn = build_tridiag(Mn, A_un[3:6])
J_uTns = empty_rec(Mn + 2, Ms + 2 + Mp + 2)
J_uTn = hstack([J_uTns, J_uTnn])
# uj
J_ujn = build_diag(Mn, A_un[6], "long")
res_u = np.hstack((bc_u0p, res_up, bc_uMp, bc_u0s, res_us, bc_uMs, bc_u0n,
res_un, bc_uMn))
J_u = hstack([
empty_rec(Mp + 2 + Ms + 2 + Mn + 2,
Mp * (Np + 2) + Mn * (Nn + 2)), # c
vstack([J_uup, J_uus, J_uun]),
hstack([
empty_rec(Mp + 2 + Ms + 2 + Mn + 2, Ma + 2), # acc
vstack([J_uTp, J_uTs, J_uTn]),
empty_rec(Mp + 2 + Ms + 2 + Mn + 2, Mz + 2) # zcc
]),
empty_rec(Mp + 2 + Ms + 2 + Mn + 2, Mp + 2 + Ms + 2 + Mn + 2), #phie
empty_rec(Mp + 2 + Ms + 2 + Mn + 2, Mp + 2 + Mn + 2), #phis
vstack([
hstack([J_ujp, empty_rec(Mp + 2, Mn)]),
empty_rec(Ms + 2, Mp + Mn),
hstack([empty_rec(Mn + 2, Mp), J_ujn])
]),
empty_rec(Mp + 2 + Ms + 2 + Mn + 2, Mp + Mn)
])
return res_u, J_u
def _rmatmat(self, X):
"""Default implementation of _rmatmat defers to rmatvec or adjoint."""
if type(self)._adjoint == LinearOperator._adjoint:
return np.hstack([self.rmatvec(col.reshape(-1, 1)) for col in X.T])
else:
return self.H.matmat(X)
def __init__(self, n, nC, deg, dim=2, basis="CP", x0=None, xf=None):
# Store givens
self._elm_classes = [
"ELMSigmoid", "ELMTanh", "ELMSin", "ELMSwish", "ELMReLU"
]
self.deg = deg
self.dim = dim
# Set N based on user input
if isinstance(n, np.ndarray):
if not n.flatten().shape[0] == dim:
TFCPrint.Error(
"n has length " + str(n.flatten().shape[0]) +
", but it should be equal to the number of dimensions, " +
str(dim) + ".")
self.n = n.astype(np.int32)
else:
if not len(n) == dim:
TFCPrint.Error(
"n has length " + str(n) +
", but it should be equal to the number of dimensions, " +
str(dim) + ".")
self.n = np.array(n, dtype=np.int32)
self.N = int(np.prod(self.n, dtype=np.int32))
self.basis = basis
# Set x0 based on user input
if x0 is None:
self.x0 = np.zeros(dim)
else:
if isinstance(x0, np.ndarray):
if not x0.flatten().shape[0] == dim:
TFCPrint.Error(
"x0 has length " + str(x0.flatten().shape[0]) +
", but it should be equal to the number of dimensions, "
+ str(dim) + ".")
self.x0 = x0
else:
if not len(x0) == dim:
TFCPrint.Error(
"x0 has length " + len(x0) +
", but it should be equal to the number of dimensions, "
+ str(dim) + ".")
self.x0 = np.array(x0).flatten()
if not self.x0.shape[0] == dim:
TFCPrint.Error(
"x0 has length " + str(x0.flatten().shape[0]) +
", but it should be equal to the number of dimensions, "
+ str(dim) + ".")
# Set xf based on user input
if xf is None:
self.xf = np.zeros(dim)
else:
if isinstance(xf, np.ndarray):
if not xf.flatten().shape[0] == dim:
TFCPrint.Error(
"xf has length " + str(xf.flatten().shape[0]) +
", but it should be equal to the number of dimensions, "
+ str(dim) + ".")
self.xf = xf
else:
if not len(xf) == dim:
TFCPrint.Error(
"xf has length " + len(xf) +
", but it should be equal to the number of dimensions, "
+ str(dim) + ".")
self.xf = np.array(xf).flatten()
if not self.xf.shape[0] == dim:
TFCPrint.Error(
"xf has length " + str(xf.flatten().shape[0]) +
", but it should be equal to the number of dimensions, "
+ str(dim) + ".")
# Create nC matrix based on user input
if basis in self._elm_classes:
if isinstance(nC, int):
self.nC = onp.arange(nC, dtype=onp.int32)
elif isinstance(nC, np.ndarray):
self.nC = nC.astype(onp.int32)
elif isinstance(nC, list):
self.nC = np.array(nC, dtype=np.int32)
if self.nC.shape[0] > self.deg:
TFCPrint.Error(
"Number of basis functions is less than number of constraints!"
)
if np.any(self.nC < 0):
TFCPrint.Error(
"To set nC to -1 (no constraints) either use nC = -1 or nC = 0 (i.e., use an integer not a list or array). Do not put only -1 in a list or array, this will cause issues in the C++ layer."
)
else:
if isinstance(nC, np.ndarray) and len(nC.shape) > 1:
if not nC.shape[0] == self.dim:
TFCPrint.Error(
"nC has " + str(nC.flatten().shape[0]) +
" rows, but the row number should be equal to the number of dimensions, "
+ str(dim) + ".")
self.nC = nC.astype(np.int32)
else:
if isinstance(nC, np.ndarray):
nC = nC.tolist()
if not len(nC) == dim:
TFCPrint.Error(
"nC has length " + str(len(nC)) +
", but it should be equal to the number of dimensions, "
+ str(dim) + ".")
nCmax = 0
for k in range(dim):
if isinstance(nC[k], np.ndarray):
nCk = np.array(nC[k]).flatten()
else:
nCk = np.array([nC[k]]).flatten()
if nCk.shape[0] == 1:
maxk = nCk[0]
else:
maxk = nCk.shape[0]
if maxk > nCmax:
nCmax = maxk
if nCmax > self.deg:
TFCPrint.Error(
"Number of basis functions is less than the number of constraints!"
)
onC = onp.zeros((dim, nCmax))
for k in range(dim):
if isinstance(nC[k], np.ndarray):
nCk = np.array(nC[k]).flatten()
else:
nCk = onp.array([nC[k]]).flatten()
n = nCk.shape[0]
if n == 1:
nCk = onp.arange(nCk[0])
n = nCk.shape[0]
if n < nCmax:
if n == 0:
nCk = -1.0 * onp.ones(nCmax)
else:
nCk = np.hstack([nCk, -1 * np.ones(nCmax - n)])
onC[k, :] = nCk.astype(np.int32)
self.nC = np.array(onC.tolist(), dtype=np.int32)
# Setup the basis function
if self.basis == "CP":
from .utils.BF import nCP
self.basisClass = nCP(self.x0, self.xf, self.nC, self.deg + 1)
z0 = -1.0
zf = 1.0
elif self.basis == "LeP":
from .utils.BF import nLeP
self.basisClass = nLeP(self.x0, self.xf, self.nC, self.deg + 1)
z0 = -1.0
zf = 1.0
elif self.basis == "FS":
from .utils.BF import nFS
self.basisClass = nFS(self.x0, self.xf, self.nC, self.deg + 1)
z0 = -np.pi
zf = np.pi
elif self.basis == "ELMSigmoid":
from .utils.BF import nELMSigmoid
self.basisClass = nELMSigmoid(self.x0, self.xf, self.nC,
self.deg + 1)
z0 = 0.0
zf = 1.0
elif self.basis == "ELMTanh":
from .utils.BF import nELMTanh
self.basisClass = nELMTanh(self.x0, self.xf, self.nC, self.deg + 1)
z0 = 0.0
zf = 1.0
elif self.basis == "ELMSin":
from .utils.BF import nELMSin
self.basisClass = nELMSin(self.x0, self.xf, self.nC, self.deg + 1)
z0 = 0.0
zf = 1.0
elif self.basis == "ELMSwish":
from .utils.BF import nELMSwish
self.basisClass = nELMSwish(self.x0, self.xf, self.nC,
self.deg + 1)
z0 = 0.0
zf = 1.0
elif self.basis == "ELMReLU":
from .utils.BF import nELMReLU
self.basisClass = nELMReLU(self.x0, self.xf, self.nC, self.deg + 1)
z0 = 0.0
zf = 1.0
else:
TFCPrint.Error(
"Invalid basis selection. Please select a valid basis")
if self.basisClass.numBasisFunc > self.N:
TFCPrint.Warning(
"Warning, you have more basis functions than points!\nThis may lead to large solution errors!"
)
self.c = self.basisClass.c
# Calculate z points and corresponding x
self.z = onp.zeros((self.dim, self.N))
x = tuple([onp.zeros(self.N) for x in range(self.dim)])
if self.basis in ["CP", "LeP"]:
for k in range(self.dim):
nProd = onp.prod(self.n[k + 1:])
nStack = onp.prod(self.n[0:k])
n = self.n[k] - 1
I = onp.linspace(0, n, n + 1).reshape((n + 1, 1))
dark = onp.cos(np.pi * (n - I) / float(n))
dark = onp.hstack([dark] * nProd).flatten()
self.z[k, :] = onp.array([dark] * nStack).flatten()
x[k][:] = (self.z[k, :] - z0) / self.c[k] + self.x0[k]
else:
for k in range(self.dim):
nProd = onp.prod(self.n[k + 1:])
nStack = onp.prod(self.n[0:k])
dark = onp.linspace(z0, zf, num=self.n[k]).reshape(
(self.n[k], 1))
dark = onp.hstack([dark] * nProd).flatten()
self.z[k, :] = onp.array([dark] * nStack).flatten()
x[k][:] = (self.z[k, :] - z0) / self.c[k] + self.x0[k]
self.z = np.array(self.z.tolist())
self.x = tuple([np.array(x[k].tolist()) for k in range(self.dim)])
self.SetupJAX()
def __init__(self, placebo=0, autodiff=True, method="3-point"):
self.nsamples = 0
self.m, self.l, self.g, self.dt, self.H = 0.1, 0.2, 9.81, 0.05, 100
self.initial_state, self.goal_state, self.goal_action = (
jnp.array([1.0, 1.0, 0.0, 0.0, 0.0, 0.0]),
jnp.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
jnp.array([self.m * self.g / 2.0, self.m * self.g / 2.0]),
)
self.goal_action = jnp.hstack((self.goal_action, jnp.zeros(placebo)))
self.viewer = None
self.action_dim, self.state_dim = 2 + placebo, 6
self.last_u = jnp.zeros((2, ))
def f(x, u):
self.nsamples += 1
state = x
x, y, th, xdot, ydot, thdot = state
u1, u2 = u[:2]
m, g, l, dt = self.m, self.g, self.l, self.dt
xddot = -(u1 + u2) * jnp.sin(th) / m
yddot = (u1 + u2) * jnp.cos(th) / m - g
thddot = l * (u2 - u1) / (m * l**2)
state_dot = jnp.array([xdot, ydot, thdot, xddot, yddot, thddot])
new_state = state + state_dot * dt
return new_state
def c(x, u):
return 0.1 * (u - self.goal_action) @ (u - self.goal_action) + (
x - self.goal_state) @ (x - self.goal_state)
def f_x(x, u):
return approx_derivative(lambda x: f(x, u), x, method=method)
def f_u(x, u):
return approx_derivative(lambda u: f(x, u), u, method=method)
def c_x(x, u):
return 2 * (x - self.goal_state)
def c_u(x, u):
return 2 * 0.1 * (u - self.goal_action)
def c_xx(x, u):
return 2 * jnp.eye(self.state_dim)
def c_uu(x, u):
return 2 * 0.1 * jnp.eye(self.action_dim)
if autodiff:
self.f, self.f_x, self.f_u = (
f,
jax.jacfwd(f, argnums=0),
jax.jacfwd(f, argnums=1),
)
self.c, self.c_x, self.c_u, self.c_xx, self.c_uu = (
c,
jax.grad(c, argnums=0),
jax.grad(c, argnums=1),
jax.hessian(c, argnums=0),
jax.hessian(c, argnums=1),
)
else:
self.f, self.f_x, self.f_u = f, f_x, f_u
self.c, self.c_x, self.c_u, self.c_xx, self.c_uu = c, c_x, c_u, c_xx, c_uu
def _dminput(xr, xc):
if xc is None:
x = xr
else:
x = jnp.hstack((xr, xc))
return x
def jacres_phis(acc, pe, sep, ne, zcc):
Mp = peq.M
Mn = neq.M
Ms = sepq.M
Ma = accq.M
Mz = zccq.M
Np = peq.N
Nn = neq.N
# def solid_poten(self,phisn, phisc, phisp, j):
# hx = self.hx; a = self.a
# sigeff = self.sigma*(1-self.eps-self.epsf)
# ans = sigeff*( phisn - 2*phisc + phisp)/hx**2 - a*F*j
# return ans.reshape()
#
# def bc_phis(self,phis0, phis1, source):
# sigeff = self.sigma*(1-pe.eps-pe.epsf)
# bc = sigeff*( phis1 - phis0 )/pe.hx - source
# return bc.reshape()
""" Positive electrode"""
arg_phis0p = [pe.phisvec[0], pe.phisvec[1], -Iapp]
arg_phisp = [
pe.phisvec[0:Mp], pe.phisvec[1:Mp + 1], pe.phisvec[2:Mp + 2],
pe.jvec[0:Mp]
]
arg_phisMp = [pe.phisvec[Mp], pe.phisvec[Mp + 1]]
bc_phis0p = peq.bc_phis(*arg_phis0p)
res_phisp = vmap(peq.solid_poten)(*arg_phisp)
bc_phisMp = peq.bc_zero_neumann(*arg_phisMp)
bc_phis0p_grad = grad(peq.bc_phis, range(0, 2))(*arg_phis0p)
A_phisp = vmap(grad(peq.solid_poten, range(0, 4)))(*arg_phisp)
bc_phisMp_grad = grad(peq.bc_zero_neumann, range(0, 2))(*arg_phisMp)
bcphisp = {"right": bc_phis0p_grad, "left": bc_phisMp_grad}
J_phisphisp = build_tridiag(Mp, A_phisp[0:3], **bcphisp)
J_phisjp = build_diag(Mp, A_phisp[3], "long")
""" Negative electrode"""
arg_phis0n = [ne.phisvec[0], ne.phisvec[1]]
arg_phisn = [
ne.phisvec[0:Mn], ne.phisvec[1:Mn + 1], ne.phisvec[2:Mn + 2],
ne.jvec[0:Mn]
]
arg_phisMn = [ne.phisvec[Mn], ne.phisvec[Mn + 1], -Iapp]
bc_phis0n = neq.bc_zero_neumann(*arg_phis0n)
res_phisn = vmap(neq.solid_poten)(*arg_phisn)
bc_phisMn = neq.bc_phis(*arg_phisMn)
bc_phis0n_grad = grad(neq.bc_zero_neumann, range(0, 2))(*arg_phis0n)
A_phisn = vmap(grad(neq.solid_poten, range(0, 4)))(*arg_phisn)
bc_phisMn_grad = grad(neq.bc_phis, range(0, 2))(*arg_phisMn)
bcphisn = {"right": bc_phis0n_grad, "left": bc_phisMn_grad}
J_phisphisn = build_tridiag(Mn, A_phisn[0:3], **bcphisn)
J_phisjn = build_diag(Mn, A_phisn[3], "long")
J_phisphis = block_diag((J_phisphisp, J_phisphisn))
J_phisj = block_diag((J_phisjp, J_phisjn))
res_phis = np.hstack(
(bc_phis0p, res_phisp, bc_phisMp, bc_phis0n, res_phisn, bc_phisMn))
J_phis = hstack([
empty_rec(Mp + 2 + Mn + 2,
Mp * (Np + 2) + Mn * (Nn + 2)), #c
empty_rec(Mp + 2 + Mn + 2, Mp + 2 + Ms + 2 + Mn + 2), #u
empty_rec(Mp + 2 + Mn + 2,
Ma + 2 + Mp + 2 + Ms + 2 + Mn + 2 + Mz + 2), #T
empty_rec(Mp + 2 + Mn + 2, Mp + 2 + Ms + 2 + Mn + 2), #phie
J_phisphis,
J_phisj,
empty_rec(Mp + 2 + Mn + 2, Mp + Mn)
])
return res_phis, J_phis
Ntot_pe = (Np + 2) * (Mp) + 4 * (Mp + 2) + 2 * (Mp)
Ntot_ne = (Nn + 2) * (Mn) + 4 * (Mn + 2) + 2 * (Mn)
Ntot_sep = 3 * (Ms + 2)
Ntot_acc = Ma + 2
Ntot_zcc = Mz + 2
Ntot = Ntot_pe + Ntot_ne + Ntot_sep + Ntot_acc + Ntot_zcc
U = np.hstack([
peq.cavg * np.ones(Mp * (Np + 2)), neq.cavg * np.ones(Mn * (Nn + 2)),
1000 + np.zeros(Mp + 2), 1000 + np.zeros(Ms + 2), 1000 + np.zeros(Mn + 2),
np.zeros(Mp),
np.zeros(Mn),
np.zeros(Mp),
np.zeros(Mn),
np.zeros(Mp + 2) +
peq.open_circuit_poten(peq.cavg, peq.cavg, Tref, peq.cmax),
np.zeros(Mn + 2) +
neq.open_circuit_poten(neq.cavg, neq.cavg, Tref, neq.cmax),
np.zeros(Mp + 2) + 0,
np.zeros(Ms + 2) + 0,
np.zeros(Mn + 2) + 0, Tref + np.zeros(Ma + 2), Tref + np.zeros(Mp + 2),
Tref + np.zeros(Ms + 2), Tref + np.zeros(Mn + 2), Tref + np.zeros(Mz + 2)
])
#from https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html
def indirect(fn):
def jacfn(U, Uold):
y, vjp_fun = jax.vjp(fn, U, Uold)
J = jax.vmap(jax.jit(vjp_fun))(np.eye(len(U)))
return J
def image_grid(nrow, ncol, imagevecs, imshape):
"""Reshape a stack of image vectors into an image grid for plotting."""
images = iter(imagevecs.reshape((-1,) + imshape))
return np.vstack([np.hstack([next(images).T for _ in range(ncol)][::-1])
for _ in range(nrow)]).T
def test_Cdmo(plot=False):
def generate_donut(nmeans=10, nsamps_per_mean=50):
from scipy.stats import multivariate_normal
from numpy import exp
def pol2cart(theta, rho):
x = (rho * np.cos(theta)).reshape(-1, 1)
y = (rho * np.sin(theta)).reshape(-1, 1)
return np.concatenate([x, y], axis=1)
comp_distribution = multivariate_normal(np.zeros(2), np.eye(2) / 100)
means = pol2cart(np.linspace(0, 2 * 3.141, nmeans + 1)[:-1], 1)
rvs = comp_distribution.rvs(nmeans * nsamps_per_mean) + np.repeat(
means, nsamps_per_mean, 0)
true_dens = lambda samps: exp(
location_mixture_logpdf(samps, means,
np.ones(nmeans) / nmeans, comp_distribution
))
return rvs, means, true_dens
x1 = np.ones((1, 1))
x2 = np.zeros((1, 1))
(rvs, means, true_dens) = generate_donut(50, 10)
invec = FiniteVec(GaussianKernel(0.5), rvs[:, :1])
outvec = FiniteVec(GaussianKernel(0.5), rvs[:, 1:])
refervec = FiniteVec(outvec.k, np.linspace(-4, 4, 5000)[:, None])
cd = Cdo(invec, outvec, refervec, 0.1)
cm = Cmo(invec, outvec, 0.1)
(true_x1, est_x1, este_x1, true_x2, est_x2, este_x2) = [
lambda samps: true_dens(
np.hstack([np.repeat(x1, len(samps), 0), samps])),
lambda samps: np.squeeze(
inner(
multiply(cd, FiniteVec.construct_RKHS_Elem(invec.k, x1)).
normalized().unsigned_projection().normalized(),
FiniteVec(refervec.k, samps, prefactors=np.ones(len(samps))))),
lambda samps: np.squeeze(
inner(
multiply(cm, FiniteVec.construct_RKHS_Elem(invec.k, x1)).
normalized().unsigned_projection().normalized(),
FiniteVec(refervec.k, samps, prefactors=np.ones(len(samps))))),
lambda samps: true_dens(
np.hstack([np.repeat(x2, len(samps), 0), samps])),
lambda samps: np.squeeze(
inner(
multiply(cd, FiniteVec.construct_RKHS_Elem(invec.k, x2)).
normalized().unsigned_projection().normalized(),
FiniteVec(refervec.k, samps, prefactors=np.ones(len(samps))))),
lambda samps: np.squeeze(
inner(
multiply(cm, FiniteVec.construct_RKHS_Elem(invec.k, x2)).
normalized().unsigned_projection().normalized(),
FiniteVec(refervec.k, samps, prefactors=np.ones(len(samps)))))
]
t = np.array(
(true_x1(refervec.inspace_points), true_x2(refervec.inspace_points)))
e = np.array(
(est_x1(refervec.inspace_points), est_x2(refervec.inspace_points)))
if plot:
import pylab as pl
(fig, ax) = pl.subplots(1, 3, False, False)
ax[0].plot(refervec.inspace_points, t[0])
ax[0].plot(refervec.inspace_points, e[0], "--", label="dens")
ax[0].plot(refervec.inspace_points,
este_x1(refervec.inspace_points),
"-.",
label="emb")
ax[1].plot(refervec.inspace_points, t[1])
ax[1].plot(refervec.inspace_points, e[1], "--", label="dens")
ax[1].plot(refervec.inspace_points,
este_x2(refervec.inspace_points),
"-.",
label="emb")
ax[2].scatter(*rvs.T)
fig.legend()
fig.show()
assert (np.allclose(e, t, atol=0.5))
def p2d_fn(Np, Nn, Mp, Mn, Ms, Ma, Mz, fn, jac_fn):
peq, neq, sepq, accq, zccq = get_battery_sections(Np, Nn, Mp, Ms, Mn, Ma,
Mz)
# grid = pack_grid(Mp,Np,Mn,Nn,Ms,Ma,Mz)
U = np.hstack([
peq.cavg * np.ones(Mp * (Np + 2)), neq.cavg * np.ones(Mn * (Nn + 2)),
1000 + np.zeros(Mp + 2), 1000 + np.zeros(Ms + 2),
1000 + np.zeros(Mn + 2),
np.zeros(Mp),
np.zeros(Mn),
np.zeros(Mp),
np.zeros(Mn),
np.zeros(Mp + 2) +
peq.open_circuit_poten(peq.cavg, peq.cavg, Tref, peq.cmax),
np.zeros(Mn + 2) +
neq.open_circuit_poten(neq.cavg, neq.cavg, Tref, neq.cmax),
np.zeros(Mp + 2) + 0,
np.zeros(Ms + 2) + 0,
np.zeros(Mn + 2) + 0, Tref + np.zeros(Ma + 2), Tref + np.zeros(Mp + 2),
Tref + np.zeros(Ms + 2), Tref + np.zeros(Mn + 2),
Tref + np.zeros(Mz + 2)
])
Tf = 100
steps = Tf / delta_t
# steps = 2
# jac_fn = jax.jit(jacfwd(fn))
voltages = []
temps = []
start = timeit.default_timer()
# for i in range(0,int(steps)):
#
# [U, fail] = newton(fn, jac_fn, body_fun, U)
#
# cmat_pe, cmat_ne,uvec_pe, uvec_sep, uvec_ne, \
# Tvec_acc, Tvec_pe, Tvec_sep, Tvec_ne, Tvec_zcc, \
# phie_pe, phie_sep, phie_ne, phis_pe, phis_ne, jvec_pe, jvec_ne, eta_pe, eta_ne = unpack(U,Mp, Np, Mn, Nn, Ms, Ma, Mz)
# volt = phis_pe[1] - phis_ne[Mn]
# voltages.append(volt)
# temps.append(np.mean(Tvec_pe[1:Mp+1]))
# if (fail == 0):
# pass
# # print("timestep:", i)
# else:
# print('Premature end of run\n')
# break
# for i in range(0,int(steps)):
#
# [U, fail] = lax_newton(fn, jac_fn, U, maxit=5, tol=1e-8)
#
# cmat_pe, cmat_ne,uvec_pe, uvec_sep, uvec_ne, \
# Tvec_acc, Tvec_pe, Tvec_sep, Tvec_ne, Tvec_zcc, \
# phie_pe, phie_sep, phie_ne, phis_pe, phis_ne, jvec_pe, jvec_ne, eta_pe, eta_ne = unpack(U,Mp, Np, Mn, Nn, Ms, Ma, Mz)
# volt = phis_pe[1] - phis_ne[Mn]
# voltages.append(volt)
# temps.append(np.mean(Tvec_pe[1:Mp+1]))
# if (fail == 0):
# pass
# # print("timestep:", i)
# else:
# print('Premature end of run\n')
# break
print("entering newton")
[state, fail] = newton(fn, jac_fn, U)
# state = lax_newton(fn, jac_fn, U, maxit=5, tol=1e-7)
end = timeit.default_timer()
time = end - start
# return U, voltages, temps,time
return state, time
def test_Cdo_timeseries(plot=False):
if plot:
import pylab as pl
x = np.linspace(0, 40, 400).reshape((-1, 1))
y = np.sin(x) + randn(len(x)).reshape((-1, 1)) * 0.2
proc_data = np.hstack([x, y])
if plot:
pl.plot(x.flatten(), y.flatten())
invec = FiniteVec(GaussianKernel(0.5),
np.array([y.squeeze()[i:i + 10] for i in range(190)]))
outvec = FiniteVec(GaussianKernel(0.5), y[10:200])
refervec = FiniteVec(
outvec.k,
np.linspace(y[:-201].min() - 2, y[:-201].max() + 2, 5000)[:, None])
cd = Cdo(invec, outvec, refervec, 0.1)
cd = Cmo(invec, outvec, 0.1)
sol2 = np.array([
multiply(cd, FiniteVec(invec.k,
y[end - 10:end].T)).normalized().get_mean_var()
for end in range(200, 400)
])
if plot:
pl.plot(x[200:].flatten(), sol2.T[0].flatten())
invec = CombVec(
FiniteVec(PeriodicKernel(np.pi, 5), x[:200, :]),
SpVec(SplitDimsKernel([0, 1, 2],
[PeriodicKernel(np.pi, 5),
GaussianKernel(0.1)]),
proc_data[:200, :],
np.array([200]),
use_subtrajectories=True), np.multiply)
outvec = FiniteVec(GaussianKernel(0.5), y[1:-199])
#cd = Cdo(invec, outvec, refervec, 0.1)
cd = Cmo(invec, outvec, 0.1)
#sol = (cd.inp_feat.inner(SpVec(invec.k, proc_data[:230], np.array([230]), use_subtrajectories=True)))
#sol = [(cd.inp_feat.inner(SiEdSpVec(invec.k_obs, y[:end], np.array([end]), invec.k_idx, use_subtrajectories=False ))) for end in range(200,400) ]
#pl.plot(np.array([sol[i][-1] for i in range(len(sol))]))
#sol = np.array([multiply (cd, SpVec(invec.k, proc_data[:end], np.array([end]), use_subtrajectories=False)).normalized().get_mean_var() for end in range(200,400) ])
sol = multiply(
cd,
CombVec(
FiniteVec(invec.v1.k, x),
SpVec(invec.v2.k,
proc_data[:400],
np.array([400]),
use_subtrajectories=True),
np.multiply)).normalized().get_mean_var()
print(sol)
return sol2.T[0], sol.T[0][200:], y[200:]
(true_x1, est_x1, este_x1, true_x2, est_x2, este_x2) = [
lambda samps: true_dens(
np.hstack([np.repeat(x1, len(samps), 0), samps])),
lambda samps: np.squeeze(
inner(
multiply(cd, FiniteVec.construct_RKHS_Elem(invec.k, x1)).
normalized().unsigned_projection().normalized(),
FiniteVec(refervec.k, samps, prefactors=np.ones(len(samps))))),
lambda samps: np.squeeze(
inner(
multiply(cm, FiniteVec.construct_RKHS_Elem(invec.k, x1)).
normalized().unsigned_projection().normalized(),
FiniteVec(refervec.k, samps, prefactors=np.ones(len(samps))))),
lambda samps: true_dens(
np.hstack([np.repeat(x2, len(samps), 0), samps])),
lambda samps: np.squeeze(
inner(
multiply(cd, FiniteVec.construct_RKHS_Elem(invec.k, x2)).
normalized().unsigned_projection().normalized(),
FiniteVec(refervec.k, samps, prefactors=np.ones(len(samps))))),
lambda samps: np.squeeze(
inner(
multiply(cm, FiniteVec.construct_RKHS_Elem(invec.k, x2)).
normalized().unsigned_projection().normalized(),
FiniteVec(refervec.k, samps, prefactors=np.ones(len(samps)))))
]
t = np.array(
(true_x1(refervec.inspace_points), true_x2(refervec.inspace_points)))
e = np.array(
(est_x1(refervec.inspace_points), est_x2(refervec.inspace_points)))
if plot:
import pylab as pl
(fig, ax) = pl.subplots(1, 3, False, False)
ax[0].plot(refervec.inspace_points, t[0])
ax[0].plot(refervec.inspace_points, e[0], "--", label="dens")
ax[0].plot(refervec.inspace_points,
este_x1(refervec.inspace_points),
"-.",
label="emb")
ax[1].plot(refervec.inspace_points, t[1])
ax[1].plot(refervec.inspace_points, e[1], "--", label="dens")
ax[1].plot(refervec.inspace_points,
este_x2(refervec.inspace_points),
"-.",
label="emb")
ax[2].scatter(*rvs.T)
fig.legend()
fig.show()
assert (np.allclose(e, t, atol=0.5))
def get_min_accepted(self,
rng,
ϵ,
accepted,
n_simulations=1,
max_iterations=10,
smoothing=None,
verbose=True):
"""Iteratively run ABC until a minimum number of samples are accepted
Setting a maximum distance that simulation summaries can be allowed to
be from the summaries of each target an iterative scheme can be used to
get ``n_samples`` simulations at a time, compress them and calculate
the distances until the minimum number of desired samples are within
the allowed cutoff distance. If multiple targets are being inferred
then the simulations are run until all targets have at least the
minimum number of accepted samples.
Parameters
----------
rng : int(2,)
A random number generator to draw the parameters and make
simulations with
ϵ : float or float(n_targets)
The acceptance distance between summaries from simulations and the
summary of the target data. A different epsilon can be passed for
each target.
accepted : int
The minimum number of samples to be accepted within the `ϵ` cutoff
n_simulations : int, default=1
The number of simulations to do at once
max_iterations : int, default=10
Maximum number of iterations in the while loop to prevent infinite
runs. Note if max_iterations is reached then there are probably
not the required number of accepted samples
smoothing : float or None, default=None
A Gaussian smoothing for the marginal distributions
verbose : bool, default=True
Whether to print out the running stats (number accepted,
iterations, etc)
Returns
-------
float(any, n_params)
All parameters values drawn
float(any, n_summaries)
All summaries of simulations made
float(n_target, any)
The distance of every summary to each target
"""
@jax.jit
def loop(inputs):
return jax.lax.while_loop(loop_cond, loop_body, inputs)
def loop_cond(inputs):
return np.logical_and(np.less(np.min(inputs[-2]), accepted),
np.less(inputs[-1], max_iterations))
def loop_body(inputs):
rng, parameters, summaries, distances, n_accepted, iteration = \
inputs
rng, key = jax.random.split(rng)
parameter_samples = self.prior.sample(n_simulations, seed=key)
rng, key = jax.random.split(rng)
summary_samples = self.compressor(
self.simulator(key, parameter_samples))
distance_samples = jax.vmap(
lambda target, F: self.distance_measure(
summary_samples, target, F))(self.target_summaries, self.F)
indices = jax.lax.dynamic_slice(
np.arange(n_simulations * max_iterations),
[n_simulations * iteration], [n_simulations])
parameters = jax.ops.index_update(parameters,
jax.ops.index[indices],
parameter_samples)
summaries = jax.ops.index_update(summaries, jax.ops.index[indices],
summary_samples)
distances = jax.ops.index_update(distances, jax.ops.index[:,
indices],
distance_samples)
n_accepted = np.int32(np.less(distances, ϵ).sum(1))
return rng, parameters, summaries, distances, n_accepted, \
iteration + np.int32(1)
parameters = np.ones((max_iterations * n_simulations, self.n_params))
summaries = np.ones((max_iterations * n_simulations, self.n_summaries))
distances = np.ones(
(self.n_targets, max_iterations * n_simulations)) * np.inf
if self.parameters.all is not None:
parameters = np.vstack([parameters, self.parameters.all])
summaries = np.vstack([summaries, self.summaries.all])
distances = np.hstack([distances, self.distances.all])
if self.parameters.n_accepted is None:
self.set_accepted(ϵ, smoothing=smoothing)
n_accepted = np.int32(self.parameters.n_accepted)
else:
n_accepted = np.zeros(self.n_targets, dtype=np.int32)
current_accepted = n_accepted
iteration = 0
_, parameters, summaries, distances, n_accepted, iteration = loop(
(rng, parameters, summaries, distances, n_accepted, iteration))
keep = ~np.any(np.isinf(distances), 0)
parameters = parameters[keep]
summaries = summaries[keep]
distances = distances[:, keep]
if verbose:
print(f"{n_accepted - current_accepted} accepted in last ",
f"{iteration} iterations ",
f"({n_simulations * iteration} simulations done).")
return parameters, summaries, distances
def jacres_phie(acc, pe, sep, ne, zcc):
Mp = peq.M
Mn = neq.M
Ms = sepq.M
Ma = accq.M
Mz = zccq.M
Np = peq.N
Nn = neq.N
""" Positive Electrode"""
#electrolyte_poten(self,un, uc, up, phien, phiec, phiep, Tn, Tc, Tp,j):
arg_phie0p = [pe.phievec[0], pe.phievec[1]]
arg_phiep = [
pe.uvec[0:Mp], pe.uvec[1:Mp + 1], pe.uvec[2:Mp + 2], pe.phievec[0:Mp],
pe.phievec[1:Mp + 1], pe.phievec[2:Mp + 2], pe.Tvec[0:Mp],
pe.Tvec[1:Mp + 1], pe.Tvec[2:Mp + 2], pe.jvec[0:Mp]
]
#bc_phie_p(self,phie0_p, phie1_p, phie0_s, phie1_s, u0_p, u1_p, u0_s, u1_s, T0_p, T1_p, T0_s, T1_s)
arg_phieMp = [
pe.phievec[Mp], pe.phievec[Mp + 1], sep.phievec[0], sep.phievec[1],
pe.uvec[Mp], pe.uvec[Mp + 1], sep.uvec[0], sep.uvec[1], pe.Tvec[Mp],
pe.Tvec[Mp + 1], sep.Tvec[0], sep.Tvec[1]
]
bc_phie0p = peq.bc_zero_neumann(*arg_phie0p)
res_phiep = vmap(peq.electrolyte_poten)(*arg_phiep)
bc_phieMp = peq.bc_phie_p(*arg_phieMp)
bc_phie0p_grad = np.array([[-1, 1]]).T
A_phiep = vmap((grad(peq.electrolyte_poten,
range(0, len(arg_phiep)))))(*arg_phiep)
bc_phieMp_grad = grad(peq.bc_phie_p, range(0,
len(arg_phieMp)))(*arg_phieMp)
bcphiep = {"right": bc_phie0p_grad, "left": bc_phieMp_grad[0:2]}
J_phiephiepp = build_tridiag(Mp, A_phiep[3:6], **bcphiep)
J_phieps = coo_matrix((np.ravel(np.asarray(bc_phieMp_grad[2:4])),
([Mp + 1, Mp + 1], [0, 1])),
shape=(Mp + 2, Ms + 2 + Mn + 2))
J_phiephiep = hstack([J_phiephiepp, J_phieps])
bc_phieup = {"left": bc_phieMp_grad[4:6]}
J_phieupp = build_tridiag(Mp, A_phiep[0:3], **bc_phieup)
J_phieups = coo_matrix((np.ravel(np.asarray(bc_phieMp_grad[6:8])),
([Mp + 1, Mp + 1], [0, 1])),
shape=(Mp + 2, Ms + 2 + Mn + 2))
J_phieup = hstack([J_phieupp, J_phieups])
bc_phieTp = {"left": bc_phieMp_grad[8:10]}
J_phieTpp = build_tridiag(Mp, A_phiep[6:9], **bc_phieTp)
J_phieTps = coo_matrix((np.ravel(np.asarray(bc_phieMp_grad[10:12])),
([Mp + 1, Mp + 1], [0, 1])),
shape=(Mp + 2, Ms + 2 + Mn + 2))
J_phieTp = hstack([J_phieTpp, J_phieTps])
J_phiejp = build_diag(Mp, A_phiep[9], "long")
""" Separator """
arg_phie0s = [
sep.phievec[0], sep.phievec[1], pe.phievec[Mp], pe.phievec[Mp + 1]
]
arg_phies = [
sep.uvec[0:Ms], sep.uvec[1:Ms + 1], sep.uvec[2:Ms + 2],
sep.phievec[0:Ms], sep.phievec[1:Ms + 1], sep.phievec[2:Ms + 2],
sep.Tvec[0:Ms], sep.Tvec[1:Ms + 1], sep.Tvec[2:Ms + 2]
]
# bc_phie_sn(self,phie0_n, phie1_n, phie0_s, phie1_s, u0_n, u1_n, u0_s, u1_s, T0_n, T1_n, T0_s, T1_s):
arg_phieMs = [
ne.phievec[0], ne.phievec[1], sep.phievec[Ms], sep.phievec[Ms + 1],
ne.uvec[0], ne.uvec[1], sep.uvec[Ms], sep.uvec[Ms + 1], ne.Tvec[0],
ne.Tvec[1], sep.Tvec[Ms], sep.Tvec[Ms + 1]
]
bc_phie0s = peq.bc_inter_cont(*arg_phie0s)
res_phies = vmap(sepq.electrolyte_poten)(*arg_phies)
bc_phieMs = sepq.bc_phie_sn(*arg_phieMs)
bc_phie0s_grad = grad(peq.bc_inter_cont,
range(0, len(arg_phie0s)))(*arg_phie0s)
A_phies = vmap(grad(sepq.electrolyte_poten,
range(0, len(arg_phies))))(*arg_phies)
bc_phieMs_grad = grad(sepq.bc_phie_sn, range(0,
len(arg_phieMs)))(*arg_phieMs)
bcphies = {"right": bc_phie0s_grad[0:2], "left": bc_phieMs_grad[2:4]}
J_phiephiess = build_tridiag(Ms, A_phies[3:6], **bcphies)
J_phiephiesp = coo_matrix(
(np.ravel(np.asarray(bc_phie0s_grad[2:4])), ([0, 0], [Mp, Mp + 1])),
shape=(Ms + 2, Mp + 2))
J_phiephiesn = coo_matrix((np.ravel(np.asarray(bc_phieMs_grad[0:2])),
([Ms + 1, Ms + 1], [0, 1])),
shape=(Ms + 2, Mn + 2))
J_phiephies = hstack([J_phiephiesp, J_phiephiess, J_phiephiesn])
bcphieus = {"left": bc_phieMs_grad[6:8]}
J_phieuss = build_tridiag(Ms, A_phies[0:3], **bcphieus)
J_phieusp = empty_rec(Ms + 2, Mp + 2)
J_phieusn = coo_matrix((np.ravel(np.asarray(bc_phieMs_grad[4:6])),
([Ms + 1, Ms + 1], [0, 1])),
shape=(Ms + 2, Mn + 2))
J_phieus = hstack([J_phieusp, J_phieuss, J_phieusn])
bcphieTs = {"left": bc_phieMs_grad[10:12]}
J_phieTss = build_tridiag(Ms, A_phies[6:9], **bcphieTs)
J_phieTsp = empty_rec(Ms + 2, Mp + 2)
J_phieTsn = coo_matrix((np.ravel(np.asarray(bc_phieMs_grad[8:10])),
([Ms + 1, Ms + 1], [0, 1])),
shape=(Ms + 2, Mn + 2))
J_phieTs = hstack([J_phieTsp, J_phieTss, J_phieTsn])
""" Negative Electrode"""
arg_phie0n = [
ne.phievec[0], ne.phievec[1], sep.phievec[Ms], sep.phievec[Ms + 1]
]
arg_phien = [
ne.uvec[0:Mn], ne.uvec[1:Mn + 1], ne.uvec[2:Mn + 2], ne.phievec[0:Mn],
ne.phievec[1:Mn + 1], ne.phievec[2:Mn + 2], ne.Tvec[0:Mn],
ne.Tvec[1:Mn + 1], ne.Tvec[2:Mn + 2], ne.jvec[0:Mn]
]
arg_phieMn = [ne.phievec[Mn], ne.phievec[Mn + 1]]
bc_phie0n = neq.bc_inter_cont(*arg_phie0n)
res_phien = vmap(neq.electrolyte_poten)(*arg_phien)
bc_phieMn = neq.bc_zero_dirichlet(*arg_phieMn)
bc_phie0n_grad = grad(neq.bc_inter_cont,
range(0, len(arg_phie0n)))(*arg_phie0n)
A_phien = vmap(grad(neq.electrolyte_poten,
range(0, len(arg_phien))))(*arg_phien)
bc_phieMn_grad = np.array([[1 / 2, 1 / 2]]).T
bcphien = {"right": bc_phie0n_grad[0:2], "left": bc_phieMn_grad}
J_phiephienn = build_tridiag(Mn, A_phien[3:6], **bcphien)
J_phiephiens = coo_matrix((np.ravel(np.asarray(bc_phie0n_grad[2:4])),
([0, 0], [Mp + 2 + Ms, Mp + 2 + Ms + 1])),
shape=(Mn + 2, Mp + 2 + Ms + 2))
J_phiephien = hstack([J_phiephiens, J_phiephienn])
J_phieunn = build_tridiag(Mn, A_phien[0:3])
J_phieuns = empty_rec(Mn + 2, Mp + 2 + Ms + 2)
J_phieun = hstack([J_phieuns, J_phieunn])
J_phieTnn = build_tridiag(Mn, A_phien[6:9])
J_phieTns = empty_rec(Mn + 2, Mp + 2 + Ms + 2)
J_phieTn = hstack([J_phieTns, J_phieTnn])
J_phiejn = build_diag(Mn, A_phien[9], "long")
J_phieu = vstack((J_phieup, J_phieus, J_phieun))
J_phiephie = vstack((J_phiephiep, J_phiephies, J_phiephien))
J_phieT = hstack([
empty_rec(Mp + 2 + Mn + 2 + Ms + 2, Ma + 2),
vstack((J_phieTp, J_phieTs, J_phieTn)),
empty_rec(Mp + 2 + Mn + 2 + Ms + 2, Mz + 2)
])
J_phiej = vstack([
hstack([J_phiejp, empty_rec(Mp + 2, Mn)]),
empty_rec(Ms + 2, Mp + Mn),
hstack([empty_rec(Mn + 2, Mp), J_phiejn])
])
J_phiec = empty_rec(Mp + 2 + Ms + 2 + Mn + 2,
Mp * (Np + 2) + Mn * (Nn + 2))
res_phie = np.hstack(
(bc_phie0p, res_phiep, bc_phieMp, bc_phie0s, res_phies, bc_phieMs,
bc_phie0n, res_phien, bc_phieMn))
J_phie = hstack([
J_phiec, J_phieu, J_phieT, J_phiephie,
empty_rec(Mp + 2 + Mn + 2 + Ms + 2, Mp + 2 + Mn + 2), J_phiej,
empty_rec(Mp + 2 + Ms + 2 + Mn + 2, Mp + Mn)
])
return res_phie, J_phie
