def jac(v):
alf = self.alpha[0, 1] / self.alpha[0, 0]
gdc = self._g_idx
A, R = unvec(v)
x_mat = np.array(np.bmat([[2 * alf * A, -R.T], [R,
zeros((k, k))]]))
exh = expm(x_mat)
ex = expm((1 - 2 * alf) * A)
blk = np.zeros_like(exh)
blk[:d, :] = (ex * lbd[None, :]) @ ex.T @ exh[:, :d].T
blkA = (lbd[:, None] * ex.T) @ exh[:, :d].T @ X2 @ exh[:, :d]
fexh = 2 * expm_frechet(x_mat, blk @ X2)[1]
fex = 2 * expm_frechet((1 - 2 * alf) * A, blkA)[1]
for r in range(1, p + 1):
if r not in gdc:
continue
br, er = gdc[r]
fexh[br:br, br:br] = 0
fex[br:br, br:br] = 0
return vec(
(1 - 2 * alf) * asym(fex) + 2 * alf * asym(fexh[:d, :d]),
fexh[d:, :d] - fexh[:d, d:].T)
def _compute_prop_grad(self, k, j, compute_prop=True):
"""
Calculate the gradient of propagator wrt the control amplitude
in the timeslot using the expm_frechet method
The propagtor is calculated (almost) for 'free' in this method
and hence it is returned if compute_prop==True
Returns:
[prop], prop_grad
"""
dyn = self.parent
if dyn.oper_dtype == Qobj:
A = dyn._get_phased_dyn_gen(k).full() * dyn.tau[k]
E = dyn._get_phased_ctrl_dyn_gen(k, j).full() * dyn.tau[k]
if compute_prop:
prop_dense, prop_grad_dense = la.expm_frechet(A, E)
prop = Qobj(prop_dense, dims=dyn.dyn_dims)
prop_grad = Qobj(prop_grad_dense, dims=dyn.dyn_dims)
else:
prop_grad_dense = la.expm_frechet(A, E, compute_expm=False)
prop_grad = Qobj(prop_grad_dense, dims=dyn.dyn_dims)
else:
A = dyn._get_phased_dyn_gen(k) * dyn.tau[k]
E = dyn._get_phased_ctrl_dyn_gen(k, j) * dyn.tau[k]
if compute_prop:
prop, prop_grad = la.expm_frechet(A, E)
else:
prop_grad = la.expm_frechet(A, E, compute_expm=False)
if compute_prop:
return prop, prop_grad
else:
return prop_grad
def em_objective_for_aitken(
T, node_to_idx, site_weights,
m,
transq_unscaled, transp,
interact_trans, interact_dwell,
data,
root_distn1d,
trans_out, dwell_out,
scale,
):
"""
Recast EM as a fixed-point problem.
This approach is inspired by the introduction of the following paper.
A QUASI-NEWTON ACCELERATION OF THE EM ALGORITHM
Kenneth Lange
1995
"""
# Unpack some stuff.
nsites, nnodes, nstates = data.shape
n = nstates
# Scale the rate matrices according to the edge ratios.
transq = transq_unscaled * scale[:, None, None]
# Compute the probability transition matrix arrays
# and the interaction matrix arrays.
trans_indicator = np.ones((n, n)) - np.identity(n)
dwell_indicator = np.identity(n)
for edge in T.edges():
na, nb = edge
eidx = node_to_idx[nb] - 1
Q = transq[eidx]
transp[eidx] = expm(Q)
interact_trans[eidx] = expm_frechet(
Q, Q * trans_indicator, compute_expm=False)
interact_dwell[eidx] = expm_frechet(
Q, Q * dwell_indicator, compute_expm=False)
# Compute the expectations.
validation = 0
expectation_step(
m.indices, m.indptr,
transp, transq,
interact_trans, interact_dwell,
data,
root_distn1d,
trans_out, dwell_out,
validation,
)
# Compute the per-edge ratios.
trans_sum = (trans_out * site_weights[:, None]).sum(axis=0)
dwell_sum = (dwell_out * site_weights[:, None]).sum(axis=0)
scaling_ratios = trans_sum / -dwell_sum
# Return the new scaling factors.
return scale * scaling_ratios
def test_problematic_matrix(self):
# this test case uncovered a bug which has since been fixed
A = np.array([[1.50591997, 1.93537998], [0.41203263, 0.23443516]], dtype=float)
E = np.array([[1.87864034, 2.07055038], [1.34102727, 0.67341123]], dtype=float)
A_norm_1 = scipy.linalg.norm(A, 1)
sps_expm, sps_frechet = expm_frechet(A, E, method="SPS")
blockEnlarge_expm, blockEnlarge_frechet = expm_frechet(A, E, method="blockEnlarge")
assert_allclose(sps_expm, blockEnlarge_expm)
assert_allclose(sps_frechet, blockEnlarge_frechet)
def test_medium_matrix(self):
# profile this to see the speed difference
n = 1000
A = np.random.exponential(size=(n, n))
E = np.random.exponential(size=(n, n))
sps_expm, sps_frechet = expm_frechet(A, E, method="SPS")
blockEnlarge_expm, blockEnlarge_frechet = expm_frechet(A, E, method="blockEnlarge")
assert_allclose(sps_expm, blockEnlarge_expm)
assert_allclose(sps_frechet, blockEnlarge_frechet)
def test_medium_matrix(self):
# profile this to see the speed difference
n = 1000
A = np.random.exponential(size=(n, n))
E = np.random.exponential(size=(n, n))
sps_expm, sps_frechet = expm_frechet(A, E, method='SPS')
blockEnlarge_expm, blockEnlarge_frechet = expm_frechet(
A, E, method='blockEnlarge')
assert_allclose(sps_expm, blockEnlarge_expm)
assert_allclose(sps_frechet, blockEnlarge_frechet)
def get_unitary_derivative(self, angles, term_index=0):
"""
Compute the derivative of the block's unitary with respect to its
free parameter, assuming it is of the form :math:`e^{-i \\theta H}`
for a free parameter theta. If the block's operator is a
:obj:`ParameterizedHamiltonian`, use the Frechet derivative of the
exponential function.
Parameters
----------
angle: list of float
free parameters to take derivatives with respect to
term_index: int, optional
Index of Parameterized Hamiltonian term that specifies the matrix
direction in which to take the derivative.
Returns
-------
derivative: float
"""
if self.is_unitary or self.is_native_gate:
raise ValueError(
"Can only take derivative of block specified "
"by Hamiltonians or ParameterizedHamiltonian instances.")
if isinstance(self.operator, ParameterizedHamiltonian):
arg = -1j * self.operator.get_hamiltonian(angles)
direction = -1j * self.operator.p_terms[term_index]
return Qobj(
expm_frechet(arg.full(), direction.full(), compute_expm=False),
dims=direction.dims,
)
if len(angles) != 1:
raise ValueError("Expected a single angle for non-"
"ParameterizedHamiltonian instance.")
return self.get_unitary(angles) * -1j * self.operator
def _check_hky_transition_expectations(t):
n = 4
kappa = 3.3
nt_probs = np.array([0.1, 0.2, 0.3, 0.4])
# Get an HKY rate matrix with arbitrary expected rate.
pre_Q = hkymodel.get_pre_Q(kappa, nt_probs)
# Rescale the rate matrix to have expected rate of 1.0.
rates = pre_Q.sum(axis=1)
expected_rate = rates.dot(nt_probs)
pre_Q /= expected_rate
# Convert the pre-rate matrix to an actual rate matrix
# by subtracting row sums from the diagonal.
rates = pre_Q.sum(axis=1)
Q = pre_Q - np.diag(rates)
# Create the transition probability matrix over time t.
P = expm(Q*t)
assert_allclose(P.sum(axis=1), 1)
# Create a joint state distribution matrix.
J = np.diag(nt_probs).dot(P)
assert_allclose(J.sum(), 1)
# Get the expm frechet matrix.
C = pre_Q * t
S = expm_frechet(Q*t, C, compute_expm=False)
# Get the weighted sum of entries of S.
expectation_a = ((S / P) * J).sum()
assert_allclose(expectation_a, t)
# Try an equivalent calculation which does not use P or J.
expectation_b = np.diag(nt_probs).dot(S).sum()
assert_allclose(expectation_b, t)
# Use the library function.
T = nx.DiGraph()
root = 'N0'
edge = ('N0', 'N1')
T.add_edge(*edge)
edge_to_Q = {edge : Q * t}
edge_to_combination = {edge : pre_Q * t}
root_distn = nt_probs
data_weight_pairs = []
for sa, sb in product(range(n), repeat=2):
vec_a = np.zeros(n)
vec_a[sa] = 1
vec_b = np.zeros(n)
vec_b[sb] = 1
data = {'N0' : vec_a, 'N1' : vec_b}
weight = J[sa, sb]
data_weight_pairs.append((data, weight))
edge_to_expectation = get_edge_to_expectation(
T, root, edge_to_Q, edge_to_combination,
root_distn, data_weight_pairs)
expectation_c = edge_to_expectation[edge]
assert_allclose(expectation_c, t)
def test_small_norm_expm_frechet(self):
# methodically test matrices with a range of norms, for better coverage
M_original = np.array([
[1, 2, 3, 4],
[5, 6, 7, 8],
[0, 0, 1, 2],
[0, 0, 5, 6],
], dtype=float)
A_original = np.array([
[1, 2],
[5, 6],
], dtype=float)
E_original = np.array([
[3, 4],
[7, 8],
], dtype=float)
A_original_norm_1 = scipy.linalg.norm(A_original, 1)
selected_m_list = [1, 3, 5, 7, 9, 11, 13, 15]
m_neighbor_pairs = zip(selected_m_list[:-1], selected_m_list[1:])
for ma, mb in m_neighbor_pairs:
ell_a = scipy.linalg._expm_frechet.ell_table_61[ma]
ell_b = scipy.linalg._expm_frechet.ell_table_61[mb]
target_norm_1 = 0.5 * (ell_a + ell_b)
scale = target_norm_1 / A_original_norm_1
M = scale * M_original
A = scale * A_original
E = scale * E_original
expected_expm = scipy.linalg.expm(A)
expected_frechet = scipy.linalg.expm(M)[:2, 2:]
observed_expm, observed_frechet = expm_frechet(A, E)
assert_allclose(expected_expm, observed_expm)
assert_allclose(expected_frechet, observed_frechet)
def test_fuzz(self):
# try a bunch of crazy inputs
rfuncs = (
np.random.uniform,
np.random.normal,
np.random.standard_cauchy,
np.random.exponential)
ntests = 100
for i in range(ntests):
rfunc = random.choice(rfuncs)
target_norm_1 = random.expovariate(1.0)
n = random.randrange(2, 16)
A_original = rfunc(size=(n,n))
E_original = rfunc(size=(n,n))
A_original_norm_1 = scipy.linalg.norm(A_original, 1)
scale = target_norm_1 / A_original_norm_1
A = scale * A_original
E = scale * E_original
M = np.vstack([
np.hstack([A, E]),
np.hstack([np.zeros_like(A), A])])
expected_expm = scipy.linalg.expm(A)
expected_frechet = scipy.linalg.expm(M)[:n, n:]
observed_expm, observed_frechet = expm_frechet(A, E)
assert_allclose(expected_expm, observed_expm)
assert_allclose(expected_frechet, observed_frechet)
def test_small_norm_expm_frechet(self):
# methodically test matrices with a range of norms, for better coverage
M_original = np.array([
[1, 2, 3, 4],
[5, 6, 7, 8],
[0, 0, 1, 2],
[0, 0, 5, 6],
], dtype=float)
A_original = np.array([
[1, 2],
[5, 6],
], dtype=float)
E_original = np.array([
[3, 4],
[7, 8],
], dtype=float)
A_original_norm_1 = scipy.linalg.norm(A_original, 1)
selected_m_list = [1, 3, 5, 7, 9, 11, 13, 15]
m_neighbor_pairs = zip(selected_m_list[:-1], selected_m_list[1:])
for ma, mb in m_neighbor_pairs:
ell_a = scipy.linalg._expm_frechet.ell_table_61[ma]
ell_b = scipy.linalg._expm_frechet.ell_table_61[mb]
target_norm_1 = 0.5 * (ell_a + ell_b)
scale = target_norm_1 / A_original_norm_1
M = scale * M_original
A = scale * A_original
E = scale * E_original
expected_expm = scipy.linalg.expm(A)
expected_frechet = scipy.linalg.expm(M)[:2, 2:]
observed_expm, observed_frechet = expm_frechet(A, E)
assert_allclose(expected_expm, observed_expm)
assert_allclose(expected_frechet, observed_frechet)
def test_fuzz(self):
# try a bunch of crazy inputs
rfuncs = (
np.random.uniform,
np.random.normal,
np.random.standard_cauchy,
np.random.exponential)
ntests = 100
for i in range(ntests):
rfunc = random.choice(rfuncs)
target_norm_1 = random.expovariate(1.0)
n = random.randrange(2, 16)
A_original = rfunc(size=(n,n))
E_original = rfunc(size=(n,n))
A_original_norm_1 = scipy.linalg.norm(A_original, 1)
scale = target_norm_1 / A_original_norm_1
A = scale * A_original
E = scale * E_original
M = np.vstack([
np.hstack([A, E]),
np.hstack([np.zeros_like(A), A])])
expected_expm = scipy.linalg.expm(A)
expected_frechet = scipy.linalg.expm(M)[:n, n:]
observed_expm, observed_frechet = expm_frechet(A, E)
assert_allclose(expected_expm, observed_expm)
assert_allclose(expected_frechet, observed_frechet)
def test_problematic_matrix(self):
# this test case uncovered a bug which has since been fixed
A = np.array([
[1.50591997, 1.93537998],
[0.41203263, 0.23443516],
], dtype=float)
E = np.array([
[1.87864034, 2.07055038],
[1.34102727, 0.67341123],
], dtype=float)
A_norm_1 = scipy.linalg.norm(A, 1)
sps_expm, sps_frechet = expm_frechet(
A, E, method='SPS')
blockEnlarge_expm, blockEnlarge_frechet = expm_frechet(
A, E, method='blockEnlarge')
assert_allclose(sps_expm, blockEnlarge_expm)
assert_allclose(sps_frechet, blockEnlarge_frechet)
def compute_prop_grad(self, k, j, compute_prop=True):
"""
Calculate the gradient of propagator wrt the control amplitude
in the timeslot using the expm_frechet method
The propagtor is calculated (almost) for 'free' in this method
and hence it is returned if compute_prop==True
Returns:
[prop], prop_grad
"""
dyn = self.parent
A = dyn.get_dyn_gen(k)*dyn.tau[k]
E = dyn.get_ctrl_dyn_gen(j)*dyn.tau[k]
if compute_prop:
prop, propGrad = la.expm_frechet(A, E)
return prop, propGrad
else:
propGrad = la.expm_frechet(A, E, compute_expm=False)
return propGrad
def compute_prop_grad(self, k, j, compute_prop=True):
"""
Calculate the gradient of propagator wrt the control amplitude
in the timeslot using the expm_frechet method
The propagtor is calculated (almost) for 'free' in this method
and hence it is returned if compute_prop==True
Returns:
[prop], prop_grad
"""
dyn = self.parent
A = dyn.get_dyn_gen(k) * dyn.tau[k]
E = dyn.get_ctrl_dyn_gen(j) * dyn.tau[k]
if compute_prop:
prop, propGrad = la.expm_frechet(A, E)
return prop, propGrad
else:
propGrad = la.expm_frechet(A, E, compute_expm=False)
return propGrad
def test_expm_frechet(self):
# a test of the basic functionality
M = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [0, 0, 1, 2], [0, 0, 5, 6]], dtype=float)
A = np.array([[1, 2], [5, 6]], dtype=float)
E = np.array([[3, 4], [7, 8]], dtype=float)
expected_expm = scipy.linalg.expm(A)
expected_frechet = scipy.linalg.expm(M)[:2, 2:]
for kwargs in ({}, {"method": "SPS"}, {"method": "blockEnlarge"}):
observed_expm, observed_frechet = expm_frechet(A, E, **kwargs)
assert_allclose(expected_expm, observed_expm)
assert_allclose(expected_frechet, observed_frechet)
def jac(v):
A, R = unvec(v)
ex1 = expm((1 - 2 * alf) * A)
mat = np.array(bmat([[2 * alf * A, -R.T], [R, np.zeros((k, k))]]))
E = np.array(
bmat([[ex1 @ Y1.T @ Y, ex1 @ Y1.T @ Q],
[np.zeros_like(R), np.zeros((k, k))]]))
ex2, fe2 = expm_frechet(mat, E)
M = ex2[:d, :d]
N = ex2[d:, :d]
YMQN = (Y @ M + Q @ N)
partA = asym((1 - 2 * alf) * expm_frechet(
(1 - 2 * alf) * A, Y1.T @ YMQN)[1])
partA += 2 * alf * asym(fe2[:d, :d])
partR = -(fe2[:d, d:].T - fe2[d:, :d])
return vec(partA, partR)
def test_dexpm_2x2(random_2x2):
X = random_2x2
dX = np.random.random((10, 2, 2))
dM_truth = np.zeros((10, 2, 2))
for n in range(10):
M_truth, dM_truth[n] = expm_frechet(X, dX[n])
M, dM = dexpm_2x2(X, dX)
assert np.allclose(M_truth, M)
assert np.allclose(dM_truth, dM)
def jac(v):
eta_P, R = unvec(v)
A = beta/alpha[1]*(X.Pinv@eta_P - [email protected])
x_mat = np.array(
np.bmat([[2*alf*A, -R.T], [R, zeros((k, k))]]))
ex1 = expm(X.Pinv@eta_P)
ex3 = expm((1-2*alf)*A)
exmat = expm(x_mat)
efPinvD = expm_frechet(
X.Pinv@eta_P,
X.P @ ex1 @ X.P - ex3.T @ exmat[:, :p].T @ YPY_rl @
exmat[:, :p] @ ex3 @ X.P)
efxmat = expm_frechet(
x_mat,
np.bmat([[
ex3 @ X.P @ ex1 @
ex3.T @ exmat[:, :p].T @ YPY_rl +
ex3 @ X.P @ ex1 @ ex3.T @ exmat[:, :p].T @ YPY_rl],
[np.zeros((k, p+k))]]))
efA1 = expm_frechet(
(1-2*alf)*A,
(1-2*alf)*beta/alpha[1]*(
X.P @ ex1 @
ex3.T @ exmat[:, :p].T @ YPY_rl @ exmat[:, :p]))
efA2 = expm_frechet(
(1-2*alf)*A, exmat[:, :p].T @ YPY_rl @
exmat[:, :p] @ ex3 @ X.P @ ex1)
grP = 2*efPinvD[1] @ X.Pinv
grP += -2*2*alf*beta/alpha[1]*(efxmat[1][:p, :p] @ X.Pinv -
X.Pinv @ efxmat[1][:p, :p])
grP += -2*efA1[1] @ X.Pinv + 2*X.Pinv @ efA1[1]
grP += 2*(1-2*alf)*beta/alpha[1] * (
ex3.T @ efA2[1]@ ex3.T @ X.Pinv - X.Pinv @ ex3.T @ efA2[1]@ ex3.T)
grR = -2*(-efxmat[1][p:, :p] + efxmat[1][:p, p:].T)
return vec(sym(grP), grR)
def test_expm_jacobian_vector_product(self):
n = 4
x = numpy.random.randn(n, n)
E = numpy.random.randn(n, n)
# use algopy to get the jacobian vector product
ax = UTPM.init_jac_vec(x.flatten(), E.flatten())
ay = expm(ax.reshape((n, n))).reshape((n*n,))
g1 = UTPM.extract_jac_vec(ay)
# compute the jacobian vector product directly using expm_frechet
M = expm_frechet(x, E, compute_expm=False).flatten()
assert_allclose(g1, M, rtol=1e-6)
def test_expm_jacobian_vector_product(self):
n = 4
x = numpy.random.randn(n, n)
E = numpy.random.randn(n, n)
# use algopy to get the jacobian vector product
ax = UTPM.init_jac_vec(x.flatten(), E.flatten())
ay = expm(ax.reshape((n, n))).reshape((n * n, ))
g1 = UTPM.extract_jac_vec(ay)
# compute the jacobian vector product directly using expm_frechet
M = expm_frechet(x, E, compute_expm=False).flatten()
assert_allclose(g1, M, rtol=1e-6)
def get_traj_stats(self, J_other):
n = self.nstates
# compute the observed initial distribution
distn = J_other.sum(axis=1)
# compute conditional expected dwell times
dwell = np.zeros(n)
for i in range(n):
E = np.zeros((n, n), dtype=float)
E[i, i] = 1
interact = expm_frechet(self.Q, E, compute_expm=False)
dwell[i] = (J_other * interact / self.P).sum()
assert_allclose(dwell.sum(), 1)
# compute conditional expected transition counts
trans = np.zeros((n, n), dtype=float)
for i in range(n):
E = np.zeros((n, n), dtype=float)
E[i, 1-i] = 1
interact = expm_frechet(self.Q, self.Q*E, compute_expm=False)
trans[i, 1-i] = (J_other * interact / self.P).sum()
return distn, dwell, trans
def _expm_grad(op, grad):
# We want the backward-mode gradient (left multiplication).
# Let X be the NxN input matrix.
# Let J(X) be the the N^2xN^2 complete Jacobian matrix of expm at X.
# Let Y be the NxN previous gradient in the backward AD (left multiplication)
# We have
# unvec( ( vec(Y)^T . J(X) )^T )
# = unvec( J(X)^T . vec(Y) )
# = unvec( J(X^T) . vec(Y) )
# where the last part (if I am not mistaken) holds in the case of the
# exponential and other matrix power series.
# It can be seen that this is now the forward-mode derivative
# (right multiplication) applied to the Jacobian of the transpose.
grad_func = lambda x, y: expm_frechet(x, y, compute_expm=False)
return tf.py_func(grad_func, [tf.transpose(op.inputs[0]), grad],
tf.float64)
def test_expm_jacobian(self):
n = 4
x = numpy.random.randn(n, n)
# use algopy to get the jacobian
ax = UTPM.init_jacobian(x)
ay = expm(ax)
g1 = UTPM.extract_jacobian(ay)
# compute the jacobian directly using expm_frechet
M = numpy.zeros((n, n, n*n))
ident = numpy.identity(n*n)
for i in range(n*n):
E = ident[i].reshape(n, n)
M[:, :, i] = expm_frechet(x, E, compute_expm=False)
assert_allclose(g1, M, rtol=1e-6)
def test_expm_jacobian(self):
n = 4
x = numpy.random.randn(n, n)
# use algopy to get the jacobian
ax = UTPM.init_jacobian(x)
ay = expm(ax)
g1 = UTPM.extract_jacobian(ay)
# compute the jacobian directly using expm_frechet
M = numpy.zeros((n, n, n * n))
ident = numpy.identity(n * n)
for i in range(n * n):
E = ident[i].reshape(n, n)
M[:, :, i] = expm_frechet(x, E, compute_expm=False)
assert_allclose(g1, M, rtol=1e-6)
def disc_d_state(A: np.ndarray, dA: np.ndarray, dt: float = 1.0) -> Tuple[np.ndarray, np.ndarray]:
"""Discretize the state matrix and its derivative
Args:
A: State matrix
dA: Derivative state matrix
dt: Sampling time
Returns:
2-elements tuple containing
- Ad: Discrete state matrix
- dAd: Derivative discrete state matrix
"""
nj, nx, _ = dA.shape
dAd = np.zeros((nj, nx, nx))
for n in range(nj):
if dA[n].any() or n == 0:
Ad, dAd[n] = expm_frechet(A * dt, dA[n] * dt)
return Ad, dAd
def _disc_dQ_mfd(self, dt: np.ndarray, QQ: np.ndarray, dA: np.ndarray,
dQQ: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
"""Discretization partial derivative of the diffusion matrix by MFD
Args:
dt: Sampling time
QQ: Diffusion matrix
dA: Jacobian state matrix
dQQ: Jacobian diffusion matrix
Returns:
2-elements tuple containing
- **QQd**: Process noise covariance matrix
- **dQQd**: Jacobian process noise covariance matrix
"""
N = dA.shape[0]
AA = np.zeros((2 * self.nx, 2 * self.nx))
AA[:self.nx, :self.nx] = self.A
AA[:self.nx, self.nx:] = QQ
AA[self.nx:, self.nx:] = -self.A.T
dAA = np.zeros((N, 2 * self.nx, 2 * self.nx))
dAA[:, :self.nx, :self.nx] = dA
dAA[:, :self.nx, self.nx:] = dQQ
dAA[:, self.nx:, self.nx:] = -dA.swapaxes(1, 2)
dAAd = np.zeros_like(dAA)
for n in range(N):
if np.any(dAA[n]):
AAd, dAAd[n] = expm_frechet(AA * dt, dAA[n] * dt)
AdT = AAd[:self.nx, :self.nx].T
QQd = AAd[:self.nx, self.nx:] @ AdT
dAd = dAAd[:, :self.nx, :self.nx]
dQQd = dAAd[:, :self.nx, self.
nx:] @ AdT + AAd[:self.nx, self.nx:] @ dAd.swapaxes(1, 2)
return QQd, dQQd
def _compute_prop_grad(self, k, j, compute_prop=True):
"""
Calculate the gradient of propagator wrt the control amplitude
in the timeslot using the expm_frechet method
The propagtor is calculated (almost) for 'free' in this method
and hence it is returned if compute_prop==True
Returns:
[prop], prop_grad
"""
dyn = self.parent
if dyn.oper_dtype == Qobj:
A = dyn._get_phased_dyn_gen(k).full()*dyn.tau[k]
E = dyn._get_phased_ctrl_dyn_gen(k, j).full()*dyn.tau[k]
if compute_prop:
prop_dense, prop_grad_dense = la.expm_frechet(A, E)
prop = Qobj(prop_dense, dims=dyn.dyn_dims)
prop_grad = Qobj(prop_grad_dense,
dims=dyn.dyn_dims)
else:
prop_grad_dense = la.expm_frechet(A, E, compute_expm=False)
prop_grad = Qobj(prop_grad_dense,
dims=dyn.dyn_dims)
elif dyn.oper_dtype == np.ndarray:
A = dyn._get_phased_dyn_gen(k)*dyn.tau[k]
E = dyn._get_phased_ctrl_dyn_gen(k, j)*dyn.tau[k]
if compute_prop:
prop, prop_grad = la.expm_frechet(A, E)
else:
prop_grad = la.expm_frechet(A, E,
compute_expm=False)
else:
# Assuming some sparse matrix
spcls = dyn._dyn_gen[k].__class__
A = (dyn._get_phased_dyn_gen(k)*dyn.tau[k]).toarray()
E = (dyn._get_phased_ctrl_dyn_gen(k, j)*dyn.tau[k]).toarray()
if compute_prop:
prop_dense, prop_grad_dense = la.expm_frechet(A, E)
prop = spcls(prop_dense)
prop_grad = spcls(prop_grad_dense)
else:
prop_grad_dense = la.expm_frechet(A, E, compute_expm=False)
prop_grad = spcls(prop_grad_dense)
if compute_prop:
return prop, prop_grad
else:
return prop_grad
def disc_d_diffusion_mfd(
A: np.ndarray, Q: np.ndarray, dA: np.ndarray, dQ: np.ndarray, dt: float = 1.0
) -> Tuple[np.ndarray, np.ndarray]:
"""Discretize diffusion matrix and its derivative by matrix fraction decomposition
Args:
A: State matrix
Q: Diffusion matrix
dA: Derivative state matrix
dQ: Derivative diffusion matrix
dt: Sampling time
Returns:
2-elements tuple containing
- Process noise covariance matrix
- Derivative process noise covariance matrix
"""
if not Q.any():
return Q
nj, nx, _ = dA.shape
F = block_diag(A, -A.T)
F[:nx, nx:] = Q
Fd = expm(F * dt)[:nx, :]
dF = np.zeros((nj, 2 * nx, 2 * nx))
dF[:, :nx, :nx] = dA
dF[:, :nx, nx:] = dQ
dF[:, nx:, nx:] = -dA.swapaxes(1, 2)
dFd = np.zeros((nj, nx, 2 * nx))
for n in range(nj):
if dF[n].any():
dFd[n] = expm_frechet(F * dt, dF[n] * dt, compute_expm=False)[:nx, :]
Fd12 = Fd[:, nx:]
Fd11 = Fd[:, :nx].T
return Fd12 @ Fd11, dFd[:, :, nx:] @ Fd11 + Fd12 @ dFd[:, :, :nx].swapaxes(1, 2)
def _compute_prop_grad(self, k, j, compute_prop=True):
"""
Calculate the gradient of propagator wrt the control amplitude
in the timeslot using the expm_frechet method
The propagtor is calculated (almost) for 'free' in this method
and hence it is returned if compute_prop==True
Returns:
[prop], prop_grad
"""
dyn = self.parent
if dyn.oper_dtype == Qobj:
A = dyn._get_phased_dyn_gen(k).full() * dyn.tau[k]
E = dyn._get_phased_ctrl_dyn_gen(j).full() * dyn.tau[k]
if compute_prop:
prop_dense, prop_grad_dense = la.expm_frechet(A, E)
dyn._prop[k] = Qobj(prop_dense, dims=dyn.dyn_dims)
dyn._prop_grad[k, j] = Qobj(prop_grad_dense, dims=dyn.dyn_dims)
else:
prop_grad_dense = la.expm_frechet(A, E, compute_expm=False)
dyn._prop_grad[k, j] = Qobj(prop_grad_dense, dims=dyn.dyn_dims)
elif dyn.oper_dtype == np.ndarray:
A = dyn._get_phased_dyn_gen(k) * dyn.tau[k]
E = dyn._get_phased_ctrl_dyn_gen(j) * dyn.tau[k]
if compute_prop:
dyn._prop[k], dyn._prop_grad[k, j] = la.expm_frechet(A, E)
else:
dyn._prop_grad[k, j] = la.expm_frechet(A,
E,
compute_expm=False)
else:
# Assuming some sparse matrix
spcls = dyn._dyn_gen[k].__class__
A = (dyn._get_phased_dyn_gen(k) * dyn.tau[k]).toarray()
E = (dyn._get_phased_ctrl_dyn_gen(j) * dyn.tau[k]).toarray()
if compute_prop:
prop_dense, prop_grad_dense = la.expm_frechet(A, E)
dyn._prop[k] = spcls(prop_dense)
dyn._prop_grad[k, j] = spcls(prop_grad_dense)
else:
prop_grad_dense = la.expm_frechet(A, E, compute_expm=False)
dyn._prop_grad[k, j] = spcls(prop_grad_dense)
if compute_prop:
return dyn._prop[k], dyn._prop_grad[k, j]
else:
return dyn._prop_grad[k, j]
def test_expm_frechet(self):
# a test of the basic functionality
M = np.array([
[1, 2, 3, 4],
[5, 6, 7, 8],
[0, 0, 1, 2],
[0, 0, 5, 6],
], dtype=float)
A = np.array([
[1, 2],
[5, 6],
], dtype=float)
E = np.array([
[3, 4],
[7, 8],
], dtype=float)
expected_expm = scipy.linalg.expm(A)
expected_frechet = scipy.linalg.expm(M)[:2, 2:]
for kwargs in ({}, {'method':'SPS'}, {'method':'blockEnlarge'}):
observed_expm, observed_frechet = expm_frechet(A, E, **kwargs)
assert_allclose(expected_expm, observed_expm)
assert_allclose(expected_frechet, observed_frechet)
def do_cythonized_em(T, root,
edge_to_rate, edge_to_Q, root_distn1d,
data_prob_pairs, guess_edge_to_rate):
"""
Try the Cython implementation.
def expectation_step(
idx_t[:] csr_indices, # (nnodes-1,)
idx_t[:] csr_indptr, # (nnodes+1,)
cnp.float_t[:, :, :] transp, # (nnodes-1, nstates, nstates)
cnp.float_t[:, :, :] transq, # (nnodes-1, nstates, nstates)
cnp.float_t[:, :, :] interact_trans, # (nnodes-1, nstates, nstates)
cnp.float_t[:, :, :] interact_dwell, # (nnodes-1, nstates, nstates)
cnp.float_t[:, :, :] data, # (nsites, nnodes, nstates)
cnp.float_t[:] root_distn, # (nstates,)
cnp.float_t[:, :] trans_out, # (nsites, nnodes-1)
cnp.float_t[:, :] dwell_out, # (nsites, nnodes-1)
int validation=1,
):
"""
# Define a toposort node ordering and a corresponding csr matrix.
nodes = nx.topological_sort(T, [root])
node_to_idx = dict((na, i) for i, na in enumerate(nodes))
m = nx.to_scipy_sparse_matrix(T, nodes)
# Stack the transition rate matrices into a single array.
nnodes = len(nodes)
nstates = root_distn1d.shape[0]
n = nstates
transq = np.empty((nnodes-1, nstates, nstates), dtype=float)
for (na, nb), Q in edge_to_Q.items():
edge_idx = node_to_idx[nb] - 1
transq[edge_idx] = Q
# Allocate a transition probability matrix array
# and some interaction matrix arrays.
transp = np.empty_like(transq)
interact_trans = np.empty_like(transq)
interact_dwell = np.empty_like(transq)
# Stack the data into a single array,
# and construct an array of site weights.
nsites = len(data_prob_pairs)
datas, probs = zip(*data_prob_pairs)
site_weights = np.array(probs, dtype=float)
data = np.empty((nsites, nnodes, nstates), dtype=float)
for site_index, site_data in enumerate(datas):
for i, na in enumerate(nodes):
data[site_index, i] = site_data[na]
# Initialize expectation arrays.
trans_out = np.empty((nsites, nnodes-1), dtype=float)
dwell_out = np.empty((nsites, nnodes-1), dtype=float)
# Initialize the per-edge rate matrix scaling factor guesses.
scaling_guesses = np.empty(nnodes-1, dtype=float)
scaling_ratios = np.ones(nnodes-1, dtype=float)
for (na, nb), rate in guess_edge_to_rate.items():
eidx = node_to_idx[nb] - 1
scaling_guesses[eidx] = rate
# Pre-scale the rate matrix.
transq *= scaling_guesses[:, None, None]
# Do the EM iterations.
nsteps = 1000
for em_iteration_index in range(nsteps):
# Scale the rate matrices according to the edge ratios.
transq *= scaling_ratios[:, None, None]
# Compute the probability transition matrix arrays
# and the interaction matrix arrays.
trans_indicator = np.ones((n, n)) - np.identity(n)
dwell_indicator = np.identity(n)
for edge in T.edges():
na, nb = edge
eidx = node_to_idx[nb] - 1
Q = transq[eidx]
#print(edge, 'Q:')
#print(Q)
transp[eidx] = expm(Q)
interact_trans[eidx] = expm_frechet(
Q, Q * trans_indicator, compute_expm=False)
interact_dwell[eidx] = expm_frechet(
Q, Q * dwell_indicator, compute_expm=False)
# Compute the expectations.
validation = 1
expectation_step(
m.indices, m.indptr,
transp, transq,
interact_trans, interact_dwell,
data,
root_distn1d,
trans_out, dwell_out,
validation,
)
# Compute the per-edge ratios.
trans_sum = (trans_out * site_weights[:, None]).sum(axis=0)
dwell_sum = (dwell_out * site_weights[:, None]).sum(axis=0)
scaling_ratios = trans_sum / -dwell_sum
scaling_guesses *= scaling_ratios
# Report the guesses.
if not (em_iteration_index+1) % 100:
print(em_iteration_index+1)
for edge in T.edges():
na, nb = edge
eidx = node_to_idx[nb] - 1
print(edge, scaling_guesses[eidx])
print()
def disc_d_state_input_expm(
A: np.ndarray,
B: np.ndarray,
dA: np.ndarray,
dB: np.ndarray,
dt: float = 1.0,
order_hold: int = 0,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Discretize the state and input matrices, and their derivatives with the matrix exponential
Args:
A: State matrix
B: Input matrix
dA: Derivative state matrix
dB: Derivative input matrix
dt: Sampling time
order_hold: zero order hold = 0 or first order hold = 1
Returns:
6-elements tuple containing
- Ad: Discrete state matrix
- B0d: Discrete input matrix (zero order hold)
- B1d: Discrete input matrix (first order hold)
- dAd: Derivative discrete state matrix
- dB0d: Derivative discrete input matrix (zero order hold)
- dB1d: Derivative discrete input matrix (first order hold)
"""
nj, nx, nu = dB.shape
if order_hold == 0:
F = np.zeros((nx + nu, nx + nu))
dF = np.zeros((nj, nx + nu, nx + nu))
dFd = np.zeros((nj, nx + nu, nx + nu))
F[:nx, :nx] = A
F[:nx, nx:] = B
dF[:, :nx, :nx] = dA
dF[:, :nx, nx:] = dB
for n in range(nj):
if dF[n].any() or n == 0:
Fd, dFd[n] = expm_frechet(F * dt, dF[n] * dt)
Ad, B0d = np.split(Fd[:nx, :], indices_or_sections=[nx], axis=1)
dAd, dB0d = np.split(dFd[:, :nx, :], indices_or_sections=[nx], axis=2)
B1d = np.zeros((nx, nu))
dB1d = np.zeros((nj, nx, nu))
else:
F = np.zeros((nx + 2 * nu, nx + 2 * nu))
dF = np.zeros((nj, nx + 2 * nu, nx + 2 * nu))
dFd = np.zeros((nj, nx + 2 * nu, nx + 2 * nu))
F[:nx, :nx] = A
F[:nx, nx : nx + nu] = B
F[nx : nx + nu, nx + nu :] = np.eye(nu)
dF[:, :nx, :nx] = dA
dF[:, :nx, nx : nx + nu] = dB
for n in range(nj):
if dF[n].any() or n == 0:
Fd, dFd[n] = expm_frechet(F * dt, dF[n] * dt)
Ad, B0d, B1d = np.split(Fd[:nx, :], indices_or_sections=[nx, nx + nu], axis=1)
dAd, dB0d, dB1d = np.split(dFd[:, :nx, :], indices_or_sections=[nx, nx + nu], axis=2)
return Ad, B0d, B1d, dAd, dB0d, dB1d
def dexp(X, G):
"""Evaluate the derivative of the matrix exponential at
X in direction G.
"""
return np.array([expm_frechet(X[i], G[i])[1] for i in range(X.shape[0])])
def dexp(self,
direction: 'DenseOperator',
tau: complex = 1,
compute_expm: bool = False,
method: str = "spectral",
is_skew_hermitian: bool = False,
epsilon: float = 1e-10) \
-> Union['DenseOperator', Tuple['DenseOperator']]:
"""
Frechet derivative of the matrix exponential.
Parameters
----------
direction: DenseOperator
Direction in which the frechet derivative is calculated. Must be of
the same shape as self.
tau: complex
The matrix is multiplied by tau before calculating the exponential.
compute_expm: bool
If true, then the matrix exponential is calculated and returned as
well.
method: string
Numerical method used for the calculation of the matrix
exponential.
Currently the following are implemented:
- 'Frechet': Uses the scipy linalg matrix exponential for
simultaniously calculation of the frechet derivative expm_frechet
- 'approx': Approximates the Derivative by finite differences.
- 'first_order': First order taylor approximation
- 'second_order': Second order taylor approximation
- 'third_order': Third order taylor approximation
- 'spectral': Use the self implemented spectral decomposition
is_skew_hermitian: bool
Only required, for the method 'spectral'. If set to True, then the
matrix is assumed to be skew hermitian in the spectral
decomposition.
epsilon: float
Width of the finite difference. Only relevant for the method
'approx'.
Returns
-------
prop: DenseOperator
The matrix exponential. Only returned if compute_expm is True!
prop_grad: DenseOperator
The frechet derivative d exp(Ax + B)/dx at x=0 where A is the
direction and B is the matrix stored in self.
Raises
------
NotImplementedError:
If the method given as parameter is not implemented.
"""
prop = None
if type(direction) != DenseOperator:
direction = DenseOperator(direction)
if method == "Frechet":
a = self.data * tau
e = direction.data * tau
if compute_expm:
prop, prop_grad = la.expm_frechet(a, e, compute_expm=True)
else:
prop_grad = la.expm_frechet(a, e, compute_expm=False)
elif method == "spectral":
if compute_expm:
prop, prop_grad = self._dexp_diagonalization(
direction=direction,
tau=tau,
is_skew_hermitian=is_skew_hermitian,
compute_expm=compute_expm)
else:
prop_grad = self._dexp_diagonalization(
direction=direction,
tau=tau,
is_skew_hermitian=is_skew_hermitian,
compute_expm=compute_expm)
elif method == "approx":
d_m = (self.data + epsilon * direction.data) * tau
dprop = la.expm(d_m)
prop = self.exp(tau)
prop_grad = (dprop - prop) * (1 / epsilon)
elif method == "first_order":
if compute_expm:
prop = self.exp(tau)
prop_grad = direction.data * tau
elif method == "second_order":
if compute_expm:
prop = self.exp(tau)
prop_grad = direction.data * tau
prop_grad += (self.data @ direction.data +
direction.data @ self.data) * (tau * tau * 0.5)
elif method == "third_order":
if compute_expm:
prop = self.exp(tau)
prop_grad = direction.data * tau
prop_grad += (self.data @ direction.data +
direction.data @ self.data) * tau * tau * 0.5
prop_grad += (self.data @ self.data @ direction.data +
direction.data @ self.data @ self.data +
self.data @ direction.data @ self.data) * (
tau * tau * tau * 0.16666666666666666)
else:
raise NotImplementedError('The specified method ' + method +
"is not implemented!")
if compute_expm:
if type(prop) != DenseOperator:
prop = DenseOperator(prop)
if type(prop_grad) != DenseOperator:
prop_grad = DenseOperator(prop_grad)
if compute_expm:
return prop, prop_grad
else:
return prop_grad
def _lti_jacobian_disc(self, dt):
"""Discretization of augmented LTI state-space model
Args:
dt: sampling time
Returns:
Ad: Discrete state matrix
B0d, B1d: Discrete input matrix
Qd: Upper Cholesky factor process noise covariance
dAd: Derivative discrete state matrix
dB0d, dB1d: Derivative discrete input matrix
dQd: Derivative upper Cholesky factor process noise covariance
"""
names = self.parameters.names_free
N = len(names)
# Discrete state matrix and its derivative
self._AA[:N, :self.Nx, :self.Nx] = self.A
self._AA[:N, self.Nx:, self.Nx:] = self.A
self._AA[:N, self.Nx:, :self.Nx] = [self.dA[n] for n in names]
dA = self._AA[:N, self.Nx:, :self.Nx]
Ad = expm(self.A * dt)
for i in range(N):
self._AAd[i, :self.Nx, :self.Nx] = Ad
self._AAd[i, self.Nx:, self.Nx:] = Ad
if not np.all(dA[i, :, :] == 0):
self._AAd[i, self.Nx:, :self.Nx] = (expm_frechet(
self.A * dt, dA[i, :, :] * dt, 'SPS', False))
dAd = self._AAd[:N, self.Nx:, :self.Nx]
if not np.all(self.Q == 0):
dQ = np.asarray([self.dQ[k] for k in names])
# transform to covariance matrix
Q2 = self.Q.T @ self.Q
dQ2 = dQ.swapaxes(1, 2) @ self.Q + self.Q.T @ dQ
# covariance matrix of the process noise
Qd2 = lyap(self.A, -Q2 + Ad @ Q2 @ Ad.T)
eq = (-dA @ Qd2 - Qd2 @ dA.swapaxes(1, 2) - dQ2 + dAd @ Q2 @ Ad.T +
Ad @ dQ2 @ Ad.T + Ad @ Q2 @ dAd.swapaxes(1, 2))
# Derivative process noise covariance
dQd2 = np.asarray([lyap(self.A, eq[i, :, :]) for i in range(N)])
# Get Cholesky factor
Qd = pseudo_cholesky(Qd2)
tmp = solve(Qd.T, solve(Qd.T, dQd2).swapaxes(1, 2))
dQd = (np.triu(tmp, 1) + self._Ixx[0, :self.Nx, :self.Nx] * 0.5 *
tmp.diagonal(0, 1, 2)[:, np.newaxis, :]) @ Qd
else:
Qd = self._0xx[0, :, :]
dQd = self._0xx[:N, :, :]
if self.Nu != 0:
# Discrete input matrices and their derivatives
self._BB[:N, :self.Nx, :] = self.B
self._BB[:N, self.Nx:, :] = [self.dB[k] for k in names]
AA = self._AA[:N, :, :]
AAd = self._AAd[:N, :, :]
BB = self._BB[:N, :, :]
bis = solve(AA, AAd - self._Ixx[:N, :, :])
BB0d = bis @ BB
B0d = BB0d[0, :self.Nx, :]
dB0d = BB0d[:, self.Nx:, :]
if self.hold_order == 'foh':
BBd1_free = solve(AA, -bis + AAd * dt) @ BB
B1d = BBd1_free[0, :self.Nx, :]
dB1d = BBd1_free[:, self.Nx:, :]
else:
B1d = self._0xu[0, :, :]
dB1d = self._0xu[:N, :, :]
else:
B0d = self._0xu[0, :, :]
dB0d = self._0xu[:N, :, :]
B1d = B0d
dB1d = dB0d
return Ad, B0d, B1d, Qd, dAd, dB0d, dB1d, dQd
def em_function(
T, node_to_idx, site_weights,
m,
transq_unscaled,
data,
root_distn1d,
mem,
use_log_scale,
scale,
):
"""
This is purely for edge rate scale and does not consider other parameters.
Parameters
----------
T : x
x
node_to_idx : x
x
site_weights : x
x
m : x
x
transq_unscaled : x
x
data : x
x
root_distn1d : x
x
mem : x
x
use_log_scale : bool
True if the scale argument is in logarithmic units.
scale : 1d float ndarray
Edge rate scaling factors, possibly in logarithmic units.
Returns
-------
scale : 1d float ndarray
Updated edge rate scaling factors, possibly in logarithmic units.
"""
#TODO improve docstring and add unit tests
# Transform the scaling factors if necessary.
if use_log_scale:
log_scale = scale
scale = np.exp(scale)
# Unpack some stuff.
nsites, nnodes, nstates = data.shape
# Scale the rate matrices according to the edge ratios.
# Use in-place operations to avoid unnecessary memory copies.
#transq = transq_unscaled * scale[:, None, None]
mem.transq[...] = transq_unscaled
mem.transq *= scale[:, None, None]
# Compute the probability transition matrix arrays
# and the interaction matrix arrays.
for edge in T.edges():
na, nb = edge
eidx = node_to_idx[nb] - 1
Q = mem.transq[eidx]
mem.transp[eidx] = expm(Q)
mem.interact_trans[eidx] = expm_frechet(
Q, Q * mem.trans_indicator, compute_expm=False)
mem.interact_dwell[eidx] = expm_frechet(
Q, Q * mem.dwell_indicator, compute_expm=False)
# Compute the expectations using Cython.
validation = 0
expectation_step(
m.indices, m.indptr,
mem.transp, mem.transq,
mem.interact_trans, mem.interact_dwell,
data,
root_distn1d,
mem.trans_out, mem.dwell_out,
validation,
)
# Compute the per-edge ratios.
trans_sum = (mem.trans_out * site_weights[:, None]).sum(axis=0)
dwell_sum = (mem.dwell_out * site_weights[:, None]).sum(axis=0)
if use_log_scale:
log_scaling_ratios = np.log(trans_sum) - np.log(-dwell_sum)
return log_scale + log_scaling_ratios
else:
scaling_ratios = trans_sum / -dwell_sum
return scale * scaling_ratios
def do_em(T, root, edge_to_rate, edge_to_Q, root_distn1d,
data_prob_pairs, guess_edge_to_rate):
# Extract the number of states.
n = root_distn1d.shape[0]
# Do the EM iterations.
nsteps = 3
for em_iteration_index in range(nsteps):
# Compute the scaled edge-specific transition rate matrices.
edge_to_scaled_Q = {}
for edge in T.edges():
rate = guess_edge_to_rate[edge]
Q = edge_to_Q[edge]
edge_to_scaled_Q[edge] = rate * Q
# Compute the edge-specific transition probability matrices.
edge_to_P = {}
for edge in T.edges():
scaled_Q = edge_to_scaled_Q[edge]
#print(edge, 'Q:')
#print(scaled_Q)
P = expm(scaled_Q)
edge_to_P[edge] = P
# For each edge, compute the interaction matrices
# corresponding to all transition counts and dwell times.
trans_indicator = np.ones((n, n)) - np.identity(n)
dwell_indicator = np.identity(n)
edge_to_interact_trans = {}
edge_to_interact_dwell = {}
for edge in T.edges():
# extract edge-specific transition matrices
Q = edge_to_scaled_Q[edge]
P = edge_to_P[edge]
# compute the transition interaction matrix
interact = expm_frechet(Q, Q * trans_indicator, compute_expm=False)
edge_to_interact_trans[edge] = interact
# compute the dwell interaction matrix
interact = expm_frechet(Q, Q * dwell_indicator, compute_expm=False)
edge_to_interact_dwell[edge] = interact
# Initialize edge-specific summaries.
edge_to_trans_expectation = defaultdict(float)
edge_to_dwell_expectation = defaultdict(float)
# Compute the edge-specific summaries
# conditional on the current edge-specific rate guesses
# and on the leaf state distribution computed
# from the true edge-specific rates.
for node_to_data_fvec1d, lhood in data_prob_pairs:
# Compute the joint endpoint state distribution for each edge.
edge_to_J = get_edge_to_distn2d(
T, edge_to_P, root, root_distn1d, node_to_data_fvec1d)
# For each edge, compute the transition expectation contribution
# and compute the dwell expectation contribution.
# These will be scaled by the data likelihood.
for edge in T.edges():
# extract some edge-specific matrices
P = edge_to_P[edge]
J = edge_to_J[edge]
#print(edge)
#print(P)
#print(J)
#print()
# transition contribution
interact = edge_to_interact_trans[edge]
total = 0
for i in range(n):
for j in range(n):
if J[i, j]:
coeff = J[i, j] / P[i, j]
total += coeff * interact[i, j]
edge_to_trans_expectation[edge] += lhood * total
# dwell contribution
interact = edge_to_interact_dwell[edge]
total = 0
for i in range(n):
for j in range(n):
if J[i, j]:
coeff = J[i, j] / P[i, j]
total += coeff * interact[i, j]
edge_to_dwell_expectation[edge] += lhood * total
# According to EM, update each edge-specific rate guess
# using a ratio of transition and dwell expectations.
for edge in T.edges():
trans = edge_to_trans_expectation[edge]
dwell = edge_to_dwell_expectation[edge]
ratio = trans / -dwell
old_rate = guess_edge_to_rate[edge]
new_rate = old_rate * ratio
guess_edge_to_rate[edge] = new_rate
#print(edge, trans, dwell, ratio, old_rate, new_rate)
print(edge, trans, dwell, new_rate)
print()
def _lti_jacobian_disc(
self, dt: float, dA: np.ndarray, dB: np.ndarray, dQ: np.ndarray
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray,
np.ndarray, np.ndarray, np.ndarray, ]:
"""Discretization of augmented LTI state-space model
Args:
dt: Sampling time
dA: Jacobian state matrix
dB: Jacobian input matrix
dQ: Jacobian Wiener process scaling matrix
Returns:
8-elements tuple containing
- **Ad**: Discrete state matrix
- **B0d**: Discrete input matrix (zero order hold)
- **B1d**: Discrete input matrix (first order hold)
- **Qd**: Upper Cholesky factor of the process noise covariance matrix
- **dAd**: Jacobian discrete state matrix
- **dB0d**: Jacobian discrete input matrix (zero order hold)
- **dB1d**: Jacobian discrete input matrix (first order hold)
- **dQd**: Jacobian of the upper Cholesky factor of the process noise covariance
"""
N = dA.shape[0]
if self.nu == 0:
B0d = np.zeros((self.nx, self.nu))
B1d = B0d
Ad = expm(self.A * dt)
dAd = np.zeros((N, self.nx, self.nx))
for n in range(N):
if np.any(dA[n]):
dAd[n] = expm_frechet(self.A * dt,
dA[n] * dt,
compute_expm=False)
dB0d = np.zeros((N, self.nx, self.nu))
dB1d = dB0d
else:
if self.hold_order == 0:
AA = np.zeros((self.nx + self.nu, self.nx + self.nu))
AA[:self.nx, :self.nx] = self.A
AA[:self.nx, self.nx:] = self.B
dAA = np.zeros((N, self.nx + self.nu, self.nx + self.nu))
dAA[:, :self.nx, :self.nx] = dA
dAA[:, :self.nx, self.nx:] = dB
dAAd = np.zeros_like(dAA)
for n in range(N):
if np.any(dAA[n]):
AAd, dAAd[n] = expm_frechet(AA * dt, dAA[n] * dt)
Ad = AAd[:self.nx, :self.nx]
B0d = AAd[:self.nx, self.nx:]
dAd = dAAd[:, :self.nx, :self.nx]
dB0d = dAAd[:, :self.nx, self.nx:]
B1d = np.zeros((self.nx, self.nu))
dB1d = np.zeros((N, self.nx, self.nu))
else:
AA = np.zeros((self.nx + 2 * self.nu, self.nx + 2 * self.nu))
AA[:self.nx, :self.nx] = self.A
AA[:self.nx, self.nx:self.nx + self.nu] = self.B
AA[self.nx:self.nx + self.nu,
self.nx + self.nu:] = np.eye(self.nu)
dAA = np.zeros(
(N, self.nx + 2 * self.nu, self.nx + 2 * self.nu))
dAA[:, :self.nx, :self.nx] = dA
dAA[:, :self.nx, self.nx:self.nx + self.nu] = dB
dAAd = np.zeros_like(dAA)
for n in range(N):
if np.any(dAA[n]):
AAd, dAAd[n] = expm_frechet(AA * dt, dAA[n] * dt)
Ad = AAd[:self.nx, :self.nx]
B0d = AAd[:self.nx, self.nx:self.nx + self.nu]
B1d = AAd[:self.nx, self.nx + self.nu:]
dAd = dAAd[:, :self.nx, :self.nx]
dB0d = dAAd[:, :self.nx, self.nx:self.nx + self.nu]
dB1d = dAAd[:, :self.nx, self.nx + self.nu:]
if np.any(self.Q):
# transform to covariance matrix
QQ = self.Q.T @ self.Q
dQQ = dQ.swapaxes(1, 2) @ self.Q + self.Q.T @ dQ
if self.method == 'mfd':
QQd, dQQd = self._disc_dQ_mfd(dt, QQ, dA, dQQ)
elif self.method == 'lyapunov':
QQd, dQQd = self._disc_dQ_lyapunov(dt, QQ, dA, dQQ, Ad, dAd)
else:
raise ValueError('`Invalid discretization method`')
Qd = nearest_cholesky(QQd)
try:
tmp = np.linalg.solve(
Qd.T,
np.linalg.solve(Qd.T, dQQd).swapaxes(1, 2))
except (LinAlgError, RuntimeError):
inv_Qd = np.linalg.pinv(Qd)
tmp = inv_Qd.T @ dQQd @ inv_Qd
dQd = (np.triu(tmp, 1) + np.eye(self.nx) / 2 *
tmp.diagonal(0, 1, 2)[:, None, :]) @ Qd
else:
Qd = np.zeros((self.nx, self.nx))
dQd = np.zeros((N, self.nx, self.nx))
return Ad, B0d, B1d, Qd, dAd, dB0d, dB1d, dQd
def s_expm(s_y, y, a):
return sla.expm_frechet(a, s_y.T, compute_expm=False).T
