RCP<const Basic> Add::subs(const map_basic_basic &subs_dict) const
{
RCP<const Add> self = rcp_const_cast<Add>(rcp(this));
auto it = subs_dict.find(self);
if (it != subs_dict.end())
return it->second;
CSymPy::umap_basic_num d;
RCP<const Number> coef=coef_, coef2;
RCP<const Basic> t;
for (auto &p: dict_) {
RCP<const Basic> term = p.first->subs(subs_dict);
if (term == p.first) {
Add::dict_add_term(d, p.second, p.first);
} else if (is_a<Integer>(*term) &&
rcp_static_cast<const Integer>(term)->is_zero()) {
continue;
} else if (is_a_Number(*term)) {
iaddnum(outArg(coef),
mulnum(p.second, rcp_static_cast<const Number>(term)));
} else if (is_a<Add>(*term)) {
for (auto &q: (rcp_static_cast<const Add>(term))->dict_)
Add::dict_add_term(d, q.second, q.first);
iaddnum(outArg(coef), rcp_static_cast<const Add>(term)->coef_);
} else {
Add::as_coef_term(mul(p.second, term), outArg(coef2), outArg(t));
Add::dict_add_term(d, coef2, t);
}
}
return Add::from_dict(coef, std::move(d));
}
RCP<const Basic> mul(const RCP<const Basic> &a, const RCP<const Basic> &b)
{
SymEngine::map_basic_basic d;
RCP<const Number> coef = one;
if (SymEngine::is_a<Mul>(*a) && SymEngine::is_a<Mul>(*b)) {
RCP<const Mul> A = rcp_static_cast<const Mul>(a);
RCP<const Mul> B = rcp_static_cast<const Mul>(b);
// This is important optimization, as coef=1 if Mul is inside an Add.
// To further speed this up, the upper level code could tell us that we
// are inside an Add, then we don't even have can simply skip the
// following two lines.
if (!(A->coef_->is_one()) || !(B->coef_->is_one()))
coef = mulnum(A->coef_, B->coef_);
d = A->dict_;
for (auto &p: B->dict_)
Mul::dict_add_term_new(outArg(coef), d, p.second, p.first);
} else if (SymEngine::is_a<Mul>(*a)) {
RCP<const Basic> exp;
RCP<const Basic> t;
coef = (rcp_static_cast<const Mul>(a))->coef_;
d = (rcp_static_cast<const Mul>(a))->dict_;
if (is_a_Number(*b)) {
imulnum(outArg(coef), rcp_static_cast<const Number>(b));
} else {
Mul::as_base_exp(b, outArg(exp), outArg(t));
Mul::dict_add_term_new(outArg(coef), d, exp, t);
}
} else if (SymEngine::is_a<Mul>(*b)) {
RCP<const Basic> exp;
RCP<const Basic> t;
coef = (rcp_static_cast<const Mul>(b))->coef_;
d = (rcp_static_cast<const Mul>(b))->dict_;
if (is_a_Number(*a)) {
imulnum(outArg(coef), rcp_static_cast<const Number>(a));
} else {
Mul::as_base_exp(a, outArg(exp), outArg(t));
Mul::dict_add_term_new(outArg(coef), d, exp, t);
}
} else {
RCP<const Basic> exp;
RCP<const Basic> t;
if (is_a_Number(*a)) {
imulnum(outArg(coef), rcp_static_cast<const Number>(a));
} else {
Mul::as_base_exp(a, outArg(exp), outArg(t));
Mul::dict_add_term_new(outArg(coef), d, exp, t);
}
if (is_a_Number(*b)) {
imulnum(outArg(coef), rcp_static_cast<const Number>(b));
} else {
Mul::as_base_exp(b, outArg(exp), outArg(t));
Mul::dict_add_term_new(outArg(coef), d, exp, t);
}
}
return Mul::from_dict(coef, std::move(d));
}
void bvisit(const Number &x)
{
if (not is_a_Number(*pow(the_base, x.rcp_from_this()))) {
if (x.is_positive())
gen_set[one] = x.rcp_from_this_cast<const Number>();
else
gen_set[minus_one]
= mulnum(x.rcp_from_this_cast<const Number>(), minus_one);
}
}
void Add::coef_dict_add_term(const Ptr<RCP<const Number>> &coef, umap_basic_num &d,
const RCP<const Number> &c, const RCP<const Basic> &term)
{
if (is_a_Number(*term)) {
iaddnum(coef, mulnum(c, rcp_static_cast<const Number>(term)));
} else if (is_a<Add>(*term)) {
if (c->is_one()) {
for (const auto &q: (rcp_static_cast<const Add>(term))->dict_)
Add::dict_add_term(d, q.second, q.first);
iaddnum(coef, rcp_static_cast<const Add>(term)->coef_);
} else {
Add::dict_add_term(d, c, term);
}
} else {
RCP<const Number> coef2;
RCP<const Basic> t;
Add::as_coef_term(term, outArg(coef2), outArg(t));
Add::dict_add_term(d, mulnum(c, coef2), t);
}
}
RCP<const Number> Complex::powcomp(const Integer &other) const {
if (this->is_re_zero()) {
// Imaginary Number raised to an integer power.
RCP<const Number> im = Rational::from_mpq(this->imaginary_);
long rem = mod(other, *integer(4))->as_int();
RCP<const Number> res;
if (rem == 0) {
res = one;
} else if (rem == 1) {
res = I;
} else if (rem == 2) {
res = minus_one;
} else {
res = mulnum(I, minus_one);
}
return mulnum(im->pow(other), res);
} else if (other.is_positive()) {
return pow_number(*this, other.as_int());
} else {
return one->div(*pow_number(*this, -1 * other.as_int()));
}
};
RCP<const Basic> add_expand(const RCP<const Add> &self)
{
umap_basic_num d;
RCP<const Number> coef_overall = self->coef_;
RCP<const Number> coef;
RCP<const Basic> tmp, tmp2;
for (auto &p: self->dict_) {
tmp = expand(p.first);
if (is_a<Add>(*tmp)) {
for (auto &q: (rcp_static_cast<const Add>(tmp))->dict_) {
Add::as_coef_term(q.first, outArg(coef), outArg(tmp2));
Add::dict_add_term(d,
mulnum(mulnum(p.second, q.second), coef), tmp2);
}
iaddnum(outArg(coef_overall), mulnum(p.second,
rcp_static_cast<const Add>(tmp)->coef_));
} else {
Add::as_coef_term(tmp, outArg(coef), outArg(tmp));
Add::dict_add_term(d, mulnum(p.second, coef), tmp);
}
}
return Add::from_dict(coef_overall, std::move(d));
}
static RCP<const Basic> diff(const Add &self,
const RCP<const Symbol> &x) {
SymEngine::umap_basic_num d;
RCP<const Number> coef=zero, coef2;
RCP<const Basic> t;
for (auto &p: self.dict_) {
RCP<const Basic> term = p.first->diff(x);
if (is_a<Integer>(*term) && rcp_static_cast<const Integer>(term)->is_zero()) {
continue;
} else if (is_a_Number(*term)) {
iaddnum(outArg(coef),
mulnum(p.second, rcp_static_cast<const Number>(term)));
} else if (is_a<Add>(*term)) {
for (auto &q: (rcp_static_cast<const Add>(term))->dict_)
Add::dict_add_term(d, mulnum(q.second, p.second), q.first);
iaddnum(outArg(coef), mulnum(p.second, rcp_static_cast<const Add>(term)->coef_));
} else {
Add::as_coef_term(mul(p.second, term), outArg(coef2), outArg(t));
Add::dict_add_term(d, coef2, t);
}
}
return Add::from_dict(coef, std::move(d));
}
RCP<Basic> Add::diff(const RCP<Symbol> &x) const
{
CSymPy::umap_basic_int d;
RCP<Number> coef=zero, coef2;
RCP<Basic> t;
for (auto &p: dict_) {
RCP<Basic> term = p.first->diff(x);
if (is_a<Integer>(*term) && rcp_static_cast<Integer>(term)->is_zero()) {
continue;
} else if (is_a_Number(*term)) {
iaddnum(outArg(coef),
mulnum(p.second, rcp_static_cast<Number>(term)));
} else if (is_a<Add>(*term)) {
for (auto &q: (rcp_static_cast<Add>(term))->dict_)
Add::dict_add_term(d, q.second, q.first);
iaddnum(outArg(coef), rcp_static_cast<Add>(term)->coef_);
} else {
Add::as_coef_term(mul(p.second, term), outArg(coef2), outArg(t));
Add::dict_add_term(d, coef2, t);
}
}
return Add::from_dict(coef, d);
}
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 (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) &&
(rcp_static_cast<const Rational>(b)->i.get_num() == 1) &&
(rcp_static_cast<const Rational>(b)->i.get_den() == 2)) {
return I;
}
}
if (is_a_Number(*a) && 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 {
throw std::runtime_error("Not implemented");
}
} else if (is_a<Rational>(*b)) {
mpz_class q, r, num, den;
num = rcp_static_cast<const Rational>(b)->i.get_num();
den = rcp_static_cast<const Rational>(b)->i.get_den();
if (num > den || num < 0) {
mpz_fdiv_qr(q.get_mpz_t(), r.get_mpz_t(), num.get_mpz_t(),
den.get_mpz_t());
} else {
return rcp(new Pow(a, b));
}
// Here we make the exponent postive and a fraction between
// 0 and 1. We multiply numerator and denominator appropriately
// to achieve this
if (is_a<Rational>(*a)) {
RCP<const Rational> exp_new = rcp_static_cast<const Rational>(a);
RCP<const Basic> frac =
div(exp_new->powrat(Integer(q)), integer(exp_new->i.get_den()));
RCP<const Basic> surds =
mul(rcp(new Pow(integer(exp_new->i.get_num()), div(integer(r), integer(den)))),
rcp(new Pow(integer(exp_new->i.get_den()), sub(one, div(integer(r), integer(den))))));
return mul(frac, surds);
} else if (is_a<Integer>(*a)) {
RCP<const Integer> exp_new = rcp_static_cast<const Integer>(a);
RCP<const Number> frac = exp_new->powint(Integer(q));
map_basic_basic surd;
if ((exp_new->is_negative()) && (2 * r == den)) {
frac = mulnum(frac, I);
exp_new = exp_new->mulint(*minus_one);
// if exp_new is one, no need to add it to dict
if (exp_new->is_one())
return frac;
surd[exp_new] = div(integer(r), integer(den));
} else {
surd[exp_new] = div(integer(r), integer(den));
}
return rcp(new Mul(frac, std::move(surd)));
} else if (is_a<Complex>(*a)) {
return rcp(new Pow(a, b));
} else {
throw std::runtime_error("Not implemented");
}
} else if (is_a<Complex>(*b)) {
return rcp(new Pow(a, b));
} else {
throw std::runtime_error("Not implemented");
}
}
if (is_a<Mul>(*a) && is_a<Integer>(*b)) {
// Convert (x*y)^b = x^b*y^b, where 'b' is an integer. This holds for
// any complex 'x', 'y' and integer 'b'.
return rcp_static_cast<const Mul>(a)->power_all_terms(b);
}
if (is_a<Pow>(*a) && 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->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;
}
RCP<const Basic> mul_expand_two(const RCP<const Basic> &a, const RCP<const Basic> &b)
{
// Both a and b are assumed to be expanded
if (is_a<Add>(*a) && is_a<Add>(*b)) {
RCP<const Number> coef = mulnum(rcp_static_cast<const Add>(a)->coef_,
rcp_static_cast<const Add>(b)->coef_);
umap_basic_num d;
// Improves (x+1)^3(x+2)^3...(x+350)^3 expansion from 0.97s to 0.93s:
d.reserve((rcp_static_cast<const Add>(a))->dict_.size()*
(rcp_static_cast<const Add>(b))->dict_.size());
// Expand dicts first:
for (auto &p: (rcp_static_cast<const Add>(a))->dict_) {
for (auto &q: (rcp_static_cast<const Add>(b))->dict_) {
// The main bottleneck here is the mul(p.first, q.first) command
RCP<const Basic> term = mul(p.first, q.first);
if (is_a_Number(*term)) {
iaddnum(outArg(coef), rcp_static_cast<const Number>(term));
} else {
Add::dict_add_term(d, mulnum(p.second, q.second), term);
}
}
Add::dict_add_term(d,
mulnum(rcp_static_cast<const Add>(b)->coef_, p.second),
p.first);
}
// Handle the coefficient of "a":
for (auto &q: (rcp_static_cast<const Add>(b))->dict_) {
Add::dict_add_term(d,
mulnum(rcp_static_cast<const Add>(a)->coef_, q.second),
q.first);
}
return Add::from_dict(coef, std::move(d));
} else if (is_a<Add>(*a)) {
return mul_expand_two(b, a);
} else if (is_a<Add>(*b)) {
RCP<const Number> a_coef;
RCP<const Basic> a_term;
Add::as_coef_term(a, outArg(a_coef), outArg(a_term));
RCP<const Number> coef = zero;
umap_basic_num d;
d.reserve((rcp_static_cast<const Add>(b))->dict_.size());
for (auto &q: (rcp_static_cast<const Add>(b))->dict_) {
RCP<const Basic> term = mul(a_term, q.first);
if (is_a_Number(*term)) {
iaddnum(outArg(coef), rcp_static_cast<const Number>(term));
} else {
Add::dict_add_term(d, mulnum(a_coef, q.second), term);
}
}
if (eq(a_term, one)) {
iaddnum(outArg(coef),
mulnum(rcp_static_cast<const Add>(b)->coef_, a_coef));
} else {
Add::dict_add_term(d,
mulnum(rcp_static_cast<const Add>(b)->coef_, a_coef),
a_term);
}
return Add::from_dict(coef, std::move(d));
}
return mul(a, b);
}
RCP<const Basic> mul_expand_two(const RCP<const Basic> &a, const RCP<const Basic> &b)
{
// Both a and b are assumed to be expanded
if (is_a<Add>(*a) && is_a<Add>(*b)) {
RCP<const Number> coef = mulnum(rcp_static_cast<const Add>(a)->coef_,
rcp_static_cast<const Add>(b)->coef_);
umap_basic_num d;
// Improves (x+1)**3*(x+2)**3*...(x+350)**3 expansion from 0.97s to 0.93s:
d.reserve((rcp_static_cast<const Add>(a))->dict_.size()*
(rcp_static_cast<const Add>(b))->dict_.size());
// Expand dicts first:
for (auto &p: (rcp_static_cast<const Add>(a))->dict_) {
for (auto &q: (rcp_static_cast<const Add>(b))->dict_) {
// The main bottleneck here is the mul(p.first, q.first) command
RCP<const Basic> term = mul(p.first, q.first);
if (is_a_Number(*term)) {
iaddnum(outArg(coef),
mulnum(mulnum(p.second, q.second), rcp_static_cast<const Number>(term)));
} else {
if (is_a<Mul>(*term) &&
!(rcp_static_cast<const Mul>(term)->coef_->is_one())) {
// Tidy up things like {2x: 3} -> {x: 6}
RCP<const Number> 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(d, mulnum(mulnum(p.second, q.second), coef2), term);
} else {
Add::dict_add_term(d, mulnum(p.second, q.second), term);
}
}
}
Add::dict_add_term(d,
mulnum(rcp_static_cast<const Add>(b)->coef_, p.second),
p.first);
}
// Handle the coefficient of "a":
for (auto &q: (rcp_static_cast<const Add>(b))->dict_) {
Add::dict_add_term(d,
mulnum(rcp_static_cast<const Add>(a)->coef_, q.second),
q.first);
}
return Add::from_dict(coef, std::move(d));
} else if (is_a<Add>(*a)) {
return mul_expand_two(b, a);
} else if (is_a<Add>(*b)) {
RCP<const Number> a_coef;
RCP<const Basic> a_term;
Add::as_coef_term(a, outArg(a_coef), outArg(a_term));
RCP<const Number> coef = zero;
umap_basic_num d;
d.reserve((rcp_static_cast<const Add>(b))->dict_.size());
for (auto &q: (rcp_static_cast<const Add>(b))->dict_) {
RCP<const Basic> term = mul(a_term, q.first);
if (is_a_Number(*term)) {
iaddnum(outArg(coef), mulnum(mulnum(q.second, a_coef),
rcp_static_cast<const Number>(term)));
} else {
if (is_a<Mul>(*term) &&
!(rcp_static_cast<const Mul>(term)->coef_->is_one())) {
// Tidy up things like {2x: 3} -> {x: 6}
RCP<const Number> 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(d, mulnum(mulnum(q.second, a_coef), coef2), term);
} else {
Add::dict_add_term(d, mulnum(a_coef, q.second), term);
}
}
}
if (eq(*a_term, *one)) {
iaddnum(outArg(coef),
mulnum(rcp_static_cast<const Add>(b)->coef_, a_coef));
} else {
Add::dict_add_term(d,
mulnum(rcp_static_cast<const Add>(b)->coef_, a_coef),
a_term);
}
return Add::from_dict(coef, std::move(d));
}
return mul(a, b);
}
