def setUp(self):
np.random.seed(42)
self.horizon = tf.constant(100.)
self.batch_size = 1
self.n_particles = 5
self.dimension = 2
choice = np.random.binomial(
1, 0.25, [self.batch_size, self.n_particles]).astype(bool)
w = np.random.uniform(
0., 1., [self.batch_size, self.n_particles]).astype(np.float32)
w[choice] = 1e-3
self.w = w / w.sum(1, keepdims=True)
self.log_w = tf.math.log(self.w)
self.ones = np.ones_like(self.w) / self.n_particles
self.log_uniform = tf.math.log(self.ones)
self.np_x = np.random.normal(
0.1, 0.7,
[self.batch_size, self.n_particles, self.dimension]).astype(
np.float32)
self.x = tf.constant(self.np_x)
self.transport_matrix = transport(self.x, self.log_w, tf.constant(0.5),
tf.constant(0.85), tf.constant(1e-3),
tf.constant(500),
tf.constant(self.n_particles))
self.naive_instance = NaiveSolver(horizon=self.horizon,
threshold=tf.constant(1e-3),
step_size=tf.constant(0.05))
self.petkov_instance = PetkovSolver(tf.constant(10),
use_newton_schulze=False)
def test_transport(self):
T_scaled = transport(self.x, self.degenerate_logw, self.epsilon, 0.75,
self.threshold, self.n_iter, self.n_particles)
self.assertAllClose(tf.constant(self.degenerate_weights) *
tf.cast(self.n_particles, float),
tf.reduce_sum(T_scaled, 1),
atol=1e-2)
self.assertAllClose(tf.reduce_sum(T_scaled, 2),
tf.ones_like(self.degenerate_logw),
atol=1e-2)
self.assertAllClose(tf.reduce_sum(T_scaled, [1, 2]),
tf.cast(self.n_particles, float) *
tf.ones([self.batch_size]),
atol=1e-4)
scale_x = diameter(self.x, self.x).numpy()[0]
np_x = (self.np_x[0] -
np.mean(self.np_x[0], 0, keepdims=True)) / scale_x
np_transport_matrix = ot.bregman.empirical_sinkhorn(
np_x, np_x, self.np_epsilon**0.5, b=self.degenerate_weights[0])
self.assertAllClose(T_scaled[0] @ self.x[0],
np_transport_matrix *
self.n_particles.numpy() @ self.np_x[0],
atol=1e-3)
def fun_logw(logw):
logw = normalize(logw, 1, self.n_particles)
transport_matrix = transport(self.x, logw, self.epsilon,
tf.constant(0.9), self.threshold,
self.n_iter, self.n_particles)
return tf.math.reduce_sum(
tf.linalg.matmul(transport_matrix, self.x))
def test_penalty(self):
penalties = []
for epsilon in np.linspace(0.05, 5., 50, dtype=np.float32):
T_scaled = transport(self.x, self.degenerate_logw,
tf.constant(epsilon), 0.75, self.threshold,
self.n_iter, self.n_particles) / tf.cast(
self.n_particles, float)
temp = tf.math.log(T_scaled) - tf.expand_dims(
self.uniform_logw, -1) - tf.expand_dims(
self.degenerate_logw, 1)
temp -= 1
temp *= T_scaled
penalties.append(-tf.reduce_sum(temp, [1, 2]).numpy() - 1)
import matplotlib.pyplot as plt
plt.plot(np.linspace(0.05, 2., 50), penalties)
plt.show()
def apply(self, state: State, flags: tf.Tensor, seed=None):
""" Resampling method
:param state State
Particle filter state
:param flags: tf.Tensor
Flags for resampling
:return: resampled state
:rtype: State
"""
# TODO: The real batch_size is the sum of flags. We shouldn't do more operations than we need...
if self.additional_variables_are_state:
particles_list = [state.particles]
for additional_state_variable in state.ADDITIONAL_STATE_VARIABLES:
particles_list.append(getattr(state, additional_state_variable))
particles = tf.concat(particles_list, axis=-1)
else:
particles = state.particles
transport_matrix = transport(particles, state.log_weights, self.epsilon, self.scaling,
self.convergence_threshold, self.max_iter, state.n_particles)
return apply_transport_matrix(state, transport_matrix, flags)
def apply(self, state: State, flags: tf.Tensor, seed=None):
""" Resampling method
:param state State
Particle filter state
:param flags: tf.Tensor
Flags for resampling
:return: resampled state
:rtype: State
"""
# TODO: The real batch_size is the sum of flags. We shouldn't do more operations than we need...
transport_matrix = transport(state.particles, state.log_weights,
self.epsilon, self.scaling,
self.convergence_threshold, self.max_iter,
state.n_particles)
weights = state.weights
transport_correction = self.ricatti_solver(transport_matrix, weights)
if not self.propagate_correction_gradient:
transport_correction = tf.stop_gradient(transport_correction)
res = apply_transport_matrix(state,
transport_matrix + transport_correction,
flags)
return res
def fun_x(x):
transport_matrix = transport(x, self.degenerate_logw, self.epsilon,
tf.constant(0.75), self.threshold,
self.n_iter, self.n_particles)
return tf.math.reduce_sum(tf.linalg.matmul(transport_matrix, x))