update distr_of_function
return DD or DDXRA
@@ -1001,7 +1001,7 @@ class DiscreteDistribution(Distribution):
Returns
-------
- f_dstn : DiscreteDistribution
+ f_dstn : DiscreteDistribution or DiscreteDistributionXRA
The distribution of func(dstn).
"""
# we need to add one more dimension,
@@ -1012,6 +1012,11 @@ class DiscreteDistribution(Distribution):
]
f_query = func(self.X, *args)
+ if xarray:
+ f_dstn = DiscreteDistributionXRA(
+ list(self.pmf), f_query, seed=self.seed, **kwargs
+ )
+ else:
f_dstn = DiscreteDistribution(list(self.pmf), f_query, seed=self.seed)
return f_dstn
