textplot.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162
  1. from sympy.core.numbers import Float
  2. from sympy.core.symbol import Dummy
  3. from sympy.utilities.lambdify import lambdify
  4. import math
  5. def is_valid(x):
  6. """Check if a floating point number is valid"""
  7. if x is None:
  8. return False
  9. if isinstance(x, complex):
  10. return False
  11. return not math.isinf(x) and not math.isnan(x)
  12. def rescale(y, W, H, mi, ma):
  13. """Rescale the given array `y` to fit into the integer values
  14. between `0` and `H-1` for the values between ``mi`` and ``ma``.
  15. """
  16. y_new = list()
  17. norm = ma - mi
  18. offset = (ma + mi) / 2
  19. for x in range(W):
  20. if is_valid(y[x]):
  21. normalized = (y[x] - offset) / norm
  22. if not is_valid(normalized):
  23. y_new.append(None)
  24. else:
  25. rescaled = Float((normalized*H + H/2) * (H-1)/H).round()
  26. rescaled = int(rescaled)
  27. y_new.append(rescaled)
  28. else:
  29. y_new.append(None)
  30. return y_new
  31. def linspace(start, stop, num):
  32. return [start + (stop - start) * x / (num-1) for x in range(num)]
  33. def textplot_str(expr, a, b, W=55, H=21):
  34. """Generator for the lines of the plot"""
  35. free = expr.free_symbols
  36. if len(free) > 1:
  37. raise ValueError(
  38. "The expression must have a single variable. (Got {})"
  39. .format(free))
  40. x = free.pop() if free else Dummy()
  41. f = lambdify([x], expr)
  42. a = float(a)
  43. b = float(b)
  44. # Calculate function values
  45. x = linspace(a, b, W)
  46. y = list()
  47. for val in x:
  48. try:
  49. y.append(f(val))
  50. # Not sure what exceptions to catch here or why...
  51. except (ValueError, TypeError, ZeroDivisionError):
  52. y.append(None)
  53. # Normalize height to screen space
  54. y_valid = list(filter(is_valid, y))
  55. if y_valid:
  56. ma = max(y_valid)
  57. mi = min(y_valid)
  58. if ma == mi:
  59. if ma:
  60. mi, ma = sorted([0, 2*ma])
  61. else:
  62. mi, ma = -1, 1
  63. else:
  64. mi, ma = -1, 1
  65. y_range = ma - mi
  66. precision = math.floor(math.log(y_range, 10)) - 1
  67. precision *= -1
  68. mi = round(mi, precision)
  69. ma = round(ma, precision)
  70. y = rescale(y, W, H, mi, ma)
  71. y_bins = linspace(mi, ma, H)
  72. # Draw plot
  73. margin = 7
  74. for h in range(H - 1, -1, -1):
  75. s = [' '] * W
  76. for i in range(W):
  77. if y[i] == h:
  78. if (i == 0 or y[i - 1] == h - 1) and (i == W - 1 or y[i + 1] == h + 1):
  79. s[i] = '/'
  80. elif (i == 0 or y[i - 1] == h + 1) and (i == W - 1 or y[i + 1] == h - 1):
  81. s[i] = '\\'
  82. else:
  83. s[i] = '.'
  84. if h == 0:
  85. for i in range(W):
  86. s[i] = '_'
  87. # Print y values
  88. if h in (0, H//2, H - 1):
  89. prefix = ("%g" % y_bins[h]).rjust(margin)[:margin]
  90. else:
  91. prefix = " "*margin
  92. s = "".join(s)
  93. if h == H//2:
  94. s = s.replace(" ", "-")
  95. yield prefix + " |" + s
  96. # Print x values
  97. bottom = " " * (margin + 2)
  98. bottom += ("%g" % x[0]).ljust(W//2)
  99. if W % 2 == 1:
  100. bottom += ("%g" % x[W//2]).ljust(W//2)
  101. else:
  102. bottom += ("%g" % x[W//2]).ljust(W//2-1)
  103. bottom += "%g" % x[-1]
  104. yield bottom
  105. def textplot(expr, a, b, W=55, H=21):
  106. r"""
  107. Print a crude ASCII art plot of the SymPy expression 'expr' (which
  108. should contain a single symbol, e.g. x or something else) over the
  109. interval [a, b].
  110. Examples
  111. ========
  112. >>> from sympy import Symbol, sin
  113. >>> from sympy.plotting import textplot
  114. >>> t = Symbol('t')
  115. >>> textplot(sin(t)*t, 0, 15)
  116. 14 | ...
  117. | .
  118. | .
  119. | .
  120. | .
  121. | ...
  122. | / . .
  123. | /
  124. | / .
  125. | . . .
  126. 1.5 |----.......--------------------------------------------
  127. |.... \ . .
  128. | \ / .
  129. | .. / .
  130. | \ / .
  131. | ....
  132. | .
  133. | . .
  134. |
  135. | . .
  136. -11 |_______________________________________________________
  137. 0 7.5 15
  138. """
  139. for line in textplot_str(expr, a, b, W, H):
  140. print(line)