utils.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323
  1. from sympy.core.containers import Tuple
  2. from sympy.core.basic import Basic
  3. from sympy.core.expr import Expr
  4. from sympy.core.function import AppliedUndef
  5. from sympy.core.relational import Relational
  6. from sympy.core.symbol import Dummy
  7. from sympy.core.sympify import sympify
  8. from sympy.logic.boolalg import BooleanFunction
  9. from sympy.sets.fancysets import ImageSet
  10. from sympy.sets.sets import FiniteSet
  11. from sympy.tensor.indexed import Indexed
  12. def _get_free_symbols(exprs):
  13. """Returns the free symbols of a symbolic expression.
  14. If the expression contains any of these elements, assume that they are
  15. the "free symbols" of the expression:
  16. * indexed objects
  17. * applied undefined function (useful for sympy.physics.mechanics module)
  18. """
  19. if not isinstance(exprs, (list, tuple, set)):
  20. exprs = [exprs]
  21. if all(callable(e) for e in exprs):
  22. return set()
  23. free = set().union(*[e.atoms(Indexed) for e in exprs])
  24. free = free.union(*[e.atoms(AppliedUndef) for e in exprs])
  25. return free or set().union(*[e.free_symbols for e in exprs])
  26. def extract_solution(set_sol, n=10):
  27. """Extract numerical solutions from a set solution (computed by solveset,
  28. linsolve, nonlinsolve). Often, it is not trivial do get something useful
  29. out of them.
  30. Parameters
  31. ==========
  32. n : int, optional
  33. In order to replace ImageSet with FiniteSet, an iterator is created
  34. for each ImageSet contained in `set_sol`, starting from 0 up to `n`.
  35. Default value: 10.
  36. """
  37. images = set_sol.find(ImageSet)
  38. for im in images:
  39. it = iter(im)
  40. s = FiniteSet(*[next(it) for n in range(0, n)])
  41. set_sol = set_sol.subs(im, s)
  42. return set_sol
  43. def _plot_sympify(args):
  44. """This function recursively loop over the arguments passed to the plot
  45. functions: the sympify function will be applied to all arguments except
  46. those of type string/dict.
  47. Generally, users can provide the following arguments to a plot function:
  48. expr, range1 [tuple, opt], ..., label [str, opt], rendering_kw [dict, opt]
  49. `expr, range1, ...` can be sympified, whereas `label, rendering_kw` can't.
  50. In particular, whenever a special character like $, {, }, ... is used in
  51. the `label`, sympify will raise an error.
  52. """
  53. if isinstance(args, Expr):
  54. return args
  55. args = list(args)
  56. for i, a in enumerate(args):
  57. if isinstance(a, (list, tuple)):
  58. args[i] = Tuple(*_plot_sympify(a), sympify=False)
  59. elif not (isinstance(a, (str, dict)) or callable(a)
  60. # NOTE: check if it is a vector from sympy.physics.vector module
  61. # without importing the module (because it slows down SymPy's
  62. # import process and triggers SymPy's optional-dependencies
  63. # tests to fail).
  64. or ((a.__class__.__name__ == "Vector") and not isinstance(a, Basic))
  65. ):
  66. args[i] = sympify(a)
  67. return args
  68. def _create_ranges(exprs, ranges, npar, label="", params=None):
  69. """This function does two things:
  70. 1. Check if the number of free symbols is in agreement with the type of
  71. plot chosen. For example, plot() requires 1 free symbol;
  72. plot3d() requires 2 free symbols.
  73. 2. Sometime users create plots without providing ranges for the variables.
  74. Here we create the necessary ranges.
  75. Parameters
  76. ==========
  77. exprs : iterable
  78. The expressions from which to extract the free symbols
  79. ranges : iterable
  80. The limiting ranges provided by the user
  81. npar : int
  82. The number of free symbols required by the plot functions.
  83. For example,
  84. npar=1 for plot, npar=2 for plot3d, ...
  85. params : dict
  86. A dictionary mapping symbols to parameters for interactive plot.
  87. """
  88. get_default_range = lambda symbol: Tuple(symbol, -10, 10)
  89. free_symbols = _get_free_symbols(exprs)
  90. if params is not None:
  91. free_symbols = free_symbols.difference(params.keys())
  92. if len(free_symbols) > npar:
  93. raise ValueError(
  94. "Too many free symbols.\n"
  95. + "Expected {} free symbols.\n".format(npar)
  96. + "Received {}: {}".format(len(free_symbols), free_symbols)
  97. )
  98. if len(ranges) > npar:
  99. raise ValueError(
  100. "Too many ranges. Received %s, expected %s" % (len(ranges), npar))
  101. # free symbols in the ranges provided by the user
  102. rfs = set().union([r[0] for r in ranges])
  103. if len(rfs) != len(ranges):
  104. raise ValueError("Multiple ranges with the same symbol")
  105. if len(ranges) < npar:
  106. symbols = free_symbols.difference(rfs)
  107. if symbols != set():
  108. # add a range for each missing free symbols
  109. for s in symbols:
  110. ranges.append(get_default_range(s))
  111. # if there is still room, fill them with dummys
  112. for i in range(npar - len(ranges)):
  113. ranges.append(get_default_range(Dummy()))
  114. if len(free_symbols) == npar:
  115. # there could be times when this condition is not met, for example
  116. # plotting the function f(x, y) = x (which is a plane); in this case,
  117. # free_symbols = {x} whereas rfs = {x, y} (or x and Dummy)
  118. rfs = set().union([r[0] for r in ranges])
  119. if len(free_symbols.difference(rfs)) > 0:
  120. raise ValueError(
  121. "Incompatible free symbols of the expressions with "
  122. "the ranges.\n"
  123. + "Free symbols in the expressions: {}\n".format(free_symbols)
  124. + "Free symbols in the ranges: {}".format(rfs)
  125. )
  126. return ranges
  127. def _is_range(r):
  128. """A range is defined as (symbol, start, end). start and end should
  129. be numbers.
  130. """
  131. # TODO: prange check goes here
  132. return (
  133. isinstance(r, Tuple)
  134. and (len(r) == 3)
  135. and (not isinstance(r.args[1], str)) and r.args[1].is_number
  136. and (not isinstance(r.args[2], str)) and r.args[2].is_number
  137. )
  138. def _unpack_args(*args):
  139. """Given a list/tuple of arguments previously processed by _plot_sympify()
  140. and/or _check_arguments(), separates and returns its components:
  141. expressions, ranges, label and rendering keywords.
  142. Examples
  143. ========
  144. >>> from sympy import cos, sin, symbols
  145. >>> from sympy.plotting.utils import _plot_sympify, _unpack_args
  146. >>> x, y = symbols('x, y')
  147. >>> args = (sin(x), (x, -10, 10), "f1")
  148. >>> args = _plot_sympify(args)
  149. >>> _unpack_args(*args)
  150. ([sin(x)], [(x, -10, 10)], 'f1', None)
  151. >>> args = (sin(x**2 + y**2), (x, -2, 2), (y, -3, 3), "f2")
  152. >>> args = _plot_sympify(args)
  153. >>> _unpack_args(*args)
  154. ([sin(x**2 + y**2)], [(x, -2, 2), (y, -3, 3)], 'f2', None)
  155. >>> args = (sin(x + y), cos(x - y), x + y, (x, -2, 2), (y, -3, 3), "f3")
  156. >>> args = _plot_sympify(args)
  157. >>> _unpack_args(*args)
  158. ([sin(x + y), cos(x - y), x + y], [(x, -2, 2), (y, -3, 3)], 'f3', None)
  159. """
  160. ranges = [t for t in args if _is_range(t)]
  161. labels = [t for t in args if isinstance(t, str)]
  162. label = None if not labels else labels[0]
  163. rendering_kw = [t for t in args if isinstance(t, dict)]
  164. rendering_kw = None if not rendering_kw else rendering_kw[0]
  165. # NOTE: why None? because args might have been preprocessed by
  166. # _check_arguments, so None might represent the rendering_kw
  167. results = [not (_is_range(a) or isinstance(a, (str, dict)) or (a is None)) for a in args]
  168. exprs = [a for a, b in zip(args, results) if b]
  169. return exprs, ranges, label, rendering_kw
  170. def _check_arguments(args, nexpr, npar, **kwargs):
  171. """Checks the arguments and converts into tuples of the
  172. form (exprs, ranges, label, rendering_kw).
  173. Parameters
  174. ==========
  175. args
  176. The arguments provided to the plot functions
  177. nexpr
  178. The number of sub-expression forming an expression to be plotted.
  179. For example:
  180. nexpr=1 for plot.
  181. nexpr=2 for plot_parametric: a curve is represented by a tuple of two
  182. elements.
  183. nexpr=1 for plot3d.
  184. nexpr=3 for plot3d_parametric_line: a curve is represented by a tuple
  185. of three elements.
  186. npar
  187. The number of free symbols required by the plot functions. For example,
  188. npar=1 for plot, npar=2 for plot3d, ...
  189. **kwargs :
  190. keyword arguments passed to the plotting function. It will be used to
  191. verify if ``params`` has ben provided.
  192. Examples
  193. ========
  194. .. plot::
  195. :context: reset
  196. :format: doctest
  197. :include-source: True
  198. >>> from sympy import cos, sin, symbols
  199. >>> from sympy.plotting.plot import _check_arguments
  200. >>> x = symbols('x')
  201. >>> _check_arguments([cos(x), sin(x)], 2, 1)
  202. [(cos(x), sin(x), (x, -10, 10), None, None)]
  203. >>> _check_arguments([cos(x), sin(x), "test"], 2, 1)
  204. [(cos(x), sin(x), (x, -10, 10), 'test', None)]
  205. >>> _check_arguments([cos(x), sin(x), "test", {"a": 0, "b": 1}], 2, 1)
  206. [(cos(x), sin(x), (x, -10, 10), 'test', {'a': 0, 'b': 1})]
  207. >>> _check_arguments([x, x**2], 1, 1)
  208. [(x, (x, -10, 10), None, None), (x**2, (x, -10, 10), None, None)]
  209. """
  210. if not args:
  211. return []
  212. output = []
  213. params = kwargs.get("params", None)
  214. if all(isinstance(a, (Expr, Relational, BooleanFunction)) for a in args[:nexpr]):
  215. # In this case, with a single plot command, we are plotting either:
  216. # 1. one expression
  217. # 2. multiple expressions over the same range
  218. exprs, ranges, label, rendering_kw = _unpack_args(*args)
  219. free_symbols = set().union(*[e.free_symbols for e in exprs])
  220. ranges = _create_ranges(exprs, ranges, npar, label, params)
  221. if nexpr > 1:
  222. # in case of plot_parametric or plot3d_parametric_line, there will
  223. # be 2 or 3 expressions defining a curve. Group them together.
  224. if len(exprs) == nexpr:
  225. exprs = (tuple(exprs),)
  226. for expr in exprs:
  227. # need this if-else to deal with both plot/plot3d and
  228. # plot_parametric/plot3d_parametric_line
  229. is_expr = isinstance(expr, (Expr, Relational, BooleanFunction))
  230. e = (expr,) if is_expr else expr
  231. output.append((*e, *ranges, label, rendering_kw))
  232. else:
  233. # In this case, we are plotting multiple expressions, each one with its
  234. # range. Each "expression" to be plotted has the following form:
  235. # (expr, range, label) where label is optional
  236. _, ranges, labels, rendering_kw = _unpack_args(*args)
  237. labels = [labels] if labels else []
  238. # number of expressions
  239. n = (len(ranges) + len(labels) +
  240. (len(rendering_kw) if rendering_kw is not None else 0))
  241. new_args = args[:-n] if n > 0 else args
  242. # at this point, new_args might just be [expr]. But I need it to be
  243. # [[expr]] in order to be able to loop over
  244. # [expr, range [opt], label [opt]]
  245. if not isinstance(new_args[0], (list, tuple, Tuple)):
  246. new_args = [new_args]
  247. # Each arg has the form (expr1, expr2, ..., range1 [optional], ...,
  248. # label [optional], rendering_kw [optional])
  249. for arg in new_args:
  250. # look for "local" range and label. If there is not, use "global".
  251. l = [a for a in arg if isinstance(a, str)]
  252. if not l:
  253. l = labels
  254. r = [a for a in arg if _is_range(a)]
  255. if not r:
  256. r = ranges.copy()
  257. rend_kw = [a for a in arg if isinstance(a, dict)]
  258. rend_kw = rendering_kw if len(rend_kw) == 0 else rend_kw[0]
  259. # NOTE: arg = arg[:nexpr] may raise an exception if lambda
  260. # functions are used. Execute the following instead:
  261. arg = [arg[i] for i in range(nexpr)]
  262. free_symbols = set()
  263. if all(not callable(a) for a in arg):
  264. free_symbols = free_symbols.union(*[a.free_symbols for a in arg])
  265. if len(r) != npar:
  266. r = _create_ranges(arg, r, npar, "", params)
  267. label = None if not l else l[0]
  268. output.append((*arg, *r, label, rend_kw))
  269. return output