Exemplo n.º 1
0
RCP<Basic> pow_expand(const RCP<Pow> &self)
{
    if (is_a<Integer>(*self->exp_)) {
        if (is_a<Add>(*self->base_)) {
            map_vec_mpz r;
            int n = rcp_dynamic_cast<Integer>(self->exp_)->as_int();

            RCP<Add> base = rcp_dynamic_cast<Add>(self->base_);
            int m = base->dict_.size();
            multinomial_coefficients_mpz(m, n, r);
            Dict_int rd;
            for (auto &p: r) {
                auto power = p.first.begin();
                auto i2 = base->dict_.begin();
                Dict_int d;
                for (; power != p.first.end(); ++power, ++i2) {
                    if (*power > 0) {
                        RCP<Integer> exp = rcp(new Integer(*power));
                        RCP<Basic> base = i2->first;
                        d[base] = exp;
                    }
                }
                RCP<Basic> term = Mul::from_dict(one, d);
                rd[term] = rcp(new Integer(p.second));
            }
            RCP<Basic> result = Add::from_dict(rd);
            return result;
        }
    }
    return self;
}
Exemplo n.º 2
0
RCP<const Basic> pow_expand(const RCP<const Pow> &self)
{
    RCP<const Basic> _base = expand(self->base_);
    bool negative_pow = false;
    if (! is_a<Integer>(*self->exp_) || ! is_a<Add>(*_base)) {
        if (neq(_base, self->base_)) {
            return pow(_base, self->exp_);
        } else {
            return self;
        }
    }

    map_vec_mpz r;
    int n = rcp_static_cast<const Integer>(self->exp_)->as_int();
    if (n < 0) {
        n = -n;
        negative_pow = true;
    }
    RCP<const Add> base = rcp_static_cast<const Add>(_base);
    umap_basic_num base_dict = base->dict_;
    if (! (base->coef_->is_zero())) {
        // Add the numerical coefficient into the dictionary. This
        // allows a little bit easier treatment below.
        insert(base_dict, base->coef_, one);
    }
    int m = base_dict.size();
    multinomial_coefficients_mpz(m, n, r);
    umap_basic_num rd;
    // This speeds up overall expansion. For example for the benchmark
    // (y + x + z + w)^60 it improves the timing from 135ms to 124ms.
    rd.reserve(2*r.size());
    RCP<const Number> add_overall_coeff=zero;
    for (auto &p: r) {
        auto power = p.first.begin();
        auto i2 = base_dict.begin();
        map_basic_basic d;
        RCP<const Number> overall_coeff=one;
        for (; power != p.first.end(); ++power, ++i2) {
            if (*power > 0) {
                RCP<const Integer> exp = rcp(new Integer(*power));
                RCP<const Basic> base = i2->first;
                if (is_a<Integer>(*base)) {
                    imulnum(outArg(overall_coeff),
                        rcp_static_cast<const Number>(
                        rcp_static_cast<const Integer>(base)->powint(*exp)));
                } else if (is_a<Symbol>(*base)) {
                    Mul::dict_add_term(d, exp, base);
                } else {
                    RCP<const Basic> exp2, t, tmp;
                    tmp = pow(base, exp);
                    if (is_a<Mul>(*tmp)) {
                        for (auto &p: (rcp_static_cast<const Mul>(tmp))->dict_) {
                            Mul::dict_add_term_new(outArg(overall_coeff), d,
                                    p.second, p.first);
                        }
                        imulnum(outArg(overall_coeff), (rcp_static_cast<const Mul>(tmp))->coef_);
                    } else {
                        Mul::as_base_exp(tmp, outArg(exp2), outArg(t));
                        Mul::dict_add_term_new(outArg(overall_coeff), d, exp2, t);
                    }
                }
                if (!(i2->second->is_one())) {
                    if (is_a<Integer>(*(i2->second)) || is_a<Rational>(*(i2->second))) {
                        imulnum(outArg(overall_coeff),
                        pownum(i2->second,
                            rcp_static_cast<const Number>(exp)));
                    } else if (is_a<Complex>(*(i2->second))) {
                        RCP<const Number> tmp = rcp_static_cast<const Complex>(i2->second)->pow(*exp);
                        imulnum(outArg(overall_coeff), tmp);
                    }
                }
            }
        }
        RCP<const Basic> term = Mul::from_dict(overall_coeff, std::move(d));
        RCP<const Number> coef2 = rcp(new Integer(p.second));
        if (is_a_Number(*term)) {
            iaddnum(outArg(add_overall_coeff),
                mulnum(rcp_static_cast<const Number>(term), coef2));
        } else {
            if (is_a<Mul>(*term) &&
                    !(rcp_static_cast<const Mul>(term)->coef_->is_one())) {
                // Tidy up things like {2x: 3} -> {x: 6}
                imulnum(outArg(coef2),
                        rcp_static_cast<const Mul>(term)->coef_);
                // We make a copy of the dict_:
                map_basic_basic d2 = rcp_static_cast<const Mul>(term)->dict_;
                term = Mul::from_dict(one, std::move(d2));
            }
            Add::dict_add_term(rd, coef2, term);
        }
    }
    RCP<const Basic> result = Add::from_dict(add_overall_coeff, std::move(rd));
    if (negative_pow) result = pow(result, minus_one);
    return result;
}