Add tests for sampling
@@ -69,5 +69,44 @@ def net():
return kn.BayesNet(cauchy.get_connected_graph())
-def test_metropolis(net):
- samples = kn.sample(net=net, sample_from=net.get_latent_vertices(), draws=1)
+@pytest.mark.parametrize("algo", [
+ ("metropolis"),
+ ("NUTS"),
+ ("hamiltonian")
+])
+def test_sampling_returns_dict_of_list_of_ndarrays_for_vertices_in_sample_from(algo, net):
+ draws = 5
+ sample_from = net.get_latent_vertices()
+ vertex_ids = [vertex.get_id() for vertex in sample_from]
+
+ samples = kn.sample(net=net, sample_from=sample_from, algo=algo, draws=draws)
+
+ assert len(samples) == len(vertex_ids)
+ assert type(samples) == dict
+
+ for vertex_id, vertex_samples in samples.items():
+ assert vertex_id in vertex_ids
+
+ assert len(vertex_samples) == draws
+ assert type(vertex_samples) == list
+ assert all(type(sample) == np.ndarray for sample in vertex_samples)
+
+
+def test_dropping_samples(net):
+ draws = 10
+ drop = 3
+
+ samples = kn.sample(net=net, sample_from=net.get_latent_vertices(), draws=draws, drop=drop)
+
+ expected_num_samples = draws - drop
+ assert all(len(vertex_samples) == expected_num_samples for vertex_id, vertex_samples in samples.items())
+
+
+def test_down_sample_interval(net):
+ draws = 10
+ down_sample_interval = 2
+
+ samples = kn.sample(net=net, sample_from=net.get_latent_vertices(), draws=draws, down_sample_interval=down_sample_interval)
+
+ expected_num_samples = draws / down_sample_interval
+ assert all(len(vertex_samples) == expected_num_samples for vertex_id, vertex_samples in samples.items())
