expand.cpp
1 #include <symengine/visitor.h>
2 
3 namespace SymEngine
4 {
5 
6 inline RCP<const Number> _mulnum(const RCP<const Number> &x,
7  const RCP<const Number> &y)
8 {
9  if (eq(*x, *one))
10  return y;
11  if (eq(*y, *one))
12  return x;
13  return x->mul(*y);
14 }
15 
16 inline void _imulnum(const Ptr<RCP<const Number>> &self,
17  const RCP<const Number> &other)
18 {
19  *self = _mulnum(*self, other);
20 }
21 
22 class ExpandVisitor : public BaseVisitor<ExpandVisitor>
23 {
24 private:
25  umap_basic_num d_;
26  RCP<const Number> coeff = zero;
27  RCP<const Number> multiply = one;
28  bool deep;
29 
30 public:
31  ExpandVisitor(bool deep_ = true) : deep(deep_) {}
32  RCP<const Basic> apply(const Basic &b)
33  {
34  b.accept(*this);
35  return Add::from_dict(coeff, std::move(d_));
36  }
37 
38  void bvisit(const Basic &x)
39  {
40  Add::dict_add_term(d_, multiply, x.rcp_from_this());
41  }
42 
43  void bvisit(const Number &x)
44  {
45  iaddnum(outArg(coeff),
46  _mulnum(multiply, x.rcp_from_this_cast<const Number>()));
47  }
48 
49  void bvisit(const Add &self)
50  {
51  RCP<const Number> _multiply = multiply;
52  iaddnum(outArg(coeff), _mulnum(multiply, self.get_coef()));
53  for (auto &p : self.get_dict()) {
54  multiply = _mulnum(_multiply, p.second);
55  if (deep) {
56  p.first->accept(*this);
57  } else {
58  Add::dict_add_term(d_, multiply, p.first);
59  }
60  }
61  multiply = _multiply;
62  }
63 
64  void bvisit(const Mul &self)
65  {
66  for (auto &p : self.get_dict()) {
67  if (!is_a<Symbol>(*p.first)) {
68  RCP<const Basic> a, b;
69  self.as_two_terms(outArg(a), outArg(b));
70  a = expand_if_deep(a);
71  b = expand_if_deep(b);
72  mul_expand_two(a, b);
73  return;
74  }
75  }
76  this->_coef_dict_add_term(multiply, self.rcp_from_this());
77  }
78 
79  void mul_expand_two(const RCP<const Basic> &a, const RCP<const Basic> &b)
80  {
81  // Both a and b are assumed to be expanded
82  if (is_a<Add>(*a) && is_a<Add>(*b)) {
83  iaddnum(outArg(coeff),
84  _mulnum(multiply,
85  _mulnum(down_cast<const Add &>(*a).get_coef(),
86  down_cast<const Add &>(*b).get_coef())));
87 // Improves (x+1)**3*(x+2)**3*...(x+350)**3 expansion from 0.97s to 0.93s:
88 #if defined(HAVE_SYMENGINE_RESERVE)
89  d_.reserve(d_.size()
90  + (down_cast<const Add &>(*a)).get_dict().size()
91  * (down_cast<const Add &>(*b)).get_dict().size());
92 #endif
93  // Expand dicts first:
94  for (auto &p : (down_cast<const Add &>(*a)).get_dict()) {
95  RCP<const Number> temp = _mulnum(p.second, multiply);
96  for (auto &q : (down_cast<const Add &>(*b)).get_dict()) {
97  // The main bottleneck here is the mul(p.first, q.first)
98  // command
99  RCP<const Basic> term = mul(p.first, q.first);
100  if (is_a_Number(*term)) {
101  iaddnum(outArg(coeff),
102  _mulnum(_mulnum(temp, q.second),
103  rcp_static_cast<const Number>(term)));
104  } else {
105  if (is_a<Mul>(*term)
106  && !(down_cast<const Mul &>(*term)
107  .get_coef()
108  ->is_one())) {
109  // Tidy up things like {2x: 3} -> {x: 6}
110  RCP<const Number> coef2
111  = down_cast<const Mul &>(*term).get_coef();
112  // We make a copy of the dict_:
113  map_basic_basic d2
114  = down_cast<const Mul &>(*term).get_dict();
115  term = Mul::from_dict(one, std::move(d2));
117  d_, _mulnum(_mulnum(temp, q.second), coef2),
118  term);
119  } else {
120  Add::dict_add_term(d_, _mulnum(temp, q.second),
121  term);
122  }
123  }
124  }
126  d_, _mulnum(down_cast<const Add &>(*b).get_coef(), temp),
127  p.first);
128  }
129  // Handle the coefficient of "a":
130  RCP<const Number> temp
131  = _mulnum(down_cast<const Add &>(*a).get_coef(), multiply);
132  for (auto &q : (down_cast<const Add &>(*b)).get_dict()) {
133  Add::dict_add_term(d_, _mulnum(temp, q.second), q.first);
134  }
135  return;
136  } else if (is_a<Add>(*a)) {
137  mul_expand_two(b, a);
138  return;
139  } else if (is_a<Add>(*b)) {
140  RCP<const Number> a_coef;
141  RCP<const Basic> a_term;
142  Add::as_coef_term(a, outArg(a_coef), outArg(a_term));
143  _imulnum(outArg(a_coef), multiply);
144 
145 #if defined(HAVE_SYMENGINE_RESERVE)
146  d_.reserve(d_.size()
147  + (down_cast<const Add &>(*b)).get_dict().size());
148 #endif
149  for (auto &q : (down_cast<const Add &>(*b)).get_dict()) {
150  RCP<const Basic> term = mul(a_term, q.first);
151  if (is_a_Number(*term)) {
152  iaddnum(outArg(coeff),
153  _mulnum(_mulnum(q.second, a_coef),
154  rcp_static_cast<const Number>(term)));
155  } else {
156  if (is_a<Mul>(*term)
157  && !(down_cast<const Mul &>(*term)
158  .get_coef()
159  ->is_one())) {
160  // Tidy up things like {2x: 3} -> {x: 6}
161  RCP<const Number> coef2
162  = down_cast<const Mul &>(*term).get_coef();
163  // We make a copy of the dict_:
164  map_basic_basic d2
165  = down_cast<const Mul &>(*term).get_dict();
166  term = Mul::from_dict(one, std::move(d2));
168  d_, _mulnum(_mulnum(q.second, a_coef), coef2),
169  term);
170  } else {
171  // TODO: check if it's a Add
172  Add::dict_add_term(d_, _mulnum(a_coef, q.second), term);
173  }
174  }
175  }
176  if (eq(*a_term, *one)) {
177  iaddnum(outArg(coeff),
178  _mulnum(down_cast<const Add &>(*b).get_coef(), a_coef));
179  } else {
181  d_, _mulnum(down_cast<const Add &>(*b).get_coef(), a_coef),
182  a_term);
183  }
184  return;
185  }
186  _coef_dict_add_term(multiply, mul(a, b));
187  }
188 
189  void square_expand(umap_basic_num &base_dict)
190  {
191  auto m = base_dict.size();
192 #if defined(HAVE_SYMENGINE_RESERVE)
193  d_.reserve(d_.size() + m * (m + 1) / 2);
194 #endif
195  RCP<const Basic> t;
196  RCP<const Number> coef, two = integer(2);
197  for (auto p = base_dict.begin(); p != base_dict.end(); ++p) {
198  for (auto q = p; q != base_dict.end(); ++q) {
199  if (q == p) {
200  _coef_dict_add_term(
201  _mulnum(mulnum((*p).second, (*p).second), multiply),
202  pow((*p).first, two));
203  } else {
204  _coef_dict_add_term(
205  _mulnum(multiply, _mulnum((*p).second,
206  _mulnum((*q).second, two))),
207  mul((*q).first, (*p).first));
208  }
209  }
210  }
211  }
212 
213  void pow_expand(umap_basic_num &base_dict, unsigned n)
214  {
215  map_vec_mpz r;
216  unsigned m = numeric_cast<unsigned>(base_dict.size());
217  multinomial_coefficients_mpz(m, n, r);
218 // This speeds up overall expansion. For example for the benchmark
219 // (y + x + z + w)**60 it improves the timing from 135ms to 124ms.
220 #if defined(HAVE_SYMENGINE_RESERVE)
221  d_.reserve(d_.size() + 2 * r.size());
222 #endif
223  for (auto &p : r) {
224  auto power = p.first.begin();
225  auto i2 = base_dict.begin();
226  map_basic_basic d;
227  RCP<const Number> overall_coeff = one;
228  for (; power != p.first.end(); ++power, ++i2) {
229  if (*power > 0) {
230  RCP<const Integer> exp = integer(std::move(*power));
231  RCP<const Basic> base = i2->first;
232  if (is_a<Integer>(*base)) {
233  _imulnum(outArg(overall_coeff),
234  rcp_static_cast<const Number>(
235  down_cast<const Integer &>(*base).powint(
236  *exp)));
237  } else if (is_a<Symbol>(*base)) {
238  Mul::dict_add_term(d, exp, base);
239  } else {
240  RCP<const Basic> exp2, t, tmp;
241  tmp = pow(base, exp);
242  if (is_a<Mul>(*tmp)) {
243  for (auto &p :
244  (down_cast<const Mul &>(*tmp)).get_dict()) {
245  Mul::dict_add_term_new(outArg(overall_coeff), d,
246  p.second, p.first);
247  }
248  _imulnum(outArg(overall_coeff),
249  (down_cast<const Mul &>(*tmp)).get_coef());
250  } else if (is_a_Number(*tmp)) {
251  _imulnum(outArg(overall_coeff),
252  rcp_static_cast<const Number>(tmp));
253  } else {
254  Mul::as_base_exp(tmp, outArg(exp2), outArg(t));
255  Mul::dict_add_term_new(outArg(overall_coeff), d,
256  exp2, t);
257  }
258  }
259  if (!(i2->second->is_one())) {
260  _imulnum(outArg(overall_coeff),
261  pownum(i2->second,
262  rcp_static_cast<const Number>(exp)));
263  }
264  }
265  }
266  RCP<const Basic> term = Mul::from_dict(overall_coeff, std::move(d));
267  RCP<const Number> coef2 = integer(p.second);
268  if (is_a_Number(*term)) {
269  iaddnum(outArg(coeff),
270  _mulnum(_mulnum(multiply,
271  rcp_static_cast<const Number>(term)),
272  coef2));
273  } else {
274  if (is_a<Mul>(*term)
275  && !(down_cast<const Mul &>(*term).get_coef()->is_one())) {
276  // Tidy up things like {2x: 3} -> {x: 6}
277  _imulnum(outArg(coef2),
278  down_cast<const Mul &>(*term).get_coef());
279  // We make a copy of the dict_:
280  map_basic_basic d2
281  = down_cast<const Mul &>(*term).get_dict();
282  term = Mul::from_dict(one, std::move(d2));
283  }
284  Add::dict_add_term(d_, _mulnum(multiply, coef2), term);
285  }
286  }
287  }
288 
289  void bvisit(const Pow &self)
290  {
291  RCP<const Basic> _base = expand_if_deep(self.get_base());
292  // TODO add all types of polys
293  if (is_a<Integer>(*self.get_exp()) && is_a<UExprPoly>(*_base)) {
294  unsigned q = numeric_cast<unsigned>(
295  down_cast<const Integer &>(*self.get_exp()).as_uint());
296  RCP<const UExprPoly> p = rcp_static_cast<const UExprPoly>(_base);
297  RCP<const UExprPoly> r = pow_upoly(*p, q);
298  _coef_dict_add_term(multiply, r);
299  return;
300  }
301  if (is_a<Integer>(*self.get_exp()) && is_a<UIntPoly>(*_base)) {
302  unsigned q = numeric_cast<unsigned>(
303  down_cast<const Integer &>(*self.get_exp()).as_uint());
304  RCP<const UIntPoly> p = rcp_static_cast<const UIntPoly>(_base);
305  RCP<const UIntPoly> r = pow_upoly(*p, q);
306  _coef_dict_add_term(multiply, r);
307  return;
308  }
309 
310  if (!is_a<Integer>(*self.get_exp()) || !is_a<Add>(*_base)) {
311  if (neq(*_base, *self.get_base())) {
312  Add::dict_add_term(d_, multiply, pow(_base, self.get_exp()));
313  } else {
314  Add::dict_add_term(d_, multiply, self.rcp_from_this());
315  }
316  return;
317  }
318 
319  integer_class n
320  = down_cast<const Integer &>(*self.get_exp()).as_integer_class();
321  if (n < 0)
322  return _coef_dict_add_term(
323  multiply, div(one, expand_if_deep(pow(_base, integer(-n)))));
324  RCP<const Add> base = rcp_static_cast<const Add>(_base);
325  umap_basic_num base_dict = base->get_dict();
326  if (!(base->get_coef()->is_zero())) {
327  // Add the numerical coefficient into the dictionary. This
328  // allows a little bit easier treatment below.
329  insert(base_dict, base->get_coef(), one);
330  } else
331  iaddnum(outArg(coeff), base->get_coef());
332  if (n == 2)
333  return square_expand(base_dict);
334  else
335  return pow_expand(base_dict, numeric_cast<unsigned>(mp_get_ui(n)));
336  }
337 
338  inline void _coef_dict_add_term(const RCP<const Number> &c,
339  const RCP<const Basic> &term)
340  {
341  if (is_a_Number(*term)) {
342  iaddnum(outArg(coeff),
343  _mulnum(c, rcp_static_cast<const Number>(term)));
344  } else if (is_a<Add>(*term)) {
345  for (const auto &q : (down_cast<const Add &>(*term)).get_dict())
346  Add::dict_add_term(d_, _mulnum(q.second, c), q.first);
347  iaddnum(outArg(coeff),
348  _mulnum(down_cast<const Add &>(*term).get_coef(), c));
349  } else {
350  RCP<const Number> coef2;
351  RCP<const Basic> t;
352  Add::as_coef_term(term, outArg(coef2), outArg(t));
353  Add::dict_add_term(d_, _mulnum(c, coef2), t);
354  }
355  }
356 
357 private:
358  RCP<const Basic> expand_if_deep(const RCP<const Basic> &expr)
359  {
360  if (deep) {
361  return expand(expr);
362  } else {
363  return expr;
364  }
365  }
366 };
367 
369 RCP<const Basic> expand(const RCP<const Basic> &self, bool deep)
370 {
371  ExpandVisitor v(deep);
372  return v.apply(*self);
373 }
374 
375 } // namespace SymEngine
T begin(T... args)
The base class for representing addition in symbolic expressions.
Definition: add.h:27
static RCP< const Basic > from_dict(const RCP< const Number > &coef, umap_basic_num &&d)
Create an appropriate instance from dictionary quickly.
Definition: add.cpp:140
static void dict_add_term(umap_basic_num &d, const RCP< const Number > &coef, const RCP< const Basic > &t)
Adds a new term to the expression.
Definition: add.cpp:237
static void as_coef_term(const RCP< const Basic > &self, const Ptr< RCP< const Number >> &coef, const Ptr< RCP< const Basic >> &term)
Converts a Basic self into the form of coefficient * term.
Definition: add.cpp:306
The lowest unit of symbolic representation.
Definition: basic.h:95
RCP< const T2 > rcp_from_this_cast() const
Get RCP<T2> pointer to self (it will cast the pointer to T2)
RCP< T > rcp_from_this()
Get RCP<T> pointer to self (it will cast the pointer to T)
static void as_base_exp(const RCP< const Basic > &self, const Ptr< RCP< const Basic >> &exp, const Ptr< RCP< const Basic >> &base)
Convert to a base and exponent form.
Definition: mul.cpp:315
static void dict_add_term(map_basic_basic &d, const RCP< const Basic > &exp, const RCP< const Basic > &t)
Add terms to dict.
Definition: mul.cpp:151
static RCP< const Basic > from_dict(const RCP< const Number > &coef, map_basic_basic &&d)
Create a Mul from a dict.
Definition: mul.cpp:116
T end(T... args)
T move(T... args)
Main namespace for SymEngine package.
Definition: add.cpp:19
bool is_a_Number(const Basic &b)
Definition: number.h:130
RCP< const Basic > div(const RCP< const Basic > &a, const RCP< const Basic > &b)
Division.
Definition: mul.cpp:426
std::enable_if< std::is_integral< T >::value, RCP< const Integer > >::type integer(T i)
Definition: integer.h:200
RCP< const Number > mulnum(const RCP< const Number > &self, const RCP< const Number > &other)
Multiply self and other
Definition: number.h:93
RCP< const Number > pownum(const RCP< const Number > &self, const RCP< const Number > &other)
Raise self to power other
Definition: number.h:105
bool eq(const Basic &a, const Basic &b)
Checks equality for a and b
Definition: basic-inl.h:21
RCP< const Basic > exp(const RCP< const Basic > &x)
Returns the natural exponential function E**x = pow(E, x)
Definition: pow.cpp:270
RCP< const Basic > mul(const RCP< const Basic > &a, const RCP< const Basic > &b)
Multiplication.
Definition: mul.cpp:347
void insert(T1 &m, const T2 &first, const T3 &second)
Definition: dict.h:83
bool neq(const Basic &a, const Basic &b)
Checks inequality for a and b
Definition: basic-inl.h:29
RCP< const Basic > expand(const RCP< const Basic > &self, bool deep=true)
Expands self
Definition: expand.cpp:369
T reserve(T... args)
T size(T... args)