// Mul (t^exp) to the dict "d"
void Mul::dict_add_term_new(const Ptr<RCP<const Number>> &coef, map_basic_basic &d,
const RCP<const Basic> &exp, const RCP<const Basic> &t)
{
auto it = d.find(t);
if (it == d.end()) {
// Don't check for `exp = 0` here
// `pow` for Complex is not expanded by default
if (is_a<Integer>(*exp) && (is_a<Integer>(*t) || is_a<Rational>(*t))) {
imulnum(outArg(*coef), pownum(rcp_static_cast<const Number>(t),
rcp_static_cast<const Number>(exp)));
} else if (is_a<Integer>(*exp) && is_a<Complex>(*t)) {
if (rcp_static_cast<const Integer>(exp)->is_one()) {
imulnum(outArg(*coef), rcp_static_cast<const Number>(t));
} else if (rcp_static_cast<const Integer>(exp)->is_minus_one()) {
idivnum(outArg(*coef), rcp_static_cast<const Number>(t));
} else {
insert(d, t, exp);
}
} else {
insert(d, t, exp);
}
} else {
// Very common case, needs to be fast:
if (is_a_Number(*exp) && is_a_Number(*it->second)) {
RCP<const Number> tmp = rcp_static_cast<const Number>(it->second);
iaddnum(outArg(tmp),
rcp_static_cast<const Number>(exp));
it->second = tmp;
}
else
it->second = add(it->second, exp);
if (is_a<Integer>(*it->second)) {
// `pow` for Complex is not expanded by default
if (is_a<Integer>(*t) || is_a<Rational>(*t)) {
if (!rcp_static_cast<const Integer>(it->second)->is_zero()) {
imulnum(outArg(*coef), pownum(rcp_static_cast<const Number>(t),
rcp_static_cast<const Number>(it->second)));
}
d.erase(it);
} else if (rcp_static_cast<const Integer>(it->second)->is_zero()) {
d.erase(it);
} else if (is_a<Complex>(*t)) {
if (rcp_static_cast<const Integer>(it->second)->is_one()) {
imulnum(outArg(*coef), rcp_static_cast<const Number>(t));
d.erase(it);
} else if (rcp_static_cast<const Integer>(it->second)->is_minus_one()) {
idivnum(outArg(*coef), rcp_static_cast<const Number>(t));
d.erase(it);
}
}
}
}
}
RCP<const Basic> pow(const RCP<const Basic> &a, const RCP<const Basic> &b)
{
if (eq(b, zero)) return one;
if (eq(b, one)) return a;
if (eq(a, zero)) return zero;
if (eq(a, one)) return one;
if (is_a_Number(*a) && is_a<Integer>(*b))
return pownum(rcp_static_cast<const Number>(a), rcp_static_cast<const Integer>(b));
if (is_a<Mul>(*a))
return rcp_static_cast<const Mul>(a)->power_all_terms(b);
if (is_a<Pow>(*a)) {
RCP<const Pow> A = rcp_static_cast<const Pow>(a);
return pow(A->base_, mul(A->exp_, b));
}
return rcp(new Pow(a, b));
}
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;
}
// Mul (t**exp) to the dict "d"
void Mul::dict_add_term_new(const Ptr<RCP<const Number>> &coef, map_basic_basic &d,
const RCP<const Basic> &exp, const RCP<const Basic> &t)
{
auto it = d.find(t);
if (it == d.end()) {
// Don't check for `exp = 0` here
// `pow` for Complex is not expanded by default
if (is_a<Integer>(*t) || is_a<Rational>(*t)) {
if (is_a<Integer>(*exp)) {
imulnum(outArg(*coef), pownum(rcp_static_cast<const Number>(t),
rcp_static_cast<const Number>(exp)));
} else if (is_a<Rational>(*exp)) {
// Here we make the exponent postive and a fraction between
// 0 and 1.
mpz_class q, r, num, den;
num = rcp_static_cast<const Rational>(exp)->i.get_num();
den = rcp_static_cast<const Rational>(exp)->i.get_den();
mpz_fdiv_qr(q.get_mpz_t(), r.get_mpz_t(), num.get_mpz_t(),
den.get_mpz_t());
insert(d, t, Rational::from_mpq(mpq_class(r, den)));
imulnum(outArg(*coef), pownum(rcp_static_cast<const Number>(t),
rcp_static_cast<const Number>(integer(q))));
} else {
insert(d, t, exp);
}
} else if (is_a<Integer>(*exp) && is_a<Complex>(*t)) {
if (rcp_static_cast<const Integer>(exp)->is_one()) {
imulnum(outArg(*coef), rcp_static_cast<const Number>(t));
} else if (rcp_static_cast<const Integer>(exp)->is_minus_one()) {
idivnum(outArg(*coef), rcp_static_cast<const Number>(t));
} else {
insert(d, t, exp);
}
} else {
insert(d, t, exp);
}
} else {
// Very common case, needs to be fast:
if (is_a_Number(*exp) && is_a_Number(*it->second)) {
RCP<const Number> tmp = rcp_static_cast<const Number>(it->second);
iaddnum(outArg(tmp),
rcp_static_cast<const Number>(exp));
it->second = tmp;
}
else
it->second = add(it->second, exp);
if (is_a<Integer>(*it->second)) {
// `pow` for Complex is not expanded by default
if (is_a<Integer>(*t) || is_a<Rational>(*t)) {
if (!rcp_static_cast<const Integer>(it->second)->is_zero()) {
imulnum(outArg(*coef), pownum(rcp_static_cast<const Number>(t),
rcp_static_cast<const Number>(it->second)));
}
d.erase(it);
} else if (rcp_static_cast<const Integer>(it->second)->is_zero()) {
d.erase(it);
} else if (is_a<Complex>(*t)) {
if (rcp_static_cast<const Integer>(it->second)->is_one()) {
imulnum(outArg(*coef), rcp_static_cast<const Number>(t));
d.erase(it);
} else if (rcp_static_cast<const Integer>(it->second)->is_minus_one()) {
idivnum(outArg(*coef), rcp_static_cast<const Number>(t));
d.erase(it);
}
}
} else if (is_a<Rational>(*it->second)) {
if (is_a_Number(*t)) {
mpz_class q, r, num, den;
num = rcp_static_cast<const Rational>(it->second)->i.get_num();
den = rcp_static_cast<const Rational>(it->second)->i.get_den();
// Here we make the exponent postive and a fraction between
// 0 and 1.
if (num > den || num < 0) {
mpz_fdiv_qr(q.get_mpz_t(), r.get_mpz_t(), num.get_mpz_t(),
den.get_mpz_t());
it->second = Rational::from_mpq(mpq_class(r, den));
imulnum(outArg(*coef), pownum(rcp_static_cast<const Number>(t),
rcp_static_cast<const Number>(integer(q))));
}
}
}
}
}
RCP<const Basic> pow(const RCP<const Basic> &a, const RCP<const Basic> &b)
{
if (is_a_Number(*b) and rcp_static_cast<const Number>(b)->is_zero()) {
return pownum(rcp_static_cast<const Number>(b), zero);
}
if (eq(*b, *one))
return a;
if (eq(*a, *zero))
return zero;
if (eq(*a, *one))
return one;
if (eq(*a, *minus_one)) {
if (is_a<Integer>(*b)) {
return is_a<Integer>(*div(b, integer(2))) ? one : minus_one;
} else if (is_a<Rational>(*b)
and (get_num(rcp_static_cast<const Rational>(b)->i) == 1)
and (get_den(rcp_static_cast<const Rational>(b)->i) == 2)) {
return I;
}
}
if (is_a_Number(*a) and is_a_Number(*b)) {
if (is_a<Integer>(*b)) {
if (is_a<Rational>(*a)) {
RCP<const Rational> exp_new
= rcp_static_cast<const Rational>(a);
return exp_new->powrat(*rcp_static_cast<const Integer>(b));
} else if (is_a<Integer>(*a)) {
RCP<const Integer> exp_new = rcp_static_cast<const Integer>(a);
return exp_new->powint(*rcp_static_cast<const Integer>(b));
} else if (is_a<Complex>(*a)) {
RCP<const Complex> exp_new = rcp_static_cast<const Complex>(a);
RCP<const Integer> pow_new = rcp_static_cast<const Integer>(b);
RCP<const Number> res = exp_new->pow(*pow_new);
return res;
} else {
return rcp_static_cast<const Number>(a)
->pow(*rcp_static_cast<const Number>(b));
}
} else if (is_a<Rational>(*b)) {
if (is_a<Rational>(*a)) {
return static_cast<const Rational &>(*a)
.powrat(static_cast<const Rational &>(*b));
} else if (is_a<Integer>(*a)) {
return static_cast<const Rational &>(*b)
.rpowrat(static_cast<const Integer &>(*a));
} else if (is_a<Complex>(*a)) {
return make_rcp<const Pow>(a, b);
} else {
return rcp_static_cast<const Number>(a)
->pow(*rcp_static_cast<const Number>(b));
}
} else if (is_a<Complex>(*b)) {
return make_rcp<const Pow>(a, b);
} else {
return rcp_static_cast<const Number>(a)
->pow(*rcp_static_cast<const Number>(b));
}
}
if (is_a<Mul>(*a) and is_a_Number(*b)) {
map_basic_basic d;
RCP<const Number> coef = one;
rcp_static_cast<const Mul>(a)
->power_num(outArg(coef), d, rcp_static_cast<const Number>(b));
return Mul::from_dict(coef, std::move(d));
}
if (is_a<Pow>(*a) and is_a<Integer>(*b)) {
// Convert (x**y)**b = x**(b*y), where 'b' is an integer. This holds for
// any complex 'x', 'y' and integer 'b'.
RCP<const Pow> A = rcp_static_cast<const Pow>(a);
return pow(A->get_base(), mul(A->get_exp(), b));
}
return make_rcp<const Pow>(a, b);
}
