diff --git a/python/bindings.cpp b/python/bindings.cpp
index 7b7b1b740aecbe37d12bc254b77ce63a3a26e976..46092643ba79a95d3488635ed1620b9e94aa7ec9 100644
--- a/python/bindings.cpp
+++ b/python/bindings.cpp
@@ -50,6 +50,7 @@ PYBIND11_MODULE(pygin, m) {
.def("filterLeq", &Dist::filterLeq)
.def("filterGreater", &Dist::filterGreater)
.def("filterGeq", &Dist::filterGeq)
+ .def("coefficient_iterator", &Dist::coefficientIterator)
.def(py::self == py::self)
.def(py::self != py::self)
//.def(string() * py::self)
@@ -71,6 +72,10 @@ PYBIND11_MODULE(pygin, m) {
.value("true", troolean::TRUE)
.value("unknown", troolean::UNKNOWN);
+ py::class_<Dist::CoefficientIterator>(m, "CoefficientIterator")
+ .def("next", &Dist::CoefficientIterator::next)
+ .def("rest", &Dist::CoefficientIterator::rest);
+
#ifdef VERSION_INFO
m.attr("__version__") = MACRO_STRINGIFY(VERSION_INFO);
diff --git a/python/test.py b/python/test.py
index 3aec377a6a64a801e30a48e4b2ff8319c4232b9c..2392a9ddbb99ac59b7b3cfa7781d858036eaa82f 100644
--- a/python/test.py
+++ b/python/test.py
@@ -25,3 +25,7 @@ if __name__ == "__main__":
check_if_zero(dist)
check_if_zero(pg.Dist("0"))
+
+ it = dist.coefficient_iterator(var)
+ print(f"first terms of dist are {it.next()}, {it.next()}, and {it.next()}")
+ print(f"remaining terms not consumed by iterator: {it.rest()}")
diff --git a/src/Dist.cpp b/src/Dist.cpp
index 964a7a6042bc5b5f107175f01e259b001047202d..e9fc9f3c7cc163418d33900dbcf3c15178a4d7a7 100644
--- a/src/Dist.cpp
+++ b/src/Dist.cpp
@@ -370,6 +370,7 @@ namespace prodigy {
EXPECTS(eExp.info(info_flags::nonnegint), "given expression must be non-neg int");
int order = ex_to<numeric>(eExp).to_int();
ex xExp;
+ // todo why not use parseWithKnownSymbols?
if(!isKnownAsVar(x)) {
xExp = symbol {x};
Dist::registerAsVar(ex_to<symbol>(xExp));
@@ -380,6 +381,18 @@ namespace prodigy {
return Dist{series_to_poly(series)};
}
+ Dist Dist::CoefficientIterator::next() {
+ ex resGf = _rest.subs(_var == 0);
+ _rest = (_rest - resGf) * pow(_var, -1);
+ _rest = _rest.normal(); // for canceling common factors
+ return Dist(resGf);
+ }
+
+ Dist::CoefficientIterator Dist::coefficientIterator(const std::string &var) const {
+ EXPECTS(isKnownAsVar(var), "can only expand series in known variables");
+ return CoefficientIterator(*this, ex_to<symbol>(_vars.at(var)));
+ }
+
Dist Dist::substituteParam(const std::string &p, const std::string &val) const {
EXPECTS(!isKnownAsVar(p), "given symbol is considered a variable");
if (!isKnownAsParam(p)) {
diff --git a/src/Dist.h b/src/Dist.h
index 79f12f890b6e8d8366f7b081e4ae801c1d417e15..7ac5bfb5aac55e13236532fa16d5b479f3dd121f 100644
--- a/src/Dist.h
+++ b/src/Dist.h
@@ -132,11 +132,25 @@ namespace prodigy {
const ex& gf() const;
+ /*
+ * returns the specified raw moment (possibly non-numeric due to presence of parameters)
+ * if *this is a distribution, then the raw moment is the expected value of the given monomial
+ * this function does not modify the global state (//todo why?)
+ */
+ ex moment(const std::map<std::string, int>& monomial) const;
+
/*
* returns total probability mass, possibly non-numeric due to params
+ * // todo implement in terms of moment
*/
ex mass() const;
+ /*
+ * returns expected value, possibly non-numeric due to params
+ * // todo implement in terms of moment
+ */
+ ex E(const std::string& var) const;
+
/*
* approximate size of underlying gf representation
*/
@@ -180,10 +194,9 @@ namespace prodigy {
troolean areIndependent(const std::string& x, const std::string& y) const;
/*
- * returns expected value, possibly non-numeric due to params
- * does not modify global state
+ * attempts to prove that *this is finite-support, i.e., just finitely many coefficients are non-zero
*/
- ex E(const std::string& var) const;
+ troolean isFiniteSupport() const;
/*
* update distribution according to assignment x := e
@@ -217,6 +230,29 @@ namespace prodigy {
Dist filterEq(const std::string& x, const std::string& e) const { return *this - filterLess(x, e) - filterGreater(x, e); }
Dist filterNeq(const std::string& x, const std::string& e) const { return *this - filterEq(x, e); }
+ /*
+ * this iterator regards a k-dimensional distribution as 1-dimensional in the given variable with coefficients
+ * that are k-1 dimensional, and it iterates of these coefficients
+ *
+ * in other words, the iterator ranges over the coefficients in the series expansion wrt. the given of the gf
+ * underlying the given distribution
+ */
+ class CoefficientIterator {
+ private:
+ ex _rest;
+ symbol _var;
+ public:
+ CoefficientIterator(Dist dist, symbol var) : _rest(std::move(dist._gf)), _var(std::move(var)) {}
+ Dist next();
+ Dist rest() const { return Dist(_rest); }
+ };
+
+ /*
+ * returns a coefficient iterator for the dimension specified by var
+ * throws an exception if var is not registered as a variable
+ */
+ CoefficientIterator coefficientIterator(const std::string& var) const;
+
/*
* replaces given parameter with given value or parametric expression
*/
diff --git a/test/TestDist.cpp b/test/TestDist.cpp
index 13babff0f33d727e77d2837d11dfa3c349637feb..09379b0d5787c3c8e83601ce51126eb3e5be9c49 100644
--- a/test/TestDist.cpp
+++ b/test/TestDist.cpp
@@ -69,3 +69,10 @@ TEST(Dist, gfSize) {
ASSERT_EQ(Dist("x").gfSize(), 1);
ASSERT_EQ(Dist("x+y").gfSize(), 3);
}
+
+TEST(Dist, coefficientIterator) {
+ Dist::resetSymbolCache();
+ Dist dist = Fac::geometric("x", "q");
+ Dist::CoefficientIterator it = dist.coefficientIterator("x");
+ ASSERT_EQ(it.next() + it.next() + it.next(), Dist("1-q - q^2 + q + q^2 - q^3", "q"));
+}
\ No newline at end of file