Repository: deeprob-org/deeprob-kit (language: python)

Issue:
Introduce multithreaded implementation of forward and backward evaluation of SPNs
- The forward evaluation (used for EVI, MAR and MPE queries and sampling) can be parallelized by considering a layered topological ordering of the SPN graph. That is, every leaf node can be evaluated in parallel and, after that, every parent node can be computed in parallel as well, and so on.
- The backward evaluation (used for MPE query and sampling) can be parallelized by considering a layered topological ordering of the SPN graph, as for forward evaluation.
- Moreover, introduce unit tests ensuring the correctness of the implementation.

A suitable multiprocessing library for this task is ```joblib```, for which it is possible to specify ```'threading'``` as lightweight backend.


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index c4a8697..9f210bd 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,4 +5,5 @@ __pycache__
 .coverage
 build
 dist
+*.json
 *.egg-info
diff --git a/Makefile b/Makefile
index dea684d..2003668 100644
--- a/Makefile
+++ b/Makefile
@@ -31,9 +31,7 @@ unit_tests:
 
 # Run benchmarks
 benchmarks:
-	export PYTHONPATH=. && \
-	$(PYTHON) $(BENCHMARK_DIR)/clt_queries.py && \
-	$(PYTHON) $(BENCHMARK_DIR)/spn_queries.py
+	for SCRIPT in $(wildcard $(BENCHMARK_DIR)/run_*.py); do PYTHONPATH=. $(PYTHON) $$SCRIPT; done
 
 # Upload the PIP package
 pip_upload: pip_package
diff --git a/benchmark/README.md b/benchmark/README.md
index 7c7a105..d0c924a 100644
--- a/benchmark/README.md
+++ b/benchmark/README.md
@@ -2,14 +2,13 @@
 
 The `benchmark` directory contains benchmark scripts of models and algorithms implemented in the library.
 All the scripts can be launched by command line.
-Every script will print out an estimation of the time required to run an algorithm (expressed in milliseconds).
-
-The following table contains a description about the benchmark scripts and the external libraries used for comparison.
+Every script will print an estimation of the time required to run each algorithm (in seconds) to a JSON file.
+The following table contains a description about the available benchmark scripts.
 Please install the packages in `requirements.txt` to be able to run the scripts.
 
-|       Benchmark      |                                    Description                               |  Compared Libraries |
-|----------------------|------------------------------------------------------------------------------|:-------------------:|
-| clt_queries.py       | Benchmark on Binary Chow-Liu Trees (CLTs): learning, inference and sampling. | [*SPFlow*][SPFlow]  |      
-| spn_queries.py       | Benchmark on Sum-Product Networks (SPNs): inference and sampling.            | [*SPFlow*][SPFlow]  |
-
-[SPFlow]: https://github.com/SPFlow/SPFlow 
+|       Benchmark      |                        Description                        |
+|----------------------|-----------------------------------------------------------|
+| run_deeprob_clt.py   | Benchmark on DeeProb-kit of Binary Chow-Liu Trees (CLTs). |
+| run_deeprob_spn.py   | Benchmark on DeeProb-kit of Sum-Product Networks (SPNs).  |
+| run_spflow_clt.py    | Benchmark on SPFlow of Binary Chow-Liu Trees (CLTs).      |
+| run_spflow_spn.py    | Benchmark on SPFlow of Sum-Product Networks (SPNs).       |
diff --git a/benchmark/clt_queries.py b/benchmark/clt_queries.py
deleted file mode 100644
index 06f5dd3..0000000
--- a/benchmark/clt_queries.py
+++ /dev/null
@@ -1,185 +0,0 @@
-import gc
-import time
-import numpy as np
-
-from utils import compute_mean_stddev_times, load_dataset
-
-from spn.structure.leaves.cltree.CLTree import CLTree
-from spn.structure.leaves.cltree.MLE import update_cltree_parameters_mle
-from spn.structure.leaves.cltree.Inference import cltree_log_likelihood
-from spn.structure.leaves.cltree.MPE import cltree_mpe
-from spn.structure.leaves.cltree.Sampling import sample_cltree_node
-
-import deeprob.spn.structure as spn
-
-
-def deeprob_learn_binary_clt(data: np.ndarray) -> spn.BinaryCLT:
-    n_features = data.shape[1]
-    scope = list(range(n_features))
-    domain = [[0, 1]] * n_features
-    clt = spn.BinaryCLT(scope, root=0)
-    clt.fit(data, domain, alpha=0.1, random_state=42)
-    return clt
-
-
-def spflow_learn_binary_clt(data: np.ndarray) -> CLTree:
-    n_features = data.shape[1]
-    scope = list(range(n_features))
-    clt = CLTree(scope, data)
-    update_cltree_parameters_mle(clt, data, alpha=0.1)
-    return clt
-
-
-if __name__ == '__main__':
-    gc.disable()
-    n_repetitions = 50
-    data_train, evi_data, mar_data = load_dataset('ad', n_samples=5000, seed=42)
-    evi_data = evi_data.astype(np.int64)
-    n_samples = 1000
-
-    print()
-    print("Benchmark of Binary-CLTs on DeeProb-kit ...")
-
-    # ---- DeeProb-kit Learn Binary CLT -------------------------------------------------------------------------------
-    elapsed_times = list()
-    for _ in range(n_repetitions):
-        gc.collect()
-        start_time = time.perf_counter()
-        clt = deeprob_learn_binary_clt(data_train)
-        end_time = time.perf_counter()
-        elapsed_times.append(1e3 * (end_time - start_time))
-    mean_time, stddev_time = compute_mean_stddev_times(elapsed_times)
-    print("[Chow-Liu] DeeProb-kit Avg. Time: {:.2f}ms (+- {:.2f})".format(
-        mean_time, stddev_time
-    ))
-
-    # ---- DeeProb-kit EVI Inference ----------------------------------------------------------------------------------
-    elapsed_times = list()
-    for _ in range(n_repetitions):
-        gc.collect()
-        start_time = time.perf_counter()
-        lls = clt.log_likelihood(evi_data)
-        end_time = time.perf_counter()
-        elapsed_times.append(1e3 * (end_time - start_time))
-    mean_time, stddev_time = compute_mean_stddev_times(elapsed_times)
-    print("[EVI] DeeProb-kit Avg. Time: {:.2f}ms (+- {:.2f}) - LL: {:.6f}".format(
-        mean_time, stddev_time, np.mean(lls)
-    ))
-
-    # ---- DeeProb-kit MAR Inference ----------------------------------------------------------------------------------
-    elapsed_times = list()
-    for _ in range(n_repetitions):
-        gc.collect()
-        start_time = time.perf_counter()
-        lls = clt.log_likelihood(mar_data)
-        end_time = time.perf_counter()
-        elapsed_times.append(1e3 * (end_time - start_time))
-    mean_time, stddev_time = compute_mean_stddev_times(elapsed_times)
-    print("[MAR] DeeProb-kit Avg. Time: {:.2f}ms (+- {:.2f}) - LL: {:.6f}".format(
-        mean_time, stddev_time, np.mean(lls)
-    ))
-
-    # ---- DeeProb-kit MPE Inference ----------------------------------------------------------------------------------
-    elapsed_times = list()
-    for _ in range(n_repetitions):
-        gc.collect()
-        start_time = time.perf_counter()
-        mpe_data = clt.mpe(mar_data)
-        end_time = time.perf_counter()
-        elapsed_times.append(1e3 * (end_time - start_time))
-    lls = clt.log_likelihood(mpe_data)
-    mean_time, stddev_time = compute_mean_stddev_times(elapsed_times)
-    print("[MPE] DeeProb-kit Avg. Time: {:.2f}ms (+- {:.2f}) - LL: {:.6f}".format(
-        mean_time, stddev_time, np.mean(lls)
-    ))
-
-    # ---- DeeProb-kit Sampling ---------------------------------------------------------------------------------------
-    mis_data = np.full(shape=(n_samples, evi_data.shape[1]), fill_value=np.nan, dtype=np.float32)
-    elapsed_times = list()
-    for _ in range(n_repetitions):
-        gc.collect()
-        start_time = time.perf_counter()
-        _ = clt.sample(mis_data)
-        end_time = time.perf_counter()
-        elapsed_times.append(1e3 * (end_time - start_time))
-    mean_time, stddev_time = compute_mean_stddev_times(elapsed_times)
-    print("[Sampling] DeeProb-kit Avg. Time: {:.2f}ms (+- {:.2f})".format(
-        mean_time, stddev_time
-    ))
-
-    print()
-    print("Benchmark of Binary-CLTs on SPFlow ...")
-
-    # ---- SPFlow Learn Binary CLT ------------------------------------------------------------------------------------
-    elapsed_times = list()
-    for _ in range(n_repetitions):
-        gc.collect()
-        start_time = time.perf_counter()
-        clt = spflow_learn_binary_clt(data_train)
-        end_time = time.perf_counter()
-        elapsed_times.append(1e3 * (end_time - start_time))
-    mean_time, stddev_time = compute_mean_stddev_times(elapsed_times)
-    print("[Chow-Liu] SPFlow Avg. Time: {:.2f}ms (+- {:.2f})".format(
-        mean_time, stddev_time
-    ))
-
-    # ---- SPFlow EVI Inference ---------------------------------------------------------------------------------------
-    elapsed_times = list()
-    for _ in range(n_repetitions):
-        gc.collect()
-        start_time = time.perf_counter()
-        lls = cltree_log_likelihood(clt, evi_data, dtype=np.float32)
-        end_time = time.perf_counter()
-        elapsed_times.append(1e3 * (end_time - start_time))
-    mean_time, stddev_time = compute_mean_stddev_times(elapsed_times)
-    print("[EVI] SPFlow Avg. Time: {:.2f}ms (+- {:.2f}) - LL: {:.6f}".format(
-        mean_time, stddev_time, np.mean(lls)
-    ))
-
-    # ---- SPFlow MAR Inference ---------------------------------------------------------------------------------------
-    elapsed_times = list()
-    for _ in range(n_repetitions):
-        gc.collect()
-        start_time = time.perf_counter()
-        lls = cltree_log_likelihood(clt, mar_data, dtype=np.float32)
-        end_time = time.perf_counter()
-        elapsed_times.append(1e3 * (end_time - start_time))
-    mean_time, stddev_time = compute_mean_stddev_times(elapsed_times)
-    print("[MAR] SPFlow Avg. Time: {:.2f}ms (+- {:.2f}) - LL: {:.6f}".format(
-        mean_time, stddev_time, np.mean(lls)
-    ))
-
-    # ---- SPFlow MPE Inference ---------------------------------------------------------------------------------------
-    logprobs = np.empty(len(mar_data), dtype=np.float32)
-    elapsed_times = list()
-    for _ in range(n_repetitions):
-        gc.collect()
-        start_time = time.perf_counter()
-        mar_data_copied = mar_data.copy()
-        cltree_mpe(clt, mar_data_copied, logprobs=logprobs)
-        end_time = time.perf_counter()
-        mpe_data = mar_data_copied
-        elapsed_times.append(1e3 * (end_time - start_time))
-    mpe_data = mpe_data.astype(np.int64)
-    lls = cltree_log_likelihood(clt, mpe_data, dtype=np.float32)
-    mean_time, stddev_time = compute_mean_stddev_times(elapsed_times)
-    print("[MPE] SPFlow Avg. Time: {:.2f}ms (+- {:.2f}) - LL: {:.6f}".format(
-        mean_time, stddev_time, np.mean(lls)
-    ))
-
-    # The implementation of Sampling in SPFlow have an infinite loop
-    """
-    # ---- SPFlow Sampling --------------------------------------------------------------------------------------------
-    rand_gen = np.random.RandomState(42)
-    elapsed_times = list()
-    for _ in range(n_repetitions):
-        gc.collect()
-        start_time = time.perf_counter()
-        _ = sample_cltree_node(clt, n_samples=n_samples, data=None, rand_gen=rand_gen)  # data is not used ...
-        end_time = time.perf_counter()
-        elapsed_times.append(1e3 * (end_time - start_time))
-    mean_time, stddev_time = compute_mean_stddev_times(elapsed_times)
-    print("[Sampling] SPFlow Avg. Time: {:.2f}ms (+- {:.2f})".format(
-        mean_time, stddev_time
-    ))
-    """
diff --git a/benchmark/run_deeprob_clt.py b/benchmark/run_deeprob_clt.py
new file mode 100644
index 0000000..9fe1353
--- /dev/null
+++ b/benchmark/run_deeprob_clt.py
@@ -0,0 +1,94 @@
+import gc
+import time
+import json
+import numpy as np
+
+from utils import compute_mean_stddev_times, load_benchmark_data
+
+import deeprob.spn.structure as spn
+
+
+def learn_binary_clt(data: np.ndarray) -> spn.BinaryCLT:
+    n_features = data.shape[1]
+    scope = list(range(n_features))
+    domain = [[0, 1]] * n_features
+    clt = spn.BinaryCLT(scope, root=0)
+    clt.fit(data, domain, alpha=0.1, random_state=42)
+    return clt
+
+
+if __name__ == '__main__':
+    with open('benchmark/settings.json', 'r') as f:
+        settings = json.load(f)
+    dataset = settings['dataset']
+    n_repetitions = settings['n_repetitions']
+    n_samples = settings['n_samples']
+    results = dict()
+
+    train_data, evi_data, mar_data = load_benchmark_data(dataset, seed=42)
+    n_features = train_data.shape[1]
+
+    gc.collect()
+    gc.disable()
+
+    print("Benchmarking Binary-CLTs on DeeProb-kit ...")
+    clt = learn_binary_clt(train_data)
+
+    print("\tBenchmarking Chiow-Liu Algorithm ...")
+    elapsed_times = list()
+    for _ in range(n_repetitions):
+        gc.collect()
+        start_time = time.perf_counter()
+        learn_binary_clt(train_data)
+        end_time = time.perf_counter()
+        elapsed_times.append(end_time - start_time)
+    results['chowliu'] = compute_mean_stddev_times(elapsed_times)
+
+    print("\tBenchmarking EVI Inference ...")
+    elapsed_times = list()
+    for _ in range(n_repetitions):
+        gc.collect()
+        start_time = time.perf_counter()
+        clt.log_likelihood(evi_data)
+        end_time = time.perf_counter()
+        elapsed_times.append(end_time - start_time)
+    results['evi'] = compute_mean_stddev_times(elapsed_times)
+
+    print("\tBenchmarking MAR Inference ...")
+    elapsed_times = list()
+    for _ in range(n_repetitions):
+        gc.collect()
+        start_time = time.perf_counter()
+        clt.log_likelihood(mar_data)
+        end_time = time.perf_counter()
+        elapsed_times.append(end_time - start_time)
+    results['mar'] = compute_mean_stddev_times(elapsed_times)
+
+    print("\tBenchmarking MPE Inference ...")
+    elapsed_times = list()
+    for _ in range(n_repetitions):
+        gc.collect()
+        start_time = time.perf_counter()
+        clt.mpe(mar_data)
+        end_time = time.perf_counter()
+        elapsed_times.append(end_time - start_time)
+    results['mpe'] = compute_mean_stddev_times(elapsed_times)
+
+    print("\tBenchmarking Ancestral Sampling ...")
+    mis_data = np.full((n_samples, n_features), fill_value=np.nan, dtype=np.float32)
+    elapsed_times = list()
+    for _ in range(n_repetitions):
+        gc.collect()
+        start_time = time.perf_counter()
+        clt.sample(mar_data)
+        end_time = time.perf_counter()
+        elapsed_times.append(end_time - start_time)
+    results['anc_sampling'] = compute_mean_stddev_times(elapsed_times)
+
+    # Save the benchmark results to file
+    results_filepath = 'deeprob-clt-benchmark.json'
+    with open(results_filepath, 'w') as f:
+        results = json.loads(json.dumps(results), parse_float=lambda x: round(float(x), 2))
+        json.dump({'deeprob': {'clt': results}}, f, indent=2)
+    print("Saved benchmark results to {}".format(results_filepath))
+    print()
diff --git a/benchmark/run_deeprob_spn.py b/benchmark/run_deeprob_spn.py
new file mode 100644
index 0000000..9ba274f
--- /dev/null
+++ b/benchmark/run_deeprob_spn.py
@@ -0,0 +1,93 @@
+import gc
+import time
+import json
+import numpy as np
+
+from utils import compute_mean_stddev_times, load_benchmark_data
+
+from deeprob.context import ContextState
+import deeprob.spn.structure as spn
+import deeprob.spn.algorithms as spnalg
+
+
+def learn_binary_clt(data: np.ndarray) -> spn.BinaryCLT:
+    n_features = data.shape[1]
+    scope = list(range(n_features))
+    domain = [[0, 1]] * n_features
+    clt = spn.BinaryCLT(scope, root=0)
+    clt.fit(data, domain, alpha=0.1, random_state=42)
+    return clt
+
+
+def learn_binary_spn(data: np.ndarray) -> spn.Node:
+    return learn_binary_clt(data).to_pc()
+
+
+if __name__ == '__main__':
+    with open('benchmark/settings.json', 'r') as f:
+        settings = json.load(f)
+    dataset = settings['dataset']
+    n_repetitions = settings['n_repetitions']
+    n_samples = settings['n_samples']
+    results = dict()
+
+    train_data, evi_data, mar_data = load_benchmark_data(dataset, seed=42)
+    n_features = train_data.shape[1]
+
+    gc.collect()
+    gc.disable()
+
+    print("Benchmarking SPNs on DeeProb-kit ...")
+    root = learn_binary_spn(train_data)
+
+    print("\tBenchmarking EVI Inference ...")
+    elapsed_times = list()
+    with ContextState(check_spn=False):
+        for _ in range(n_repetitions):
+            gc.collect()
+            start_time = time.perf_counter()
+            spnalg.log_likelihood(root, evi_data)
+            end_time = time.perf_counter()
+            elapsed_times.append(end_time - start_time)
+    results['evi'] = compute_mean_stddev_times(elapsed_times)
+
+    print("\tBenchmarking MAR Inference ...")
+    elapsed_times = list()
+    with ContextState(check_spn=False):
+        for _ in range(n_repetitions):
+            gc.collect()
+            start_time = time.perf_counter()
+            spnalg.log_likelihood(root, mar_data)
+            end_time = time.perf_counter()
+            elapsed_times.append(end_time - start_time)
+    results['mar'] = compute_mean_stddev_times(elapsed_times)
+
+    print("\tBenchmarking MPE Inference ...")
+    elapsed_times = list()
+    with ContextState(check_spn=False):
+        for _ in range(n_repetitions):
+            gc.collect()
+            start_time = time.perf_counter()
+            spnalg.mpe(root, mar_data, inplace=False)
+            end_time = time.perf_counter()
+            elapsed_times.append(end_time - start_time)
+    results['mpe'] = compute_mean_stddev_times(elapsed_times)
+
+    print("\tBenchmarking Conditional Sampling ...")
+    elapsed_times = list()
+    with ContextState(check_spn=False):
+        for _ in range(n_repetitions):
+            gc.collect()
+            start_time = time.perf_counter()
+            spnalg.sample(root, mar_data, inplace=False)
+            end_time = time.perf_counter()
+            elapsed_times.append(end_time - start_time)
+    results['cond_sampling'] = compute_mean_stddev_times(elapsed_times)
+
+    # Save the benchmark results to file
+    results_filepath = 'deeprob-spn-benchmark.json'
+    with open(results_filepath, 'w') as f:
+        results = json.loads(json.dumps(results), parse_float=lambda x: round(float(x), 2))
+        json.dump({'deeprob': {'spn': results}}, f, indent=2)
+    print("Saved benchmark results to {}".format(results_filepath))
+    print()
diff --git a/benchmark/run_spflow_clt.py b/benchmark/run_spflow_clt.py
new file mode 100644
index 0000000..ecc8bd2
--- /dev/null
+++ b/benchmark/run_spflow_clt.py
@@ -0,0 +1,103 @@
+import gc
+import time
+import json
+import numpy as np
+
+from utils import compute_mean_stddev_times, load_benchmark_data
+
+from spn.structure.leaves.cltree.CLTree import CLTree
+from spn.structure.leaves.cltree.MLE import update_cltree_parameters_mle
+from spn.structure.leaves.cltree.Inference import cltree_log_likelihood
+from spn.structure.leaves.cltree.MPE import cltree_mpe
+from spn.structure.leaves.cltree.Sampling import sample_cltree_node
+
+
+def learn_binary_clt(data: np.ndarray) -> CLTree:
+    n_features = data.shape[1]
+    scope = list(range(n_features))
+    clt = CLTree(scope, data)
+    update_cltree_parameters_mle(clt, data, alpha=0.1)
+    return clt
+
+
+if __name__ == '__main__':
+    with open('benchmark/settings.json', 'r') as f:
+        settings = json.load(f)
+    dataset = settings['dataset']
+    n_repetitions = settings['n_repetitions']
+    n_samples = settings['n_samples']
+    results = dict()
+
+    train_data, evi_data, mar_data = load_benchmark_data(dataset, seed=42)
+    n_features = train_data.shape[1]
+
+    gc.collect()
+    gc.disable()
+
+    print("Benchmarking Binary-CLTs on SPFlow ...")
+    clt = learn_binary_clt(train_data)
+
+    print("\tBenchmarking Chiow-Liu Algorithm ...")
+    elapsed_times = list()
+    for _ in range(n_repetitions):
+        gc.collect()
+        start_time = time.perf_counter()
+        learn_binary_clt(train_data)
+        end_time = time.perf_counter()
+        elapsed_times.append(end_time - start_time)
+    results['chowliu'] = compute_mean_stddev_times(elapsed_times)
+
+    print("\tBenchmarking EVI Inference ...")
+    elapsed_times = list()
+    for _ in range(n_repetitions):
+        gc.collect()
+        start_time = time.perf_counter()
+        cltree_log_likelihood(clt, evi_data, dtype=np.float32)
+        end_time = time.perf_counter()
+        elapsed_times.append(end_time - start_time)
+    results['evi'] = compute_mean_stddev_times(elapsed_times)
+
+    print("\tBenchmarking MAR Inference ...")
+    elapsed_times = list()
+    for _ in range(n_repetitions):
+        gc.collect()
+        start_time = time.perf_counter()
+        cltree_log_likelihood(clt, mar_data, dtype=np.float32)
+        end_time = time.perf_counter()
+        elapsed_times.append(end_time - start_time)
+    results['mar'] = compute_mean_stddev_times(elapsed_times)
+
+    print("\tBenchmarking MPE Inference ...")
+    logprobs = np.empty(len(mar_data), dtype=np.float32)
+    elapsed_times = list()
+    for _ in range(n_repetitions):
+        gc.collect()
+        start_time = time.perf_counter()
+        mar_data_copied = mar_data.copy()
+        cltree_mpe(clt, mar_data_copied, logprobs=logprobs)
+        end_time = time.perf_counter()
+        elapsed_times.append(end_time - start_time)
+    results['mpe'] = compute_mean_stddev_times(elapsed_times)
+
+    # The implementation of ancestral sampling on CLTs in SPFlow have an infinite loop
+    results['anc_sampling'] = [None, None]
+    """
+    print("\tBenchmarking Ancestral Sampling ...")
+    mis_data = np.full((n_samples, n_features), fill_value=np.nan, dtype=np.float32)
+    elapsed_times = list()
+    for _ in range(n_repetitions):
+        gc.collect()
+        start_time = time.perf_counter()
+        sample_cltree_node(clt, n_samples=n_samples, data=None, rand_gen=rand_gen)  # data is not used ...
+        end_time = time.perf_counter()
+        elapsed_times.append(end_time - start_time)
+    results['anc_sampling'] = compute_mean_stddev_times(elapsed_times)
+    """
+
+    # Save the benchmark results to file
+    results_filepath = 'spflow-clt-benchmark.json'
+    with open(results_filepath, 'w') as f:
+        results = json.loads(json.dumps(results), parse_float=lambda x: round(float(x), 2))
+        json.dump({'spflow': {'clt': results}}, f, indent=2)
+    print("Saved benchmark results to {}".format(results_filepath))
+    print()
diff --git a/benchmark/run_spflow_spn.py b/benchmark/run_spflow_spn.py
new file mode 100644
index 0000000..cde1410
--- /dev/null
+++ b/benchmark/run_spflow_spn.py
@@ -0,0 +1,123 @@
+import gc
+import time
+import json
+from typing import Union, List
+import numpy as np
+
+from utils import compute_mean_stddev_times, load_benchmark_data
+
+from spn.algorithms.Inference import log_likelihood
+from spn.algorithms.MPE import mpe
+from spn.algorithms.Sampling import sample_instances
+from spn.structure.Base import Node, Sum, Product, assign_ids, rebuild_scopes_bottom_up
+from spn.structure.leaves.parametric.Parametric import Bernoulli
+
+from deeprob.utils.graph import build_tree_structure
+import deeprob.spn.structure as spn
+
+
+def learn_binary_clt(data: np.ndarray) -> spn.BinaryCLT:
+    n_features = data.shape[1]
+    scope = list(range(n_features))
+    domain = [[0, 1]] * n_features
+    clt = spn.BinaryCLT(scope, root=0)
+    clt.fit(data, domain, alpha=0.1, random_state=42)
+    return clt
+
+
+def learn_binary_spn(data: np.ndarray) -> Node:  # This function is copy-pasted from DeeProb-kit BinaryCLT.to_pc()
+    clt = learn_binary_clt(data)
+    root = build_tree_structure(clt.tree, scope=clt.scope)
+    factors = {clt.scope[i]: np.exp(clt.params[i]) for i in range(len(clt.tree))}
+    neg_buffer, pos_buffer = [], []
+    nodes_stack = [root]
+    last_node_visited = None
+    while nodes_stack:
+        node = nodes_stack[-1]
+        if node.is_leaf() or (last_node_visited in node.get_children()):
+            leaves: List[Union[Bernoulli, Sum]] = [
+                Bernoulli(p=0.0, scope=node.get_id()),
+                Bernoulli(p=1.0, scope=node.get_id()),
+            ]
+            if not node.is_leaf():
+                neg_prod = Product(children=[leaves[0]] + neg_buffer[-len(node.get_children()):])
+                pos_prod = Product(children=[leaves[1]] + pos_buffer[-len(node.get_children()):])
+                del neg_buffer[-len(node.get_children()):]
+                del pos_buffer[-len(node.get_children()):]
+                sum_children = [neg_prod, pos_prod]
+            else:
+                sum_children = leaves
+            weights = factors[node.get_id()]
+            neg_buffer.append(Sum(children=sum_children, weights=weights[0]))
+            pos_buffer.append(Sum(children=sum_children, weights=weights[1]))
+            last_node_visited = nodes_stack.pop()
+        else:
+            nodes_stack.extend(node.get_children())
+    return rebuild_scopes_bottom_up(assign_ids(pos_buffer[0]))
+
+
+if __name__ == '__main__':
+    with open('benchmark/settings.json', 'r') as f:
+        settings = json.load(f)
+    dataset = settings['dataset']
+    n_repetitions = settings['n_repetitions']
+    n_samples = settings['n_samples']
+    results = dict()
+
+    train_data, evi_data, mar_data = load_benchmark_data(dataset, seed=42)
+    n_features = train_data.shape[1]
+
+    gc.collect()
+    gc.disable()
+
+    print("Benchmarking SPNs on SPFlow ...")
+    root = learn_binary_spn(train_data)
+
+    print("\tBenchmarking EVI Inference ...")
+    elapsed_times = list()
+    for _ in range(n_repetitions):
+        gc.collect()
+        start_time = time.perf_counter()
+        log_likelihood(root, evi_data, dtype=np.float32)
+        end_time = time.perf_counter()
+        elapsed_times.append(end_time - start_time)
+    results['evi'] = compute_mean_stddev_times(elapsed_times)
+
+    print("\tBenchmarking MAR Inference ...")
+    elapsed_times = list()
+    for _ in range(n_repetitions):
+        gc.collect()
+        start_time = time.perf_counter()
+        log_likelihood(root, mar_data, dtype=np.float32)
+        end_time = time.perf_counter()
+        elapsed_times.append(end_time - start_time)
+    results['mar'] = compute_mean_stddev_times(elapsed_times)
+
+    print("\tBenchmarking MPE Inference ...")
+    elapsed_times = list()
+    for _ in range(n_repetitions):
+        gc.collect()
+        start_time = time.perf_counter()
+        mpe(root, mar_data, in_place=False)
+        end_time = time.perf_counter()
+        elapsed_times.append(end_time - start_time)
+    results['mpe'] = compute_mean_stddev_times(elapsed_times)
+
+    print("\tBenchmarking Conditional Sampling ...")
+    rand_gen = np.random.RandomState(42)
+    elapsed_times = list()
+    for _ in range(n_repetitions):
+        gc.collect()
+        start_time = time.perf_counter()
+        sample_instances(root, mar_data, rand_gen=rand_gen, in_place=False)
+        end_time = time.perf_counter()
+        elapsed_times.append(end_time - start_time)
+    results['cond_sampling'] = compute_mean_stddev_times(elapsed_times)
+
+    # Save the benchmark results to file
+    results_filepath = 'spflow-spn-benchmark.json'
+    with open(results_filepath, 'w') as f:
+        results = json.loads(json.dumps(results), parse_float=lambda x: round(float(x), 2))
+        json.dump({'spflow': {'spn': results}}, f, indent=2)
+    print("Saved benchmark results to {}".format(results_filepath))
+    print()
diff --git a/benchmark/settings.json b/benchmark/settings.json
new file mode 100644
index 0000000..fae8e1f
--- /dev/null
+++ b/benchmark/settings.json
@@ -0,0 +1,5 @@
+{
+  "dataset": "binarized_mnist",
+  "n_repetitions": 10,
+  "n_samples": 1000
+}
diff --git a/benchmark/spn_queries.py b/benchmark/spn_queries.py
deleted file mode 100644
index 981400f..0000000
--- a/benchmark/spn_queries.py
+++ /dev/null
@@ -1,192 +0,0 @@
-import gc
-import time
-import numpy as np
-
-from typing import Union, List
-from utils import compute_mean_stddev_times, load_dataset
-
-from spn.algorithms.Inference import log_likelihood
-from spn.algorithms.MPE import mpe
-from spn.algorithms.Sampling import sample_instances
-from spn.structure.Base import Node, Sum, Product, assign_ids, rebuild_scopes_bottom_up
-from spn.structure.leaves.parametric.Parametric import Bernoulli
-
-from deeprob.context import ContextState
-import deeprob.spn.structure as spn
-import deeprob.spn.algorithms as spnalg
-import deeprob.spn.utils as spnutils
-
-
-def learn_binary_clt(data: np.ndarray) -> spn.BinaryCLT:
-    n_features = data.shape[1]
-    scope = list(range(n_features))
-    domain = [[0, 1]] * n_features
-    clt = spn.BinaryCLT(scope, root=0)
-    clt.fit(data, domain, alpha=0.1, random_state=42)
-    return clt
-
-
-def deeprob_build_spn(clt: spn.BinaryCLT) -> spn.Node:
-    return clt.to_pc()
-
-
-def spflow_build_spn(clt: spn.BinaryCLT) -> Node:  # This function is copy-pasted from BinaryCLT.to_pc()
-    from deeprob.utils.graph import build_tree_structure
-    root = build_tree_structure(clt.tree, scope=clt.scope)
-    factors = {clt.scope[i]: np.exp(clt.params[i]) for i in range(len(clt.tree))}
-    neg_buffer, pos_buffer = [], []
-    nodes_stack = [root]
-    last_node_visited = None
-    while nodes_stack:
-        node = nodes_stack[-1]
-        if node.is_leaf() or (last_node_visited in node.get_children()):
-            leaves: List[Union[Bernoulli, Sum]] = [
-                Bernoulli(p=0.0, scope=node.get_id()),
-                Bernoulli(p=1.0, scope=node.get_id()),
-            ]
-            if not node.is_leaf():
-                neg_prod = Product(children=[leaves[0]] + neg_buffer[-len(node.get_children()):])
-                pos_prod = Product(children=[leaves[1]] + pos_buffer[-len(node.get_children()):])
-                del neg_buffer[-len(node.get_children()):]
-                del pos_buffer[-len(node.get_children()):]
-                sum_children = [neg_prod, pos_prod]
-            else:
-                sum_children = leaves
-            weights = factors[node.get_id()]
-            neg_buffer.append(Sum(children=sum_children, weights=weights[0]))
-            pos_buffer.append(Sum(children=sum_children, weights=weights[1]))
-            last_node_visited = nodes_stack.pop()
-        else:
-            nodes_stack.extend(node.get_children())
-    return rebuild_scopes_bottom_up(assign_ids(pos_buffer[0]))
-
-
-if __name__ == '__main__':
-    gc.disable()
-    n_repetitions = 50
-    data_train, evi_data, mar_data = load_dataset('ad', n_samples=5000, seed=42)
-    clt = learn_binary_clt(data_train)
-
-    root = deeprob_build_spn(clt)
-    print()
-    print("CLT -> SPN compilation statistics:")
-    print(spnutils.compute_statistics(root))
-
-    print()
-    print("Benchmark of SPNs on DeeProb-kit ...")
-
-    # ---- DeeProb-kit EVI Inference ----------------------------------------------------------------------------------
-    elapsed_times = list()
-    with ContextState(check_spn=False):
-        for _ in range(n_repetitions):
-            gc.collect()
-            start_time = time.perf_counter()
-            lls = spnalg.log_likelihood(root, evi_data)
-            end_time = time.perf_counter()
-            elapsed_times.append(1e3 * (end_time - start_time))
-    mean_time, stddev_time = compute_mean_stddev_times(elapsed_times)
-    print("[EVI] DeeProb-kit Avg. Time: {:.2f}ms (+- {:.2f}) - LL: {:.6f}".format(
-        mean_time, stddev_time, np.mean(lls)
-    ))
-
-    # ---- DeeProb-kit MAR Inference ----------------------------------------------------------------------------------
-    elapsed_times = list()
-    with ContextState(check_spn=False):
-        for _ in range(n_repetitions):
-            gc.collect()
-            start_time = time.perf_counter()
-            lls = spnalg.log_likelihood(root, mar_data)
-            end_time = time.perf_counter()
-            elapsed_times.append(1e3 * (end_time - start_time))
-    mean_time, stddev_time = compute_mean_stddev_times(elapsed_times)
-    print("[MAR] DeeProb-kit Avg. Time: {:.2f}ms (+- {:.2f}) - LL: {:.6f}".format(
-        mean_time, stddev_time, np.mean(lls)
-    ))
-
-    # ---- DeeProb-kit MPE Inference ----------------------------------------------------------------------------------
-    elapsed_times = list()
-    with ContextState(check_spn=False):
-        for _ in range(n_repetitions):
-            gc.collect()
-            start_time = time.perf_counter()
-            mpe_data = spnalg.mpe(root, mar_data, inplace=False)
-            end_time = time.perf_counter()
-            elapsed_times.append(1e3 * (end_time - start_time))
-    lls = spnalg.log_likelihood(root, mpe_data)
-    mean_time, stddev_time = compute_mean_stddev_times(elapsed_times)
-    print("[MPE] DeeProb-kit Avg. Time: {:.2f}ms (+- {:.2f}) - LL: {:.6f}".format(
-        mean_time, stddev_time, np.mean(lls)
-    ))
-
-    # ---- DeeProb-kit Sampling ---------------------------------------------------------------------------------------
-    elapsed_times = list()
-    with ContextState(check_spn=False):
-        for _ in range(n_repetitions):
-            gc.collect()
-            start_time = time.perf_counter()
-            _ = spnalg.sample(root, mar_data, inplace=False)
-            end_time = time.perf_counter()
-            elapsed_times.append(1e3 * (end_time - start_time))
-    mean_time, stddev_time = compute_mean_stddev_times(elapsed_times)
-    print("[Sampling] DeeProb-kit Avg. Time: {:.2f}ms (+- {:.2f})".format(
-        mean_time, stddev_time
-    ))
-
-    print()
-    print("Benchmark of SPNs on SPFlow ...")
-
-    root = spflow_build_spn(clt)
-
-    # ---- SPFlow EVI Inference ---------------------------------------------------------------------------------------
-    elapsed_times = list()
-    for _ in range(n_repetitions):
-        gc.collect()
-        start_time = time.perf_counter()
-        lls = log_likelihood(root, evi_data, dtype=np.float32)
-        end_time = time.perf_counter()
-        elapsed_times.append(1e3 * (end_time - start_time))
-    mean_time, stddev_time = compute_mean_stddev_times(elapsed_times)
-    print("[EVI] SPFlow Avg. Time: {:.2f}ms (+- {:.2f}) - LL: {:.6f}".format(
-        mean_time, stddev_time, np.mean(lls)
-    ))
-
-    # ---- SPFlow MAR Inference ---------------------------------------------------------------------------------------
-    elapsed_times = list()
-    for _ in range(n_repetitions):
-        gc.collect()
-        start_time = time.perf_counter()
-        lls = log_likelihood(root, mar_data, dtype=np.float32)
-        end_time = time.perf_counter()
-        elapsed_times.append(1e3 * (end_time - start_time))
-    mean_time, stddev_time = compute_mean_stddev_times(elapsed_times)
-    print("[MAR] SPFlow Avg. Time: {:.2f}ms (+- {:.2f}) - LL: {:.6f}".format(
-        mean_time, stddev_time, np.mean(lls)
-    ))
-
-    # ---- SPFlow MPE Inference ---------------------------------------------------------------------------------------
-    elapsed_times = list()
-    for _ in range(n_repetitions):
-        gc.collect()
-        start_time = time.perf_counter()
-        mpe_data = mpe(root, mar_data, in_place=False)
-        end_time = time.perf_counter()
-        elapsed_times.append(1e3 * (end_time - start_time))
-    lls = log_likelihood(root, mpe_data, dtype=np.float32)
-    mean_time, stddev_time = compute_mean_stddev_times(elapsed_times)
-    print("[MPE] SPFlow Avg. Time: {:.2f}ms (+- {:.2f}) - LL: {:.6f}".format(
-        mean_time, stddev_time, np.mean(lls)
-    ))
-
-    # ---- SPFlow Sampling --------------------------------------------------------------------------------------------
-    rand_gen = np.random.RandomState(42)
-    elapsed_times = list()
-    for _ in range(n_repetitions):
-        gc.collect()
-        start_time = time.perf_counter()
-        _ = sample_instances(root, mar_data, rand_gen=rand_gen, in_place=False)
-        end_time = time.perf_counter()
-        elapsed_times.append(1e3 * (end_time - start_time))
-    mean_time, stddev_time = compute_mean_stddev_times(elapsed_times)
-    print("[Sampling] SPFlow Avg. Time: {:.2f}ms (+- {:.2f})".format(
-        mean_time, stddev_time
-    ))
diff --git a/benchmark/utils.py b/benchmark/utils.py
index b93b9ef..3b6b80b 100644
--- a/benchmark/utils.py
+++ b/benchmark/utils.py
@@ -1,6 +1,7 @@
+from typing import Tuple, List
+
 import numpy as np
 
-from typing import Tuple, List
 from experiments.datasets import load_binary_dataset
 
 
@@ -10,18 +11,16 @@ def compute_mean_stddev_times(elapsed_times: List[float]) -> Tuple[float, float]
     return mean_time, stddev_time
 
 
-def load_dataset(
+def load_benchmark_data(
     name: str,
     mar_p: float = 0.5,
-    n_samples: int = 1000,
     seed: int = 42
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
     assert 0.0 < mar_p < 1.0
-    random_state = np.random.RandomState(seed)
     data_train, _, data_test = load_binary_dataset('experiments/datasets', name, raw=True)
-    indices = random_state.choice(len(data_test), size=n_samples, replace=True)  # Resample the test set
-    evi_data = data_test[indices]
-    data_train, evi_data = data_train.astype(np.float32), evi_data.astype(np.float32)
-    mar_data = evi_data.copy()
+    data_train, data_test = data_train.astype(np.float32), data_test.astype(np.float32)
+    mar_data = data_test.copy()
+    random_state = np.random.RandomState(seed)
     mar_data[random_state.rand(*mar_data.shape) <= mar_p] = np.nan
-    return data_train, evi_data, mar_data
+    data_test = data_test.astype(np.int64)
+    return data_train, data_test, mar_data
diff --git a/deeprob/spn/algorithms/evaluation.py b/deeprob/spn/algorithms/evaluation.py
index c690254..05f1677 100644
--- a/deeprob/spn/algorithms/evaluation.py
+++ b/deeprob/spn/algorithms/evaluation.py
@@ -1,9 +1,10 @@
 from typing import Optional, Union, Tuple, Any, Callable
 
+import joblib
 import numpy as np
 
 from deeprob.spn.structure.leaf import Leaf
-from deeprob.spn.structure.node import Node, Sum, Product, topological_order
+from deeprob.spn.structure.node import Node, Sum, Product, topological_order_layered
 from deeprob.spn.utils.validity import check_spn
 
 
@@ -15,6 +16,7 @@ def eval_bottom_up(
     leaf_func_kwargs: Optional[dict] = None,
     node_func_kwargs: Optional[dict] = None,
     return_results: bool = False,
+    n_jobs: int = -1
 ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
     """
     Evaluate the SPN bottom up given some inputs and leaves and nodes evaluation functions.
@@ -26,31 +28,38 @@ def eval_bottom_up(
     :param leaf_func_kwargs: The optional parameters of the leaf evaluation function.
     :param node_func_kwargs: The optional parameters of the inner nodes evaluation function.
     :param return_results: A flag indicating if this function must return the log likelihoods of each node of the SPN.
+    :param n_jobs: The number of parallel jobs. It follows the joblib's convention.
     :return: The outputs. Additionally it returns the output of each node.
     :raises ValueError: If a parameter is out of domain.
     """
-    # Check the SPN
-    check_spn(root, labeled=True, smooth=True, decomposable=True)
-
-    nodes = topological_order(root)
-    if nodes is None:
-        raise ValueError("SPN structure is not a directed acyclic graph (DAG)")
-
     if leaf_func_kwargs is None:
         leaf_func_kwargs = dict()
     if node_func_kwargs is None:
         node_func_kwargs = dict()
 
-    n_samples, _ = x.shape
-    ls = np.empty(shape=(len(nodes), n_samples), dtype=np.float32)
+    # Check the SPN
+    check_spn(root, labeled=True, smooth=True, decomposable=True)
+
+    # Compute the layered topological ordering
+    layers = topological_order_layered(root)
+    if layers is None:
+        raise ValueError("SPN structure is not a directed acyclic graph (DAG)")
 
-    for node in reversed(nodes):
+    n_nodes, n_samples = sum(map(len, layers)), len(x)
+    ls = np.empty(shape=(n_nodes, n_samples), dtype=np.float32)
+
+    def eval_forward(node):
         if isinstance(node, Leaf):
             ls[node.id] = leaf_func(node, x[:, node.scope], **leaf_func_kwargs)
         else:
             children_ls = np.stack([ls[c.id] for c in node.children], axis=1)
             ls[node.id] = node_func(node, children_ls, **node_func_kwargs)
 
+    with joblib.parallel_backend('threading', n_jobs=n_jobs):
+        with joblib.Parallel() as parallel:
+            for layer in reversed(layers):
+                parallel(joblib.delayed(eval_forward)(node) for node in layer)
+
     if return_results:
         return ls[root.id], ls
     return ls[root.id]
@@ -64,7 +73,8 @@ def eval_top_down(
     sum_func: Callable[[Sum, np.ndarray, Any], np.ndarray],
     leaf_func_kwargs: Optional[dict] = None,
     sum_func_kwargs: Optional[dict] = None,
-    inplace: bool = False
+    inplace: bool = False,
+    n_jobs: int = -1
 ) -> np.ndarray:
     """
     Evaluate the SPN top down given some inputs, the likelihoods of each node and a leaves evaluation function.
@@ -78,30 +88,33 @@ def eval_top_down(
     :param leaf_func_kwargs: The optional parameters of the leaf evaluation function.
     :param sum_func_kwargs: The optional parameters of the sum nodes evaluation function.
     :param inplace: Whether to make inplace assignments.
+    :param n_jobs: The number of parallel jobs. It follows the joblib's convention.
     :return: The NaN-filled inputs.
     :raises ValueError: If a parameter is out of domain.
     """
-    # Check the SPN
-    check_spn(root, labeled=True, smooth=True, decomposable=True)
-
-    nodes = topological_order(root)
-    if nodes is None:
-        raise ValueError("SPN structure is not a directed acyclic graph (DAG)")
-
     if leaf_func_kwargs is None:
         leaf_func_kwargs = dict()
     if sum_func_kwargs is None:
         sum_func_kwargs = dict()
 
+    # Check the SPN
+    check_spn(root, labeled=True, smooth=True, decomposable=True)
+
+    # Compute the layered topological ordering
+    layers = topological_order_layered(root)
+    if layers is None:
+        raise ValueError("SPN structure is not a directed acyclic graph (DAG)")
+
+    # Copy the input array, if not inplace mode
     if not inplace:
         x = np.copy(x)
-    n_samples, _ = x.shape
 
     # Build the array consisting of top-down path masks
-    masks = np.zeros(shape=(len(nodes), n_samples), dtype=np.bool_)
+    n_nodes, n_samples = sum(map(len, layers)), len(x)
+    masks = np.zeros(shape=(n_nodes, n_samples), dtype=np.bool_)
     masks[root.id] = True
 
-    for node in nodes:
+    def eval_backward(node):
         if isinstance(node, Leaf):
             mask = np.ix_(masks[node.id], node.scope)
             x[mask] = leaf_func(node, x[mask], **leaf_func_kwargs)
@@ -118,4 +131,9 @@ def eval_top_down(
                 "Top down evaluation not implemented for node of type {}".format(node.__class__.__name__)
             )
 
+    with joblib.parallel_backend('threading', n_jobs=n_jobs):
+        with joblib.Parallel() as parallel:
+            for layer in layers:
+                parallel(joblib.delayed(eval_backward)(node) for node in layer)
+
     return x
diff --git a/deeprob/spn/structure/node.py b/deeprob/spn/structure/node.py
index 0240864..b0ee593 100644
--- a/deeprob/spn/structure/node.py
+++ b/deeprob/spn/structure/node.py
@@ -228,7 +228,7 @@ def topological_order(root: Node) -> Optional[List[Node]]:
     if num_outgoings[root] != 0:
         return None
 
-    # Initialize the queue by the nodes having no outgoing edges, i.e. the root node
+    # Non-layered topological ordering implementation
     queue = deque([root])
     while queue:
         node = queue.popleft()
@@ -242,3 +242,43 @@ def topological_order(root: Node) -> Optional[List[Node]]:
     if sum(num_outgoings.values()) != 0:
         return None
     return ordering
+
+
+def topological_order_layered(root: Node) -> Optional[List[List[Node]]]:
+    """
+    Compute the Topological Ordering Layered for a SPN, using the Kahn's Algorithm.
+
+    :param root: The root of the SPN.
+    :return: A list of layers that form a topological ordering.
+             If the SPN graph is not acyclic, it returns None.
+    """
+    ordering = list()
+    num_outgoings = defaultdict(int)
+    num_outgoings[root] = 0
+
+    # Initialize the number of outgoings edges for each node
+    for node in bfs(root):
+        for c in node.children:
+            num_outgoings[c] += 1
+
+    # Check the unusual case where the root node have outgoings edges, i.e. a trivial cycle has been found
+    if num_outgoings[root] != 0:
+        return None
+
+    # Layered topological ordering implementation
+    ordering.append([root])
+    while True:
+        layer = list()
+        for node in ordering[-1]:
+            for c in node.children:
+                num_outgoings[c] -= 1
+                if num_outgoings[c] == 0:
+                    layer.append(c)
+        if not layer:
+            break
+        ordering.append(layer)
+
+    # Check if a cycle has been found
+    if sum(num_outgoings.values()) != 0:
+        return None
+    return ordering
diff --git a/requirements.txt b/requirements.txt
index a415ad0..c21cc7c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
+joblib
 numpy
 scipy
 tqdm


Test patch:
diff --git a/test/test_clt.py b/test/test_clt.py
index 51d0b73..0ae2254 100644
--- a/test/test_clt.py
+++ b/test/test_clt.py
@@ -1,7 +1,7 @@
 import unittest
 import tempfile
 
-from experiments.datasets import load_binary_dataset
+from sklearn.datasets import load_diabetes
 from test.utils import *
 
 from deeprob.spn.utils.validity import is_structured_decomposable
@@ -17,11 +17,11 @@ def __init__(self, *args, **kwargs):
     @classmethod
     def setUpClass(cls):
         random_state = np.random.RandomState(42)
-        data, _, _ = load_binary_dataset('experiments/datasets', 'nltcs', raw=True)
-        data = data.astype(np.float32)
+        data, _, = load_diabetes(return_X_y=True)
+        data = (data < np.median(data, axis=0)).astype(np.float32)
         cls.root_id = 1
         cls.n_samples, cls.n_features = data.shape
-        cls.evi_data = resample_data(data, 5000, random_state)
+        cls.evi_data = resample_data(data, 1000, random_state)
         cls.mar_data = random_marginalize_data(cls.evi_data, 0.2, random_state)
 
         cls.complete_data = complete_binary_data(cls.n_features)
diff --git a/test/test_em.py b/test/test_em.py
index cffd812..9ac6338 100644
--- a/test/test_em.py
+++ b/test/test_em.py
@@ -2,7 +2,7 @@
 import numpy as np
 
 from sklearn.datasets import make_blobs
-from experiments.datasets import load_binary_dataset
+from sklearn.datasets import load_diabetes
 from test.utils import resample_data
 
 from deeprob.spn.structure.node import Sum, Product
@@ -19,10 +19,10 @@ def __init__(self, *args, **kwargs):
     @classmethod
     def setUpClass(cls):
         random_state = np.random.RandomState(42)
-        data, _, _ = load_binary_dataset('experiments/datasets', 'nltcs', raw=True)
-        data = data.astype(np.float32)
+        data, _, = load_diabetes(return_X_y=True)
+        data = (data < np.median(data, axis=0)).astype(np.float32)
         cls.n_samples, cls.n_features = data.shape
-        cls.evi_data = resample_data(data, 5000, random_state)
+        cls.evi_data = resample_data(data, 1000, random_state)
         cls.blobs_data, _ = make_blobs(
             n_samples=1000, n_features=2, random_state=1337,
             centers=[[-1.0, 1.0], [1.0, -1.0]], cluster_std=0.25
@@ -42,14 +42,14 @@ def __build_normal_spn():
 
     def __learn_spn_mle(self):
         return learn_estimator(
-            self.evi_data, [Bernoulli] * self.n_features,
+            self.evi_data, [Bernoulli] * self.n_features, [[0, 1]] * self.n_features,
             learn_leaf='mle', split_rows='gmm', split_cols='gvs', min_rows_slice=512,
             random_state=42, verbose=False
         )
 
     def __learn_spn_clt(self):
         return learn_estimator(
-            self.evi_data, [Bernoulli] * self.n_features,
+            self.evi_data, [Bernoulli] * self.n_features, [[0, 1]] * self.n_features,
             learn_leaf='binary-clt', split_rows='kmeans', split_cols='gvs', min_rows_slice=512,
             learn_leaf_kwargs={'to_pc': False},
             random_state=42, verbose=False
@@ -62,7 +62,7 @@ def test_spn_binary(self):
             random_init=False, random_state=42, verbose=False
         )
         ll = log_likelihood(spn, self.evi_data).mean()
-        self.assertAlmostEqual(ll, -6.0, places=1)
+        self.assertAlmostEqual(ll, -5.3, places=1)
 
     def test_clt_binary(self):
         spn = self.__learn_spn_clt()
@@ -71,7 +71,7 @@ def test_clt_binary(self):
             random_init=True, random_state=42, verbose=False
         )
         ll = log_likelihood(spn, self.evi_data).mean()
-        self.assertAlmostEqual(ll, -5.9, places=1)
+        self.assertAlmostEqual(ll, -5.1, places=1)
 
     def test_spn_gaussian(self):
         spn = self.__build_normal_spn()
diff --git a/test/test_spn.py b/test/test_spn.py
index dec7f9f..c911f45 100644
--- a/test/test_spn.py
+++ b/test/test_spn.py
@@ -1,14 +1,15 @@
 import unittest
 import tempfile
 
-from experiments.datasets import load_binary_dataset
+from sklearn.datasets import load_diabetes
 from test.utils import *
 
 from deeprob.spn.utils.statistics import compute_statistics
 from deeprob.spn.utils.filter import filter_nodes_by_type
 from deeprob.spn.utils.validity import check_spn
 from deeprob.spn.structure.node import Sum, Product
-from deeprob.spn.structure.node import bfs, dfs_post_order, topological_order
+from deeprob.spn.structure.node import bfs, dfs_post_order, topological_order, topological_order_layered
+from deeprob.spn.structure.cltree import BinaryCLT
 from deeprob.spn.structure.leaf import Bernoulli, Gaussian
 from deeprob.spn.structure.io import save_spn_json, load_spn_json
 from deeprob.spn.learning.learnspn import learn_spn
@@ -25,15 +26,15 @@ def __init__(self, *args, **kwargs):
     @classmethod
     def setUpClass(cls):
         random_state = np.random.RandomState(42)
-        data, _, _ = load_binary_dataset('experiments/datasets', 'nltcs', raw=True)
-        data = data.astype(np.float32)
+        data, _, = load_diabetes(return_X_y=True)
+        data = (data < np.median(data, axis=0)).astype(np.float32)
         cls.n_samples, cls.n_features = data.shape
-        cls.evi_data = resample_data(data, 5000, random_state)
+        cls.evi_data = resample_data(data, 1000, random_state)
         cls.mar_data = random_marginalize_data(cls.evi_data, 0.2, random_state)
 
-        cls.clf_index = 3
+        cls.clf_index = 2
         cls.clf_data = marginalize_data(cls.evi_data, [cls.clf_index])
-        cls.scope = [5, 7, 9, 15, 8]
+        cls.scope = [5, 9, 8]
         mar_scope = [s for s in range(cls.n_features) if s not in cls.scope]
         cls.scope_mar_data = marginalize_data(cls.evi_data, mar_scope)
 
@@ -43,6 +44,9 @@ def setUpClass(cls):
         ], axis=1)
 
         cls.complete_data = complete_binary_data(cls.n_features)
+        mar_features = [1, 2, 3, 5, 8]
+        cls.complete_mar_data = complete_marginalized_binary_data(cls.n_features, mar_features)
+        cls.complete_mpe_data = complete_posterior_binary_data(cls.n_features, mar_features)
 
     @staticmethod
     def __build_normal_spn():
@@ -80,32 +84,38 @@ def __build_cyclical_spn():
         b0a.id, b1a.id, b0b.id, b1b.id = 3, 4, 5, 6
         return s0
 
+    def __learn_binary_clt(self):
+        scope = list(range(self.n_features))
+        clt = BinaryCLT(scope, root=0)
+        clt.fit(self.evi_data, [[0, 1]] * self.n_features, alpha=0.1, random_state=42)
+        return clt
+
     def __learn_spn_unpruned(self):
         return learn_spn(
             self.evi_data, [Bernoulli] * self.n_features, [[0, 1]] * self.n_features,
-            learn_leaf='mle', split_cols='gvs', min_rows_slice=512,
+            learn_leaf='mle', split_cols='gvs', min_rows_slice=64,
             random_state=42, verbose=False
         )
 
     def __learn_spn_mle(self):
         return learn_estimator(
-            self.evi_data, [Bernoulli] * self.n_features,
-            learn_leaf='mle', split_rows='gmm', split_cols='gvs', min_rows_slice=512,
+            self.evi_data, [Bernoulli] * self.n_features, [[0, 1]] * self.n_features,
+            learn_leaf='mle', split_rows='gmm', split_cols='gvs', min_rows_slice=64,
             random_state=42, verbose=False
         )
 
     def __learn_spn_clt(self):
         return learn_estimator(
-            self.evi_data, [Bernoulli] * self.n_features,
-            learn_leaf='binary-clt', split_rows='kmeans', split_cols='gvs', min_rows_slice=512,
+            self.evi_data, [Bernoulli] * self.n_features, [[0, 1]] * self.n_features,
+            learn_leaf='binary-clt', split_rows='kmeans', split_cols='gvs', min_rows_slice=64,
             learn_leaf_kwargs={'to_pc': False},
             random_state=42, verbose=False
         )
 
     def __learn_spn_mle_classifier(self):
         return learn_classifier(
-            self.evi_data, [Bernoulli] * self.n_features, class_idx=self.clf_index,
-            learn_leaf='binary-clt', split_cols='rdc', min_rows_slice=512, learn_leaf_kwargs={'to_pc': True},
+            self.evi_data, [Bernoulli] * self.n_features, [[0, 1]] * self.n_features, class_idx=self.clf_index,
+            learn_leaf='binary-clt', split_cols='rdc', min_rows_slice=64, learn_leaf_kwargs={'to_pc': True},
             random_state=42, verbose=False
         )
 
@@ -171,12 +181,20 @@ def test_mpe_inference(self):
         self.assertFalse(np.any(np.isnan(mpe_data)))
         self.assertGreater(mpe_ll, evi_ll)
 
+    def test_mpe_complete_inference(self):
+        spn = self.__learn_binary_clt().to_pc()
+        complete_lls = log_likelihood(spn, self.complete_data)
+        mpe_data = mpe(spn, self.complete_mar_data)
+        mpe_ids = binary_data_ids(mpe_data).tolist()
+        expected_mpe_ids = compute_mpe_ids(self.complete_mpe_data, complete_lls.squeeze())
+        self.assertEqual(mpe_ids, expected_mpe_ids)
+
     def test_classifier(self):
         spn = self.__learn_spn_mle_classifier()
         clf_data = mpe(spn, self.clf_data)
         error_rate = np.mean(np.abs(clf_data[:, self.clf_index] - self.evi_data[:, self.clf_index]))
         self.assertFalse(np.any(np.isnan(clf_data)))
-        self.assertGreater(1.0 - error_rate, 0.8)
+        self.assertGreater(1.0 - error_rate, 0.7)
 
     def test_bfs(self):
         spn = self.__build_dag_spn()
@@ -197,6 +215,15 @@ def test_topological_order(self):
         ordering = topological_order(spn)
         self.assertIsNone(ordering)
 
+    def test_topological_order_layered(self):
+        spn = self.__build_dag_spn()
+        layers = topological_order_layered(spn)
+        node_layered_ids = list(map(lambda layer: list(map(lambda node: node.id, layer)), layers))
+        self.assertEqual(node_layered_ids, [[0], [1, 2, 3], [5, 6, 4, 7]])
+        spn = self.__build_cyclical_spn()
+        layers = topological_order_layered(spn)
+        self.assertIsNone(layers)
+
     def test_prune(self):
         spn = self.__learn_spn_unpruned()
         ll = log_likelihood(spn, self.evi_data).mean()