Repository: pgmpy/pgmpy

Issue:
XML Belief Network writer class hardcoded for variables with cardinality 2
The current implementation seems to be hardcoded to work only for the CPDs of variables having a cardinality of 2. The lines at: https://github.com/pgmpy/pgmpy/blob/dev/pgmpy/readwrite/XMLBeliefNetwork.py#L451. The `range` here makes the assumption of cardinality of 2.

Discussion:
ankurankan: - [ ] Add the evidence in sorted order for tests to pass.

codecov[bot]: # [Codecov](https://codecov.io/gh/pgmpy/pgmpy/pull/1215?src=pr&el=h1) Report
> Merging [#1215](https://codecov.io/gh/pgmpy/pgmpy/pull/1215?src=pr&el=desc) into [dev](https://codecov.io/gh/pgmpy/pgmpy/commit/a3a544694f55105208ae5edc1decd64be9ea913a?src=pr&el=desc) will **increase** coverage by `0.01%`.
> The diff coverage is `95.23%`.

[![Impacted file tree graph](https://codecov.io/gh/pgmpy/pgmpy/pull/1215/graphs/tree.svg?width=650&token=UaJMCdHaEF&height=150&src=pr)](https://codecov.io/gh/pgmpy/pgmpy/pull/1215?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##              dev    #1215      +/-   ##
==========================================
+ Coverage   92.02%   92.03%   +0.01%     
==========================================
  Files         135      135              
  Lines       13190    13198       +8     
==========================================
+ Hits        12138    12147       +9     
+ Misses       1052     1051       -1
```


| [Impacted Files](https://codecov.io/gh/pgmpy/pgmpy/pull/1215?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [pgmpy/tests/test\_utils/test\_state\_name.py](https://codecov.io/gh/pgmpy/pgmpy/pull/1215/diff?src=pr&el=tree#diff-cGdtcHkvdGVzdHMvdGVzdF91dGlscy90ZXN0X3N0YXRlX25hbWUucHk=) | `100% <100%> (ø)` | :arrow_up: |
| [...py/tests/test\_factors/test\_discrete/test\_Factor.py](https://codecov.io/gh/pgmpy/pgmpy/pull/1215/diff?src=pr&el=tree#diff-cGdtcHkvdGVzdHMvdGVzdF9mYWN0b3JzL3Rlc3RfZGlzY3JldGUvdGVzdF9GYWN0b3IucHk=) | `99.82% <100%> (ø)` | :arrow_up: |
| [pgmpy/tests/test\_inference/test\_Inference.py](https://codecov.io/gh/pgmpy/pgmpy/pull/1215/diff?src=pr&el=tree#diff-cGdtcHkvdGVzdHMvdGVzdF9pbmZlcmVuY2UvdGVzdF9JbmZlcmVuY2UucHk=) | `100% <100%> (ø)` | :arrow_up: |
| [pgmpy/tests/test\_models/test\_BayesianModel.py](https://codecov.io/gh/pgmpy/pgmpy/pull/1215/diff?src=pr&el=tree#diff-cGdtcHkvdGVzdHMvdGVzdF9tb2RlbHMvdGVzdF9CYXllc2lhbk1vZGVsLnB5) | `100% <100%> (ø)` | :arrow_up: |
| [pgmpy/readwrite/XMLBeliefNetwork.py](https://codecov.io/gh/pgmpy/pgmpy/pull/1215/diff?src=pr&el=tree#diff-cGdtcHkvcmVhZHdyaXRlL1hNTEJlbGllZk5ldHdvcmsucHk=) | `92.75% <100%> (+1.51%)` | :arrow_up: |
| [pgmpy/tests/test\_inference/test\_dbn\_inference.py](https://codecov.io/gh/pgmpy/pgmpy/pull/1215/diff?src=pr&el=tree#diff-cGdtcHkvdGVzdHMvdGVzdF9pbmZlcmVuY2UvdGVzdF9kYm5faW5mZXJlbmNlLnB5) | `100% <100%> (ø)` | :arrow_up: |
| [...gmpy/tests/test\_readwrite/test\_XMLBeliefNetwork.py](https://codecov.io/gh/pgmpy/pgmpy/pull/1215/diff?src=pr&el=tree#diff-cGdtcHkvdGVzdHMvdGVzdF9yZWFkd3JpdGUvdGVzdF9YTUxCZWxpZWZOZXR3b3JrLnB5) | `94.16% <100%> (+0.09%)` | :arrow_up: |
| [...y/tests/test\_models/test\_DynamicBayesianNetwork.py](https://codecov.io/gh/pgmpy/pgmpy/pull/1215/diff?src=pr&el=tree#diff-cGdtcHkvdGVzdHMvdGVzdF9tb2RlbHMvdGVzdF9EeW5hbWljQmF5ZXNpYW5OZXR3b3JrLnB5) | `100% <100%> (ø)` | :arrow_up: |
| [pgmpy/tests/test\_readwrite/test\_XMLBIF.py](https://codecov.io/gh/pgmpy/pgmpy/pull/1215/diff?src=pr&el=tree#diff-cGdtcHkvdGVzdHMvdGVzdF9yZWFkd3JpdGUvdGVzdF9YTUxCSUYucHk=) | `88.7% <100%> (ø)` | :arrow_up: |
| [pgmpy/factors/discrete/CPD.py](https://codecov.io/gh/pgmpy/pgmpy/pull/1215/diff?src=pr&el=tree#diff-cGdtcHkvZmFjdG9ycy9kaXNjcmV0ZS9DUEQucHk=) | `78.99% <66.66%> (+0.04%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pgmpy/pgmpy/pull/1215?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/pgmpy/pgmpy/pull/1215?src=pr&el=footer). Last update [a3a5446...b495921](https://codecov.io/gh/pgmpy/pgmpy/pull/1215?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/pgmpy/factors/discrete/CPD.py b/pgmpy/factors/discrete/CPD.py
index 3f51d365..5c1de436 100644
--- a/pgmpy/factors/discrete/CPD.py
+++ b/pgmpy/factors/discrete/CPD.py
@@ -125,6 +125,15 @@ class TabularCPD(DiscreteFactor):
         if values.ndim != 2:
             raise TypeError("Values must be a 2D list/array")
 
+        if evidence is None:
+            expected_cpd_shape = (variable_card, 1)
+        else:
+            expected_cpd_shape = (variable_card, np.product(evidence_card))
+        if values.shape != expected_cpd_shape:
+            raise ValueError(
+                f"values must be of shape {expected_cpd_shape}. Got shape: {values.shape}"
+            )
+
         super(TabularCPD, self).__init__(
             variables, cardinality, values.flatten("C"), state_names=state_names
         )
@@ -169,7 +178,7 @@ class TabularCPD(DiscreteFactor):
                 self.cardinality[0], np.prod(self.cardinality[1:])
             )
         else:
-            return self.values.reshape(1, np.prod(self.cardinality))
+            return self.values.reshape(np.prod(self.cardinality), 1)
 
     def __str__(self):
         return self._make_table_str(tablefmt="grid")
diff --git a/pgmpy/readwrite/XMLBeliefNetwork.py b/pgmpy/readwrite/XMLBeliefNetwork.py
index 25af415b..af457207 100644
--- a/pgmpy/readwrite/XMLBeliefNetwork.py
+++ b/pgmpy/readwrite/XMLBeliefNetwork.py
@@ -1,3 +1,5 @@
+import itertools
+
 import numpy as np
 import networkx as nx
 
@@ -200,7 +202,7 @@ class XBNReader(object):
                 )
             distribution[variable_name]["DPIS"] = np.array(
                 [list(map(float, dpi.text.split())) for dpi in dist.find("DPIS")]
-            )
+            ).transpose()
 
         return distribution
 
@@ -225,13 +227,8 @@ class XBNReader(object):
             tabular_cpds.append(cpd)
 
         model.add_cpds(*tabular_cpds)
-
-        if nx.__version__.startswith("1"):
-            for var, properties in self.variables.items():
-                model.nodes[var] = properties
-        else:
-            for var, properties in self.variables.items():
-                model._node[var] = properties
+        for var, properties in self.variables.items():
+            model._node[var] = properties
 
         return model
 
@@ -435,28 +432,28 @@ class XBNWriter(object):
         cpds = self.model.get_cpds()
         cpds.sort(key=lambda x: x.variable)
         for cpd in cpds:
-            cpd_values = cpd.values.ravel()
+            cpd_values = cpd.get_values().transpose()
             var = cpd.variable
             dist = etree.SubElement(
                 distributions, "DIST", attrib={"TYPE": self.model.nodes[var]["TYPE"]}
             )
             etree.SubElement(dist, "PRIVATE", attrib={"NAME": var})
             dpis = etree.SubElement(dist, "DPIS")
-            evidence = cpd.variables[:0:-1]
+            evidence = cpd.variables[1:]
+            evidence_card = cpd.cardinality[1:]
             if evidence:
                 condset = etree.SubElement(dist, "CONDSET")
-                for condelem in sorted(evidence):
+                for condelem in evidence:
                     etree.SubElement(condset, "CONDELEM", attrib={"NAME": condelem})
-                # TODO: Get Index value.
-                for val in range(0, len(cpd_values), 2):
-                    etree.SubElement(dpis, "DPI", attrib={"INDEXES": " "}).text = (
-                        " "
-                        + str(cpd_values[val])
-                        + " "
-                        + str(cpd_values[val + 1])
-                        + " "
-                    )
+                indexes_iter = itertools.product(
+                    *[range(card) for card in evidence_card]
+                )
+                for val in range(cpd_values.shape[0]):
+                    index_value = " " + " ".join(map(str, next(indexes_iter))) + " "
+                    etree.SubElement(
+                        dpis, "DPI", attrib={"INDEXES": index_value}
+                    ).text = (" " + " ".join(map(str, cpd_values[val])) + " ")
             else:
-                etree.SubElement(dpis, "DPI").text = " " + " ".join(
-                    map(str, cpd_values)
+                etree.SubElement(dpis, "DPI").text = (
+                    " " + " ".join(map(str, cpd_values[0])) + " "
                 )


Test patch:
diff --git a/pgmpy/tests/test_factors/test_discrete/test_Factor.py b/pgmpy/tests/test_factors/test_discrete/test_Factor.py
index c674ad49..1ed1b0e3 100644
--- a/pgmpy/tests/test_factors/test_discrete/test_Factor.py
+++ b/pgmpy/tests/test_factors/test_discrete/test_Factor.py
@@ -828,7 +828,7 @@ class TestHash:
 
 class TestTabularCPDInit(unittest.TestCase):
     def test_cpd_init(self):
-        cpd = TabularCPD("grade", 3, [[0.1, 0.1, 0.1]])
+        cpd = TabularCPD("grade", 3, [[0.1], [0.1], [0.1]])
         self.assertEqual(cpd.variable, "grade")
         self.assertEqual(cpd.variable_card, 3)
         self.assertEqual(list(cpd.variables), ["grade"])
diff --git a/pgmpy/tests/test_inference/test_Inference.py b/pgmpy/tests/test_inference/test_Inference.py
index 7f03a81c..85644b45 100644
--- a/pgmpy/tests/test_inference/test_Inference.py
+++ b/pgmpy/tests/test_inference/test_Inference.py
@@ -12,7 +12,7 @@ from collections import defaultdict
 class TestInferenceBase(unittest.TestCase):
     def setUp(self):
         self.bayesian = BayesianModel([("a", "b"), ("b", "c"), ("c", "d"), ("d", "e")])
-        a_cpd = TabularCPD("a", 2, [[0.4, 0.6]])
+        a_cpd = TabularCPD("a", 2, [[0.4], [0.6]])
         b_cpd = TabularCPD(
             "b", 2, [[0.2, 0.4], [0.8, 0.6]], evidence=["a"], evidence_card=[2]
         )
diff --git a/pgmpy/tests/test_inference/test_dbn_inference.py b/pgmpy/tests/test_inference/test_dbn_inference.py
index 981208dc..4c3d9447 100644
--- a/pgmpy/tests/test_inference/test_dbn_inference.py
+++ b/pgmpy/tests/test_inference/test_dbn_inference.py
@@ -17,7 +17,7 @@ class TestDBNInference(unittest.TestCase):
         dbn_1.add_edges_from(
             [(("Z", 0), ("X", 0)), (("Z", 0), ("Y", 0)), (("Z", 0), ("Z", 1))]
         )
-        cpd_start_z_1 = TabularCPD(("Z", 0), 2, [[0.8, 0.2]])
+        cpd_start_z_1 = TabularCPD(("Z", 0), 2, [[0.8], [0.2]])
         cpd_x_1 = TabularCPD(("X", 0), 2, [[0.9, 0.6], [0.1, 0.4]], [("Z", 0)], [2])
         cpd_y_1 = TabularCPD(("Y", 0), 2, [[0.7, 0.2], [0.3, 0.8]], [("Z", 0)], [2])
         cpd_trans_z_1 = TabularCPD(
@@ -30,7 +30,7 @@ class TestDBNInference(unittest.TestCase):
         dbn_2.add_edges_from(
             [(("Z", 0), ("X", 0)), (("X", 0), ("Y", 0)), (("Z", 0), ("Z", 1))]
         )
-        cpd_start_z_2 = TabularCPD(("Z", 0), 2, [[0.5, 0.5]])
+        cpd_start_z_2 = TabularCPD(("Z", 0), 2, [[0.5], [0.5]])
         cpd_x_2 = TabularCPD(("X", 0), 2, [[0.6, 0.9], [0.4, 0.1]], [("Z", 0)], [2])
         cpd_y_2 = TabularCPD(("Y", 0), 2, [[0.2, 0.3], [0.8, 0.7]], [("X", 0)], [2])
         cpd_z_2 = TabularCPD(("Z", 1), 2, [[0.4, 0.7], [0.6, 0.3]], [("Z", 0)], [2])
diff --git a/pgmpy/tests/test_models/test_BayesianModel.py b/pgmpy/tests/test_models/test_BayesianModel.py
index f2887601..f574ea29 100644
--- a/pgmpy/tests/test_models/test_BayesianModel.py
+++ b/pgmpy/tests/test_models/test_BayesianModel.py
@@ -448,8 +448,8 @@ class TestBayesianModelCPD(unittest.TestCase):
         self.G.add_cpds(cpd_g, cpd_s, cpd_l)
         self.assertRaises(ValueError, self.G.check_model)
 
-        cpd_d = TabularCPD("d", 2, values=[[0.8, 0.2]])
-        cpd_i = TabularCPD("i", 2, values=[[0.7, 0.3]])
+        cpd_d = TabularCPD("d", 2, values=[[0.8], [0.2]])
+        cpd_i = TabularCPD("i", 2, values=[[0.7], [0.3]])
         self.G.add_cpds(cpd_d, cpd_i)
 
         self.assertTrue(self.G.check_model())
diff --git a/pgmpy/tests/test_models/test_DynamicBayesianNetwork.py b/pgmpy/tests/test_models/test_DynamicBayesianNetwork.py
index 3c8cfaaf..b5440786 100644
--- a/pgmpy/tests/test_models/test_DynamicBayesianNetwork.py
+++ b/pgmpy/tests/test_models/test_DynamicBayesianNetwork.py
@@ -94,8 +94,8 @@ class TestDynamicBayesianNetworkMethods(unittest.TestCase):
             evidence=[("D", 0)],
             evidence_card=[2],
         )
-        self.diff_cpd = TabularCPD(("D", 0), 2, values=[[0.6, 0.4]])
-        self.intel_cpd = TabularCPD(("I", 0), 2, values=[[0.7, 0.3]])
+        self.diff_cpd = TabularCPD(("D", 0), 2, values=[[0.6], [0.4]])
+        self.intel_cpd = TabularCPD(("I", 0), 2, values=[[0.7], [0.3]])
         self.i_i_cpd = TabularCPD(
             ("I", 1),
             2,
@@ -342,9 +342,9 @@ class TestDynamicBayesianNetworkMethods2(unittest.TestCase):
             evidence_card=[2],
         )
 
-        diff_cpd = TabularCPD(("D", 0), 2, values=[[0.6, 0.4]])
+        diff_cpd = TabularCPD(("D", 0), 2, values=[[0.6], [0.4]])
 
-        intel_cpd = TabularCPD(("I", 0), 2, values=[[0.7, 0.3]])
+        intel_cpd = TabularCPD(("I", 0), 2, values=[[0.7], [0.3]])
 
         i_i_cpd = TabularCPD(
             ("I", 1),
@@ -461,12 +461,12 @@ class TestDynamicBayesianNetworkMethods2(unittest.TestCase):
         self.assertRaises(ValueError, self.G.check_model)
         self.G.remove_cpds(d_i_cpd)
 
-        diff_cpd = TabularCPD(("D", 0), 2, values=[[0.7, 0.4]])
+        diff_cpd = TabularCPD(("D", 0), 2, values=[[0.7], [0.4]])
         self.G.add_cpds(diff_cpd)
         self.assertRaises(ValueError, self.G.check_model)
         self.G.remove_cpds(diff_cpd)
 
-        intel_cpd = TabularCPD(("I", 0), 2, values=[[1.7, 0.3]])
+        intel_cpd = TabularCPD(("I", 0), 2, values=[[1.7], [0.3]])
         self.G.add_cpds(intel_cpd)
         self.assertRaises(ValueError, self.G.check_model)
         self.G.remove_cpds(intel_cpd)
diff --git a/pgmpy/tests/test_readwrite/test_XMLBIF.py b/pgmpy/tests/test_readwrite/test_XMLBIF.py
index 6ef7b758..7a14d6be 100644
--- a/pgmpy/tests/test_readwrite/test_XMLBIF.py
+++ b/pgmpy/tests/test_readwrite/test_XMLBIF.py
@@ -314,8 +314,8 @@ class TestXMLBIFWriterMethodsString(unittest.TestCase):
         self.model_stateless = BayesianModel(
             [("D", "G"), ("I", "G"), ("G", "L"), ("I", "S")]
         )
-        self.cpd_d = TabularCPD(variable="D", variable_card=2, values=[[0.6, 0.4]])
-        self.cpd_i = TabularCPD(variable="I", variable_card=2, values=[[0.7, 0.3]])
+        self.cpd_d = TabularCPD(variable="D", variable_card=2, values=[[0.6], [0.4]])
+        self.cpd_i = TabularCPD(variable="I", variable_card=2, values=[[0.7], [0.3]])
 
         self.cpd_g = TabularCPD(
             variable="G",
diff --git a/pgmpy/tests/test_readwrite/test_XMLBeliefNetwork.py b/pgmpy/tests/test_readwrite/test_XMLBeliefNetwork.py
index 6f679688..df39b842 100644
--- a/pgmpy/tests/test_readwrite/test_XMLBeliefNetwork.py
+++ b/pgmpy/tests/test_readwrite/test_XMLBeliefNetwork.py
@@ -210,43 +210,43 @@ class TestXBNReader(unittest.TestCase):
         distribution = self.reader_string.get_distributions()
         self.assertEqual(distribution["a"]["TYPE"], "discrete")
         self.assertListEqual(distribution["b"]["CONDSET"], ["a"])
-        np_test.assert_array_equal(distribution["a"]["DPIS"], np.array([[0.2, 0.8]]))
-        np_test.assert_array_equal(distribution["f"]["DPIS"], np.array([[0.3, 0.7]]))
+        np_test.assert_array_equal(distribution["a"]["DPIS"], np.array([[0.2], [0.8]]))
+        np_test.assert_array_equal(distribution["f"]["DPIS"], np.array([[0.3], [0.7]]))
         np_test.assert_array_equal(
-            distribution["e"]["DPIS"], np.array([[0.8, 0.2], [0.6, 0.4]])
+            distribution["e"]["DPIS"], np.array([[0.8, 0.6], [0.2, 0.4]])
         )
         np_test.assert_array_equal(distribution["e"]["CARDINALITY"], np.array([2]))
         np_test.assert_array_equal(
             distribution["d"]["DPIS"],
-            np.array([[0.8, 0.2], [0.9, 0.1], [0.7, 0.3], [0.05, 0.95]]),
+            np.array([[0.8, 0.9, 0.7, 0.05], [0.2, 0.1, 0.3, 0.95]]),
         )
         np_test.assert_array_equal(
             distribution["b"]["DPIS"], np.array([[0.8, 0.2], [0.2, 0.8]])
         )
         np_test.assert_array_equal(distribution["d"]["CARDINALITY"], np.array([2, 2]))
         np_test.assert_array_equal(
-            distribution["c"]["DPIS"], np.array([[0.2, 0.8], [0.05, 0.95]])
+            distribution["c"]["DPIS"], np.array([[0.2, 0.05], [0.8, 0.95]])
         )
         np_test.assert_array_equal(distribution["c"]["CARDINALITY"], np.array([2]))
         distribution = self.reader_file.get_distributions()
         self.assertEqual(distribution["a"]["TYPE"], "discrete")
         self.assertListEqual(distribution["b"]["CONDSET"], ["a"])
-        np_test.assert_array_equal(distribution["a"]["DPIS"], np.array([[0.2, 0.8]]))
-        np_test.assert_array_equal(distribution["f"]["DPIS"], np.array([[0.3, 0.7]]))
+        np_test.assert_array_equal(distribution["a"]["DPIS"], np.array([[0.2], [0.8]]))
+        np_test.assert_array_equal(distribution["f"]["DPIS"], np.array([[0.3], [0.7]]))
         np_test.assert_array_equal(
-            distribution["e"]["DPIS"], np.array([[0.8, 0.2], [0.6, 0.4]])
+            distribution["e"]["DPIS"], np.array([[0.8, 0.6], [0.2, 0.4]])
         )
         np_test.assert_array_equal(distribution["e"]["CARDINALITY"], np.array([2]))
         np_test.assert_array_equal(
             distribution["d"]["DPIS"],
-            np.array([[0.8, 0.2], [0.9, 0.1], [0.7, 0.3], [0.05, 0.95]]),
+            np.array([[0.8, 0.9, 0.7, 0.05], [0.2, 0.1, 0.3, 0.95]]),
         )
         np_test.assert_array_equal(distribution["d"]["CARDINALITY"], np.array([2, 2]))
         np_test.assert_array_equal(
             distribution["b"]["DPIS"], np.array([[0.8, 0.2], [0.2, 0.8]])
         )
         np_test.assert_array_equal(
-            distribution["c"]["DPIS"], np.array([[0.2, 0.8], [0.05, 0.95]])
+            distribution["c"]["DPIS"], np.array([[0.2, 0.05], [0.8, 0.95]])
         )
         np_test.assert_array_equal(distribution["c"]["CARDINALITY"], np.array([2]))
 
@@ -298,11 +298,11 @@ class TestXBNReader(unittest.TestCase):
         }
         cpds_expected = {
             "b": np.array([[0.8, 0.2], [0.2, 0.8]]),
-            "e": np.array([[0.8, 0.2], [0.6, 0.4]]),
+            "e": np.array([[0.8, 0.6], [0.2, 0.4]]),
             "f": np.array([[0.3], [0.7]]),
-            "c": np.array([[0.2, 0.8], [0.05, 0.95]]),
+            "c": np.array([[0.2, 0.05], [0.8, 0.95]]),
             "a": np.array([[0.2], [0.8]]),
-            "d": np.array([[0.8, 0.2, 0.9, 0.1], [0.7, 0.3, 0.05, 0.95]]),
+            "d": np.array([[0.8, 0.9, 0.7, 0.05], [0.2, 0.1, 0.3, 0.95]]),
         }
         for cpd in model.get_cpds():
             np_test.assert_array_equal(cpd.get_values(), cpds_expected[cpd.variable])
@@ -365,14 +365,14 @@ class TestXBNWriter(unittest.TestCase):
             [("b", "d"), ("a", "b"), ("a", "c"), ("c", "d"), ("c", "e")]
         )
         cpd_distribution = {
-            "a": {"TYPE": "discrete", "DPIS": np.array([[0.2, 0.8]])},
+            "a": {"TYPE": "discrete", "DPIS": np.array([[0.2], [0.8]])},
             "e": {
                 "TYPE": "discrete",
-                "DPIS": np.array([[0.8, 0.2], [0.6, 0.4]]),
+                "DPIS": np.array([[0.8, 0.6], [0.2, 0.4]]),
                 "CONDSET": ["c"],
                 "CARDINALITY": [2],
             },
-            "f": {"TYPE": "discrete", "DPIS": np.array([[0.3, 0.7]])},
+            "f": {"TYPE": "discrete", "DPIS": np.array([[0.3], [0.7]])},
             "b": {
                 "TYPE": "discrete",
                 "DPIS": np.array([[0.8, 0.2], [0.2, 0.8]]),
@@ -381,13 +381,13 @@ class TestXBNWriter(unittest.TestCase):
             },
             "c": {
                 "TYPE": "discrete",
-                "DPIS": np.array([[0.2, 0.8], [0.05, 0.95]]),
+                "DPIS": np.array([[0.2, 0.05], [0.8, 0.95]]),
                 "CONDSET": ["a"],
                 "CARDINALITY": [2],
             },
             "d": {
                 "TYPE": "discrete",
-                "DPIS": np.array([[0.8, 0.2], [0.9, 0.1], [0.7, 0.3], [0.05, 0.95]]),
+                "DPIS": np.array([[0.8, 0.9, 0.7, 0.05], [0.2, 0.1, 0.3, 0.95]]),
                 "CONDSET": ["b", "c"],
                 "CARDINALITY": [2, 2],
             },
@@ -462,14 +462,14 @@ class TestXBNWriter(unittest.TestCase):
       <DIST TYPE="discrete">
         <PRIVATE NAME="a"/>
         <DPIS>
-          <DPI> 0.2 0.8</DPI>
+          <DPI> 0.2 0.8 </DPI>
         </DPIS>
       </DIST>
       <DIST TYPE="discrete">
         <PRIVATE NAME="b"/>
         <DPIS>
-          <DPI INDEXES=" "> 0.8 0.2 </DPI>
-          <DPI INDEXES=" "> 0.2 0.8 </DPI>
+          <DPI INDEXES=" 0 "> 0.8 0.2 </DPI>
+          <DPI INDEXES=" 1 "> 0.2 0.8 </DPI>
         </DPIS>
         <CONDSET>
           <CONDELEM NAME="a"/>
@@ -478,8 +478,8 @@ class TestXBNWriter(unittest.TestCase):
       <DIST TYPE="discrete">
         <PRIVATE NAME="c"/>
         <DPIS>
-          <DPI INDEXES=" "> 0.2 0.8 </DPI>
-          <DPI INDEXES=" "> 0.05 0.95 </DPI>
+          <DPI INDEXES=" 0 "> 0.2 0.8 </DPI>
+          <DPI INDEXES=" 1 "> 0.05 0.95 </DPI>
         </DPIS>
         <CONDSET>
           <CONDELEM NAME="a"/>
@@ -488,10 +488,10 @@ class TestXBNWriter(unittest.TestCase):
       <DIST TYPE="discrete">
         <PRIVATE NAME="d"/>
         <DPIS>
-          <DPI INDEXES=" "> 0.8 0.2 </DPI>
-          <DPI INDEXES=" "> 0.9 0.1 </DPI>
-          <DPI INDEXES=" "> 0.7 0.3 </DPI>
-          <DPI INDEXES=" "> 0.05 0.95 </DPI>
+          <DPI INDEXES=" 0 0 "> 0.8 0.2 </DPI>
+          <DPI INDEXES=" 0 1 "> 0.9 0.1 </DPI>
+          <DPI INDEXES=" 1 0 "> 0.7 0.3 </DPI>
+          <DPI INDEXES=" 1 1 "> 0.05 0.95 </DPI>
         </DPIS>
         <CONDSET>
           <CONDELEM NAME="b"/>
@@ -501,8 +501,8 @@ class TestXBNWriter(unittest.TestCase):
       <DIST TYPE="discrete">
         <PRIVATE NAME="e"/>
         <DPIS>
-          <DPI INDEXES=" "> 0.8 0.2 </DPI>
-          <DPI INDEXES=" "> 0.6 0.4 </DPI>
+          <DPI INDEXES=" 0 "> 0.8 0.2 </DPI>
+          <DPI INDEXES=" 1 "> 0.6 0.4 </DPI>
         </DPIS>
         <CONDSET>
           <CONDELEM NAME="c"/>
@@ -511,13 +511,15 @@ class TestXBNWriter(unittest.TestCase):
       <DIST TYPE="discrete">
         <PRIVATE NAME="f"/>
         <DPIS>
-          <DPI> 0.3 0.7</DPI>
+          <DPI> 0.3 0.7 </DPI>
         </DPIS>
       </DIST>
     </DISTRIBUTIONS>
   </BNMODEL>
 </ANALYSISNOTEBOOK>"""
         )
+        output = str(self.writer.__str__())
+        expected = str(etree.tostring(self.expected_xml))
         self.assertEqual(
             str(self.writer.__str__()[:-1]), str(etree.tostring(self.expected_xml))
         )
diff --git a/pgmpy/tests/test_utils/test_state_name.py b/pgmpy/tests/test_utils/test_state_name.py
index 30332e9a..2710c000 100644
--- a/pgmpy/tests/test_utils/test_state_name.py
+++ b/pgmpy/tests/test_utils/test_state_name.py
@@ -55,8 +55,8 @@ class TestStateNameInit(unittest.TestCase):
         )
 
         student = BayesianModel([("diff", "grade"), ("intel", "grade")])
-        diff_cpd = TabularCPD("diff", 2, [[0.2, 0.8]])
-        intel_cpd = TabularCPD("intel", 2, [[0.3, 0.7]])
+        diff_cpd = TabularCPD("diff", 2, [[0.2], [0.8]])
+        intel_cpd = TabularCPD("intel", 2, [[0.3], [0.7]])
         grade_cpd = TabularCPD(
             "grade",
             3,
@@ -122,8 +122,8 @@ class StateNameDecorator(unittest.TestCase):
         student = BayesianModel([("diff", "grade"), ("intel", "grade")])
         student_state_names = BayesianModel([("diff", "grade"), ("intel", "grade")])
 
-        diff_cpd = TabularCPD("diff", 2, [[0.2, 0.8]])
-        intel_cpd = TabularCPD("intel", 2, [[0.3, 0.7]])
+        diff_cpd = TabularCPD("diff", 2, [[0.2], [0.8]])
+        intel_cpd = TabularCPD("intel", 2, [[0.3], [0.7]])
         grade_cpd = TabularCPD(
             "grade",
             3,
@@ -135,13 +135,13 @@ class StateNameDecorator(unittest.TestCase):
         diff_cpd_state_names = TabularCPD(
             variable="diff",
             variable_card=2,
-            values=[[0.2, 0.8]],
+            values=[[0.2], [0.8]],
             state_names={"diff": ["high", "low"]},
         )
         intel_cpd_state_names = TabularCPD(
             variable="intel",
             variable_card=2,
-            values=[[0.3, 0.7]],
+            values=[[0.3], [0.7]],
             state_names={"intel": ["poor", "good", "very good"]},
         )
         grade_cpd_state_names = TabularCPD(