implemented CoefficientIterator with bindings

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);

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()}")

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)) {

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
          */

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