Loading...
Searching...
No Matches
expand.cpp
1#include <symengine/visitor.h>
2
3namespace SymEngine
4{
5
6inline 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
16inline void _imulnum(const Ptr<RCP<const Number>> &self,
17 const RCP<const Number> &other)
18{
19 *self = _mulnum(*self, other);
20}
21
22class ExpandVisitor : public BaseVisitor<ExpandVisitor>
23{
24private:
26 RCP<const Number> coeff = zero;
27 RCP<const Number> multiply = one;
28 bool deep;
29
30public:
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_:
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_:
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();
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_:
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
357private:
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
369RCP<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
const RCP< const Number > & get_coef() const
Definition: add.h:142
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:97
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)
void as_two_terms(const Ptr< RCP< const Basic > > &a, const Ptr< RCP< const Basic > > &b) const
Rewrite as 2 terms.
Definition: mul.cpp:309
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:320
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:150
static RCP< const Basic > from_dict(const RCP< const Number > &coef, map_basic_basic &&d)
Create a Mul from a dict.
Definition: mul.cpp:115
RCP< const Basic > get_base() const
Definition: pow.h:37
RCP< const Basic > get_exp() const
Definition: pow.h:42
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:431
bool eq(const Basic &a, const Basic &b)
Checks equality for a and b
Definition: basic-inl.h:21
RCP< const Number > pownum(const RCP< const Number > &self, const RCP< const Number > &other)
Raise self to power other
Definition: number.h:105
RCP< const Basic > exp(const RCP< const Basic > &x)
Returns the natural exponential function E**x = pow(E, x)
Definition: pow.cpp:271
RCP< const Basic > mul(const RCP< const Basic > &a, const RCP< const Basic > &b)
Multiplication.
Definition: mul.cpp:352
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
std::enable_if< std::is_integral< T >::value, RCP< constInteger > >::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
T reserve(T... args)
T size(T... args)