Combine labeled distributions

Documentation/CHANGELOG.md

@@ -32,6 +32,7 @@ Release Date: TBD
 * Add methods to non stochastically simulate an economy by computing transition matrices. Functions to compute transition matrices and ergodic distribution have been added [#1155](https://github.com/econ-ark/HARK/pull/1155).
 * Fixes a bug that causes `t_age` and `t_cycle` to get out of sync when reading pre-computed mortality. [#1181](https://github.com/econ-ark/HARK/pull/1181)
 * Adds Methods to calculate Heterogenous Agent Jacobian matrices. [#1185](https://github.com/econ-ark/HARK/pull/1185)
+* Enhances `combine_indep_dstns` to work with labeled distributions (`DiscreteDistributionLabeled`). [#1191](htttps://github.com/econ-ark/HARK/pull/1191)
 
 ### Minor Changes
 

HARK/distribution.py

@@ -5,6 +5,7 @@
 import xarray as xr
 from scipy import stats
 from scipy.special import erf, erfc
+from warnings import warn
 
 
 class Distribution:
@@ -1614,6 +1615,8 @@ def combine_indep_dstns(*distributions, seed=0):
     # Get information on the distributions
     dist_lengths = ()
     dist_dims = ()
+    dist_is_labeled = ()
+    var_labels = ()
     for dist in distributions:
 
         if len(dist.dim()) > 1:
@@ -1624,8 +1627,18 @@ def combine_indep_dstns(*distributions, seed=0):
         dist_dims += (dist.dim(),)
         dist_lengths += (len(dist.pmv),)
 
+        labeled = isinstance(dist, DiscreteDistributionLabeled)
+        dist_is_labeled += (labeled,)
+        if labeled:
+            var_labels += tuple(dist.dataset.data_vars.keys())
+        else:
+            var_labels += tuple([""] * dist.dim()[0])
+
     number_of_distributions = len(distributions)
 
+    all_labeled = all(dist_is_labeled)
+    labels_are_unique = len(var_labels) == len(set(var_labels))
+
     # We need the combinations of indices of realizations in all
     # distributions
     inds = np.meshgrid(
@@ -1645,7 +1658,18 @@ def combine_indep_dstns(*distributions, seed=0):
 
     assert np.isclose(np.sum(P_out), 1), "Probabilities do not sum to 1!"
 
-    return DiscreteDistribution(P_out, atoms_out, seed=seed)
+    if all_labeled and labels_are_unique:
+        combined_dstn = DiscreteDistributionLabeled(
+            pmv=P_out, data=atoms_out, var_names=var_labels, seed=seed,
+        )
+    else:
+        if all_labeled and not labels_are_unique:
+            warn(
+                "There are duplicated labels in the provided distributions. Returning a non-labeled combination"
+            )
+        combined_dstn = DiscreteDistribution(P_out, atoms_out, seed=seed)
+
+    return combined_dstn
 
 
 def calc_expectation(dstn, func=lambda x: x, *args):

HARK/tests/test_distribution.py

@@ -533,3 +533,59 @@ def test_self_expected_value(self):
         )
 
         self.assertAlmostEqual(ce2[3], 9.518015322143837)
+
+    def test_combine_labeled_dist(self):
+
+        # Create some dstns
+        a = DiscreteDistributionLabeled(
+            pmv=np.array([0.1, 0.9]), data=np.array([-1.0, 1.0]), var_names="a"
+        )
+        b = DiscreteDistributionLabeled(
+            pmv=np.array([0.5, 0.5]), data=np.array([0.0, 1.0]), var_names="b"
+        )
+        c = DiscreteDistributionLabeled(
+            pmv=np.array([0.3, 0.7]), data=np.array([0.5, 1.0]), var_names="c"
+        )
+
+        # Test some combinations
+        abc = combine_indep_dstns(a, b, c)
+        # Check the order
+        self.assertTrue(
+            np.all(
+                abc.expected()
+                == np.concatenate([a.expected(), b.expected(), c.expected()])
+            )
+        )
+        # Check by label
+        self.assertEqual(abc.expected(lambda x: x["b"]), b.expected()[0])
+        self.assertAlmostEqual(
+            abc.expected(lambda x: x["a"] * x["c"]), a.expected()[0] * c.expected()[0]
+        )
+
+        # Combine labeled and non labeled distribution
+        x = DiscreteDistribution(pmv=np.array([0.5, 0.5]), atoms=np.array([1.0, 2.0]))
+
+        xa = combine_indep_dstns(x, a)
+        self.assertFalse(isinstance(xa, DiscreteDistributionLabeled))
+        self.assertTrue(
+            np.all(xa.expected() == np.concatenate([x.expected(), a.expected()]))
+        )
+
+        # Combine multidimensional labeled
+        d = DiscreteDistributionLabeled(
+            pmv=np.array([0.3, 0.7]), data=np.array([-0.5, -1.0]), var_names="d"
+        )
+        e = DiscreteDistributionLabeled(
+            pmv=np.array([0.3, 0.7]), data=np.array([0.0, -1.0]), var_names="e"
+        )
+        de = combine_indep_dstns(d, e)
+
+        abcde = combine_indep_dstns(abc, de)
+        self.assertTrue(
+            np.allclose(
+                abcde.expected(
+                    lambda x: np.array([x["d"], x["e"], x["a"], x["b"], x["c"]])
+                ),
+                np.concatenate([de.expected(), abc.expected()]),
+            )
+        )