def resampling_method_factory(resampling_method_enum, resampling_kwargs): if resampling_method_enum == ResamplingMethodsEnum.MULTINOMIAL: resampling_method = MultinomialResampler() elif resampling_method_enum == ResamplingMethodsEnum.SYSTEMATIC: resampling_method = SystematicResampler() elif resampling_method_enum == ResamplingMethodsEnum.STRATIFIED: resampling_method = StratifiedResampler() elif resampling_method_enum == ResamplingMethodsEnum.REGULARIZED: resampling_method = RegularisedTransform(**resampling_kwargs) elif resampling_method_enum == ResamplingMethodsEnum.VARIANCE_CORRECTED: regularized_resampler = RegularisedTransform(**resampling_kwargs) resampling_method = PartiallyCorrectedRegularizedTransform( regularized_resampler) elif resampling_method_enum == ResamplingMethodsEnum.OPTIMIZED: lr = resampling_kwargs.pop('lr', resampling_kwargs.pop('learning_rate', 0.1)) loss = SinkhornLoss(**resampling_kwargs, symmetric=True) optimizer = SGD(loss, lr=lr, decay=0.95) regularized_resampler = RegularisedTransform(**resampling_kwargs) resampling_method = OptimizedPointCloud( optimizer, intermediate_resampler=regularized_resampler) else: raise ValueError( f'resampling_method_name {resampling_method_enum} is not a valid ResamplingMethodsEnum' ) return resampling_method
def main(resampling_method_value, resampling_neff, resampling_kwargs=None, T=150, batch_size=50, n_particles=25, data_seed=0, filter_seed=555, savefig=False): transition_matrix = 0.5 * np.eye(2, dtype=np.float32) transition_covariance = np.eye(2, dtype=np.float32) observation_matrix = np.eye(2, dtype=np.float32) observation_covariance = 0.1 * np.eye(2, dtype=np.float32) resampling_method_enum = ResamplingMethodsEnum(resampling_method_value) np_random_state = np.random.RandomState(seed=data_seed) data, kf = get_data(transition_matrix, observation_matrix, transition_covariance, observation_covariance, T, np_random_state) observation_dataset = tf.data.Dataset.from_tensor_slices(data) if resampling_method_enum == ResamplingMethodsEnum.KALMAN: return kalman_main(kf, data, savefig) if resampling_kwargs is None: resampling_kwargs = {} if resampling_neff == 0.: resampling_criterion = NeverResample() elif resampling_neff == 1.: resampling_criterion = AlwaysResample() else: resampling_criterion = NeffCriterion(resampling_neff, True) if resampling_method_enum == ResamplingMethodsEnum.MULTINOMIAL: resampling_method = MultinomialResampler() elif resampling_method_enum == ResamplingMethodsEnum.SYSTEMATIC: resampling_method = SystematicResampler() elif resampling_method_enum == ResamplingMethodsEnum.STRATIFIED: resampling_method = StratifiedResampler() elif resampling_method_enum == ResamplingMethodsEnum.REGULARIZED: resampling_method = RegularisedTransform(**resampling_kwargs) elif resampling_method_enum == ResamplingMethodsEnum.VARIANCE_CORRECTED: regularized_resampler = RegularisedTransform(**resampling_kwargs) resampling_method = PartiallyCorrectedRegularizedTransform(regularized_resampler) elif resampling_method_enum == ResamplingMethodsEnum.OPTIMIZED: lr = resampling_kwargs.pop('lr', resampling_kwargs.pop('learning_rate', 0.1)) loss = SinkhornLoss(**resampling_kwargs, symmetric=True) optimizer = SGD(loss, lr=lr, decay=0.95) regularized_resampler = RegularisedTransform(**resampling_kwargs) resampling_method = OptimizedPointCloud(optimizer, intermediate_resampler=regularized_resampler) else: raise ValueError(f'resampling_method_name {resampling_method_enum} is not a valid ResamplingMethodsEnum') observation_matrix = tf.convert_to_tensor(observation_matrix) transition_covariance_chol = tf.linalg.cholesky(transition_covariance) observation_covariance_chol = tf.linalg.cholesky(observation_covariance) initial_particles = np_random_state.normal(0., 1., [batch_size, n_particles, 2]).astype(np.float32) initial_state = State(tf.constant(initial_particles)) smc = make_filter(observation_matrix, transition_matrix, observation_covariance_chol, transition_covariance_chol, resampling_method, resampling_criterion) states = get_states(smc, initial_state, observation_dataset, tf.constant(T), tf.constant(filter_seed)) stddevs = std(states, keepdims=False).numpy() stddevs_df = stddevs
def main(resampling_method_value, resampling_neff, learning_rates=(1e-4, 1e-3), resampling_kwargs=None, alpha=0.42, dx=10, dy=3, observation_covariance=1., dense=False, T=20, batch_size=1, n_particles=25, data_seed=0, n_data=50, n_iter=50, savefig=False, filter_seed=0, use_xla=False, change_seed=True): transition_matrix = get_transition_matrix(alpha, dx) transition_covariance = get_transition_covariance(dx) observation_matrix = get_observation_matrix(dx, dy, dense) observation_covariance = get_observation_covariance( observation_covariance, dy) resampling_method_enum = ResamplingMethodsEnum(resampling_method_value) np_random_state = np.random.RandomState(seed=data_seed) observation_matrix = tf.convert_to_tensor(observation_matrix) transition_covariance_chol = tf.linalg.cholesky(transition_covariance) observation_covariance_chol = tf.linalg.cholesky(observation_covariance) initial_particles = np_random_state.normal( 0., 1., [batch_size, n_particles, dx]).astype(np.float32) initial_state = State(initial_particles) if resampling_neff == 0.: resampling_criterion = NeverResample() elif resampling_neff == 1.: resampling_criterion = AlwaysResample() else: resampling_criterion = NeffCriterion(resampling_neff, True) optimal_smc = make_optimal_filter(observation_matrix, transition_matrix, observation_covariance_chol, transition_covariance_chol, MultinomialResampler(), resampling_criterion) if resampling_kwargs is None: resampling_kwargs = {} resampling_method = resampling_method_factory(resampling_method_enum, resampling_kwargs) datas = [] lls = [] observation_datasets = [] optimal_lls = [] log_phi_x_0 = tf.ones(dx) phi_y_0 = tf.zeros(dy) for _ in range(n_data): data, ll = get_data(transition_matrix, observation_matrix, transition_covariance, observation_covariance, T, np_random_state) datas.append(data) lls.append(ll / T) observation_dataset = tf.data.Dataset.from_tensor_slices(data) observation_datasets.append(observation_dataset) final_state = optimal_smc(initial_state, observation_dataset, T, None, True, filter_seed) optimal_lls.append(final_state.log_likelihoods.numpy().mean() / T) log_phi_x = tf.Variable(log_phi_x_0, trainable=True) phi_y = tf.Variable(phi_y_0, trainable=True) smc = make_filter(observation_matrix, transition_matrix, observation_covariance_chol, transition_covariance_chol, resampling_method, resampling_criterion, log_phi_x, phi_y) def optimizer_maker(learning_rate): # tf.function doesn't like creating variables. This is a way to create them outside the graph # We can't reuse the same optimizer because it would be giving a warmed-up momentum to the ones run later optimizer = tf.optimizers.Adam(learning_rate=learning_rate) return optimizer initial_values = [log_phi_x_0, phi_y_0] losses_list = [] ess_profiles_list = [] mean_errors = [] for observation_dataset in observation_datasets: try: losses, ess_profiles = compare_learning_rates( smc, initial_state, observation_dataset, T, log_phi_x, phi_y, initial_values, n_iter, optimizer_maker, learning_rates, filter_seed, use_xla, change_seed) except: print('one dataset failed, ignoring') continue losses_df = pd.DataFrame(np.stack(losses).T, columns=np.log10(learning_rates)) ess_df = pd.DataFrame(np.stack(ess_profiles).T, columns=np.log10(learning_rates)) losses_df.columns.name = 'log learning rate' losses_df.columns.epoch = 'epoch' ess_df.columns.name = 'log learning rate' ess_df.columns.epoch = 'epoch' losses_list.append(losses_df) ess_profiles_list.append(ess_df) delta_phi_m_1 = tf.linalg.diag(tf.exp(-log_phi_x)) diff_cov = optimal_smc._proposal_model._sigma - delta_phi_m_1 @ transition_covariance approx_error = tf.linalg.diag_part(diff_cov).numpy() mean_error = np.sqrt(np.nanmean(approx_error**2)) mean_errors.append(mean_error) losses_data = pd.concat(losses_list, axis=1) ess_data = pd.concat(ess_profiles_list, axis=1) mean_data = pd.DataFrame([[np.mean(mean_errors)]], index=pd.MultiIndex.from_tuples([(batch_size, n_particles)]), columns=pd.MultiIndex.from_tuples([ (resampling_method_enum.name, change_seed) ])) losses_data = losses_data.groupby(axis=1, level=0).mean() ess_data = ess_data.groupby(axis=1, level=0).mean() # plot_losses(losses_df, resampling_method_enum.name, savefig, dx, dy, dense, T, change_seed) plot_losses_vs_ess(losses_data, ess_data, resampling_method_enum.name, savefig, dx, dy, dense, T, n_particles, change_seed, batch_size, np.mean(optimal_lls), np.mean(lls), n_iter, mean_data, n_data) print(tf.exp(log_phi_x))
def main(resampling_method_value, resampling_neff, resampling_kwargs=None, T=150, batch_size=50, n_particles=25, data_seed=0, values=(0.25, 0.5, 0.75), filter_seed=555, savefig=False): transition_matrix = 0.5 * np.eye(2, dtype=np.float32) transition_covariance = np.eye(2, dtype=np.float32) observation_matrix = np.eye(2, dtype=np.float32) observation_covariance = 0.1 * np.eye(2, dtype=np.float32) values = np.array(list(zip(values, values)), dtype=np.float32) resampling_method_enum = ResamplingMethodsEnum(resampling_method_value) np_random_state = np.random.RandomState(seed=data_seed) data, kf = get_data(transition_matrix, observation_matrix, transition_covariance, observation_covariance, T, np_random_state) observation_dataset = tf.data.Dataset.from_tensor_slices(data) if resampling_method_enum == 6: return kalman_main(kf, data, values, T, savefig) if resampling_kwargs is None: resampling_kwargs = {} if resampling_neff == 0.: resampling_criterion = NeverResample() elif resampling_neff == 1.: resampling_criterion = AlwaysResample() else: resampling_criterion = NeffCriterion(resampling_neff, True) if resampling_method_enum == ResamplingMethodsEnum.MULTINOMIAL: resampling_method = MultinomialResampler() elif resampling_method_enum == ResamplingMethodsEnum.SYSTEMATIC: resampling_method = SystematicResampler() elif resampling_method_enum == ResamplingMethodsEnum.STRATIFIED: resampling_method = StratifiedResampler() elif resampling_method_enum == ResamplingMethodsEnum.REGULARIZED: resampling_method = RegularisedTransform(**resampling_kwargs) elif resampling_method_enum == ResamplingMethodsEnum.VARIANCE_CORRECTED: regularized_resampler = RegularisedTransform(**resampling_kwargs) resampling_method = PartiallyCorrectedRegularizedTransform( regularized_resampler) elif resampling_method_enum == ResamplingMethodsEnum.OPTIMIZED: lr = resampling_kwargs.pop('lr', resampling_kwargs.pop('learning_rate', 0.1)) loss = SinkhornLoss(**resampling_kwargs, symmetric=True) optimizer = SGD(loss, lr=lr, decay=0.95) regularized_resampler = RegularisedTransform(**resampling_kwargs) resampling_method = OptimizedPointCloud( optimizer, intermediate_resampler=regularized_resampler) else: raise ValueError( f'resampling_method_name {resampling_method_enum} is not a valid ResamplingMethodsEnum' ) init_transition_matrix = (0.5 * np.eye(2) + 0.1 * np_random_state.randn(2, 2)).astype( np.float32) modifiable_transition_matrix = tf.Variable(init_transition_matrix, trainable=True) observation_matrix = tf.convert_to_tensor(observation_matrix) transition_covariance_chol = tf.linalg.cholesky(transition_covariance) observation_covariance_chol = tf.linalg.cholesky(observation_covariance) initial_particles = np_random_state.normal( 0., 1., [batch_size, n_particles, 2]).astype(np.float32) initial_state = State(tf.constant(initial_particles)) smc = make_filter(observation_matrix, modifiable_transition_matrix, observation_covariance_chol, transition_covariance_chol, resampling_method, resampling_criterion) elbos = get_elbos(smc, initial_state, observation_dataset, tf.constant(T), modifiable_transition_matrix, tf.constant(values), tf.constant(filter_seed)) elbos_df = pd.DataFrame( elbos.numpy(), pd.Index(values[:, 0], name=r'$\theta_1$, $\theta_2$')) elbos_df = elbos_df.T.describe().T[['mean', 'std']].reset_index() if savefig: filename = f"{resampling_method_enum.name}_batchsize_{batch_size}_N_{n_particles}_epsilon_{resampling_kwargs.get('epsilon')}_likelihoods_values.tex" elbos_df.to_latex(buf=os.path.join('./tables/', filename), float_format='{:,.3f}'.format, escape=False, index=False) else: print( elbos_df.to_latex(float_format='{:,.3f}'.format, escape=False, index=False))
def main(resampling_method_value, resampling_neff, resampling_kwargs=None, T=100, batch_size=1, n_particles=25, phi=0.5, data_seed=0, filter_seed=1, learning_rate=0.001, n_iter=50, savefig=False, use_xla=False, batch_data=1, assume_differentiable=False, change_seed=False): transition_matrix = phi * np.eye(2, dtype=np.float32) transition_covariance = 0.5 * np.eye(2, dtype=np.float32) observation_matrix = np.eye(2, dtype=np.float32) observation_covariance = 0.1 * np.eye(2, dtype=np.float32) resampling_method_enum = ResamplingMethodsEnum(resampling_method_value) np_random_state = np.random.RandomState(seed=data_seed) data = [] np_data = [] assert batch_data > 0 for _ in range(batch_data): a_data, kf = get_data(transition_matrix, observation_matrix, transition_covariance, observation_covariance, T, np_random_state) data.append(tf.data.Dataset.from_tensor_slices(a_data)) np_data.append(a_data) if resampling_kwargs is None: resampling_kwargs = {} if resampling_neff == 0.: resampling_criterion = NeverResample() elif resampling_neff == 1.: resampling_criterion = AlwaysResample() else: resampling_criterion = NeffCriterion(resampling_neff, True) if resampling_method_enum == ResamplingMethodsEnum.MULTINOMIAL: resampling_method = MultinomialResampler() elif resampling_method_enum == ResamplingMethodsEnum.SYSTEMATIC: resampling_method = SystematicResampler() elif resampling_method_enum == ResamplingMethodsEnum.STRATIFIED: resampling_method = StratifiedResampler() elif resampling_method_enum == ResamplingMethodsEnum.REGULARIZED: resampling_method = RegularisedTransform(**resampling_kwargs) elif resampling_method_enum == ResamplingMethodsEnum.VARIANCE_CORRECTED: regularized_resampler = RegularisedTransform(**resampling_kwargs) resampling_method = PartiallyCorrectedRegularizedTransform( regularized_resampler) elif resampling_method_enum == ResamplingMethodsEnum.OPTIMIZED: lr = resampling_kwargs.pop('lr', resampling_kwargs.pop('learning_rate', 0.1)) loss = SinkhornLoss(**resampling_kwargs, symmetric=True) optimizer = SGD(loss, lr=lr, decay=0.95) regularized_resampler = RegularisedTransform(**resampling_kwargs) resampling_method = OptimizedPointCloud( optimizer, intermediate_resampler=regularized_resampler) else: raise ValueError( f'resampling_method_name {resampling_method_enum} is not a valid ResamplingMethodsEnum' ) modifiable_transition_matrix = tf.Variable(transition_matrix, trainable=True) observation_matrix = tf.convert_to_tensor(observation_matrix) transition_covariance_chol = tf.linalg.cholesky(transition_covariance) observation_covariance_chol = tf.linalg.cholesky(observation_covariance) initial_particles = np_random_state.normal( 0., 1., [batch_size, n_particles, 2]).astype(np.float32) initial_state = State(initial_particles) smc = make_filter(observation_matrix, modifiable_transition_matrix, observation_covariance_chol, transition_covariance_chol, resampling_method, resampling_criterion) x0 = np.array([2 * phi / 3] * 2).astype(np.float32) print(x0) if resampling_method.DIFFERENTIABLE or assume_differentiable: loss_fun = lambda x, observation_dataset, seed: values_and_gradient( x, modifiable_transition_matrix, smc, initial_state, observation_dataset, T, seed) else: loss_fun = lambda x, observation_dataset, seed: values_and_gradient_finite_diff( x, modifiable_transition_matrix, smc, initial_state, observation_dataset, T, seed) final_values = [] losses = [] kalman_params = [] def kf_likelihood_fun(val, data): import copy kf_copy = copy.copy(kf) kf_copy.transition_matrices = np.diag(val) return -kf_loglikelihood(kf_copy, data) fun = tf.function(loss_fun, experimental_compile=use_xla) for observation_dataset, np_dataset in tqdm(zip(data, np_data), total=batch_data): final_value, loss = gradient_descent(fun, x0, observation_dataset, tf.constant(learning_rate), tf.constant(n_iter), tf.constant(filter_seed), tf.constant(change_seed)) final_values.append(final_value.numpy()) losses.append(loss.numpy()) kf_params = minimize(kf_likelihood_fun, x0, args=(np_dataset, )) kalman_params.append(kf_params.x) losses = np.array(losses).T plt.plot(losses) plt.show() final_values = np.vstack(final_values) kalman_params = np.vstack(kalman_params) df = pd.DataFrame(final_values - kalman_params, columns=[r'$\theta_1', r'$\theta_2']) parameters_diff = np.mean(np.square(df), 0) if savefig: filename = f'theta_diff_{resampling_method_enum.name}_batch_size_{batch_size}_N_{n_particles}_batch_data_{batch_data}_changeseed_{change_seed}.csv' df.to_csv(os.path.join('./tables/', filename), float_format='%.5f') else: print(parameters_diff.to_latex(float_format='%.5f'))
def main(resampling_method_value, resampling_neff, resampling_kwargs=None, T=100, batch_size=1, n_particles=25, data_seed=0, filter_seed=1, mesh_size=10, savefig=True, use_tqdm=False, use_xla=False, diff_epsilon=1e-1, optimal_proposal=False): v = 1. t = .1 transition_matrix = np.array([[1., 1.], [0., 1.]], dtype=np.float32) transition_covariance = v**2 * np.array([[1 / 3, 1 / 2], [1 / 2, 1.]], dtype=np.float32) observation_matrix = np.array([[1., 0]], dtype=np.float32) observation_covariance = np.array([[t**2]], dtype=np.float32) resampling_method_enum = ResamplingMethodsEnum(resampling_method_value) x_linspace = np.linspace(0.95, 1., mesh_size).astype(np.float32) y_linspace = np.linspace(0.95, 1., mesh_size).astype(np.float32) mesh = np.asanyarray([(x, y) for x in x_linspace for y in y_linspace]) np_random_state = np.random.RandomState(seed=data_seed) data, kf = get_data(transition_matrix, observation_matrix, transition_covariance, observation_covariance, T, np_random_state) if resampling_method_enum == ResamplingMethodsEnum.KALMAN: return kalman_main(kf, data, mesh, mesh_size, 1e-2, use_tqdm, savefig) observation_dataset = tf.data.Dataset.from_tensor_slices(data) if resampling_kwargs is None: resampling_kwargs = {} if resampling_neff == 0.: resampling_criterion = NeverResample() elif resampling_neff == 1.: resampling_criterion = AlwaysResample() else: resampling_criterion = NeffCriterion(resampling_neff, True) if resampling_method_enum == ResamplingMethodsEnum.MULTINOMIAL: resampling_method = MultinomialResampler() elif resampling_method_enum == ResamplingMethodsEnum.SYSTEMATIC: resampling_method = SystematicResampler() elif resampling_method_enum == ResamplingMethodsEnum.STRATIFIED: resampling_method = StratifiedResampler() elif resampling_method_enum == ResamplingMethodsEnum.REGULARIZED: resampling_method = RegularisedTransform(**resampling_kwargs) elif resampling_method_enum == ResamplingMethodsEnum.VARIANCE_CORRECTED: regularized_resampler = RegularisedTransform(**resampling_kwargs) resampling_method = PartiallyCorrectedRegularizedTransform( regularized_resampler) elif resampling_method_enum == ResamplingMethodsEnum.OPTIMIZED: lr = resampling_kwargs.pop('lr', resampling_kwargs.pop('learning_rate', 0.1)) loss = SinkhornLoss(**resampling_kwargs, symmetric=True) optimizer = SGD(loss, lr=lr, decay=0.95) regularized_resampler = RegularisedTransform(**resampling_kwargs) resampling_method = OptimizedPointCloud( optimizer, intermediate_resampler=regularized_resampler) else: raise ValueError( f'resampling_method_name {resampling_method_enum} is not a valid ResamplingMethodsEnum' ) modifiable_transition_matrix = tf.Variable(transition_matrix, trainable=False) observation_matrix = tf.convert_to_tensor(observation_matrix) transition_covariance_chol = tf.linalg.cholesky(transition_covariance) observation_covariance_chol = tf.linalg.cholesky(observation_covariance) initial_particles = np_random_state.normal( 0., .01, [batch_size, n_particles, 2]).astype(np.float32) initial_state = State(initial_particles) smc = make_filter(observation_matrix, modifiable_transition_matrix, observation_covariance_chol, transition_covariance_chol, resampling_method, resampling_criterion, optimal_proposal=optimal_proposal) # if resampling_method.DIFFERENTIABLE: get_method = tf.function(get_surface, experimental_compile=use_xla) # else: # fun = partial(get_surface_finite_difference, diff_epsilon=diff_epsilon) # get_method = tf.function(fun, experimental_compile=use_xla) log_likelihoods, gradients = get_method(mesh, modifiable_transition_matrix, smc, initial_state, False, observation_dataset, T, filter_seed, use_tqdm) plot_surface(mesh, mesh_size, log_likelihoods.numpy(), resampling_method_enum.name, resampling_kwargs, n_particles, savefig) plot_vector_field(mesh, mesh_size, log_likelihoods.numpy(), gradients.numpy(), resampling_method_enum.name, resampling_kwargs, n_particles, savefig)
def main(resampling_method_value, resampling_neff, learning_rates=(1e-4, 1e-3), resampling_kwargs=None, currencies=('EUR', 'GBP', 'CAD'), batch_size=1, n_particles=25, api_key='', start_date="2019-09-02", end_date="2020-01-02", n_iter=50, savefig=False, filter_seed=0, use_xla=False, change_seed=True): data = get_data(currencies, api_key, start_date, end_date) M = len(currencies) T = len(data) resampling_method_enum = ResamplingMethodsEnum(resampling_method_value) observation_dataset = tf.data.Dataset.from_tensor_slices(data) if resampling_kwargs is None: resampling_kwargs = {} if resampling_neff == 0.: resampling_criterion = NeverResample() elif resampling_neff == 1.: resampling_criterion = AlwaysResample() else: resampling_criterion = NeffCriterion(resampling_neff, True) if resampling_method_enum == ResamplingMethodsEnum.MULTINOMIAL: resampling_method = MultinomialResampler() elif resampling_method_enum == ResamplingMethodsEnum.SYSTEMATIC: resampling_method = SystematicResampler() elif resampling_method_enum == ResamplingMethodsEnum.STRATIFIED: resampling_method = StratifiedResampler() elif resampling_method_enum == ResamplingMethodsEnum.REGULARIZED: resampling_method = RegularisedTransform(**resampling_kwargs) elif resampling_method_enum == ResamplingMethodsEnum.VARIANCE_CORRECTED: regularized_resampler = RegularisedTransform(**resampling_kwargs) resampling_method = PartiallyCorrectedRegularizedTransform( regularized_resampler) elif resampling_method_enum == ResamplingMethodsEnum.OPTIMIZED: lr = resampling_kwargs.pop('lr', resampling_kwargs.pop('learning_rate', 0.1)) loss = SinkhornLoss(**resampling_kwargs, symmetric=True) optimizer = SGD(loss, lr=lr, decay=0.95) regularized_resampler = RegularisedTransform(**resampling_kwargs) resampling_method = OptimizedPointCloud( optimizer, intermediate_resampler=regularized_resampler) elif resampling_method_enum == ResamplingMethodsEnum.CORRECTED: resampling_method = CorrectedRegularizedTransform(**resampling_kwargs) else: raise ValueError( f'resampling_method_name {resampling_method_enum} is not a valid ResamplingMethodsEnum' ) np_random_state = np.random.RandomState(seed=555) initial_particles = np_random_state.normal( 1., 0.5, [batch_size, n_particles, M]).astype(np.float32) initial_state = State(initial_particles) large_initial_particles = np_random_state.normal( 1., 0.5, [25, n_particles, M]).astype(np.float32) large_initial_state = State(large_initial_particles) mu_init = -5. * tf.ones(M) F_init = 0.9 * tf.eye(M) transition_cov_init = 0.35 * tf.eye(M) observation_cov_init = 1. * tf.eye(M) mu = tf.Variable(mu_init, trainable=True) F = tf.Variable(F_init, trainable=True) transition_cov = tf.Variable(transition_cov_init, trainable=True) observation_cov = tf.Variable(observation_cov_init, trainable=False) smc = make_filter(mu, F, transition_cov, observation_cov, resampling_method, resampling_criterion) surrogate_smc = make_filter(mu, F, transition_cov, observation_cov, SystematicResampler(), resampling_criterion) def optimizer_maker(learning_rate): # tf.function doesn't like creating variables. This is a way to create them outside the graph # We can't reuse the same optimizer because it would be giving a warmed-up momentum to the ones run later optimizer = tf.optimizers.Adam(learning_rate=learning_rate) return optimizer variables = [mu, F, transition_cov] initial_values = [mu_init, F_init, transition_cov_init] losses, ess_profiles = compare_learning_rates( smc, initial_state, observation_dataset, T, variables, initial_values, n_iter, optimizer_maker, learning_rates, filter_seed, change_seed, large_initial_state, surrogate_smc) losses_df = pd.DataFrame(np.stack(losses).T, columns=np.log10(learning_rates)) ess_df = pd.DataFrame(np.stack(ess_profiles).T, columns=np.log10(learning_rates)) losses_df.columns.name = 'log learning rate' losses_df.columns.epoch = 'epoch' ess_df.columns.name = 'log learning rate' ess_df.columns.epoch = 'epoch' # plot_losses(losses_df, resampling_method_enum.name, savefig, dx, dy, dense, T, change_seed) plot_losses_vs_ess(losses_df, ess_df, resampling_method_enum.name, savefig, M, n_particles, change_seed, batch_size, n_iter, resampling_kwargs.get("epsilon")) print(mu) print(F) print(transition_cov) print(mu_init) print(F_init) print(transition_cov_init)