funcmatrix.py 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118
  1. from .matexpr import MatrixExpr
  2. from sympy.core.function import FunctionClass, Lambda
  3. from sympy.core.symbol import Dummy
  4. from sympy.core.sympify import _sympify, sympify
  5. from sympy.matrices import Matrix
  6. from sympy.functions.elementary.complexes import re, im
  7. class FunctionMatrix(MatrixExpr):
  8. """Represents a matrix using a function (``Lambda``) which gives
  9. outputs according to the coordinates of each matrix entries.
  10. Parameters
  11. ==========
  12. rows : nonnegative integer. Can be symbolic.
  13. cols : nonnegative integer. Can be symbolic.
  14. lamda : Function, Lambda or str
  15. If it is a SymPy ``Function`` or ``Lambda`` instance,
  16. it should be able to accept two arguments which represents the
  17. matrix coordinates.
  18. If it is a pure string containing Python ``lambda`` semantics,
  19. it is interpreted by the SymPy parser and casted into a SymPy
  20. ``Lambda`` instance.
  21. Examples
  22. ========
  23. Creating a ``FunctionMatrix`` from ``Lambda``:
  24. >>> from sympy import FunctionMatrix, symbols, Lambda, MatPow
  25. >>> i, j, n, m = symbols('i,j,n,m')
  26. >>> FunctionMatrix(n, m, Lambda((i, j), i + j))
  27. FunctionMatrix(n, m, Lambda((i, j), i + j))
  28. Creating a ``FunctionMatrix`` from a SymPy function:
  29. >>> from sympy import KroneckerDelta
  30. >>> X = FunctionMatrix(3, 3, KroneckerDelta)
  31. >>> X.as_explicit()
  32. Matrix([
  33. [1, 0, 0],
  34. [0, 1, 0],
  35. [0, 0, 1]])
  36. Creating a ``FunctionMatrix`` from a SymPy undefined function:
  37. >>> from sympy import Function
  38. >>> f = Function('f')
  39. >>> X = FunctionMatrix(3, 3, f)
  40. >>> X.as_explicit()
  41. Matrix([
  42. [f(0, 0), f(0, 1), f(0, 2)],
  43. [f(1, 0), f(1, 1), f(1, 2)],
  44. [f(2, 0), f(2, 1), f(2, 2)]])
  45. Creating a ``FunctionMatrix`` from Python ``lambda``:
  46. >>> FunctionMatrix(n, m, 'lambda i, j: i + j')
  47. FunctionMatrix(n, m, Lambda((i, j), i + j))
  48. Example of lazy evaluation of matrix product:
  49. >>> Y = FunctionMatrix(1000, 1000, Lambda((i, j), i + j))
  50. >>> isinstance(Y*Y, MatPow) # this is an expression object
  51. True
  52. >>> (Y**2)[10,10] # So this is evaluated lazily
  53. 342923500
  54. Notes
  55. =====
  56. This class provides an alternative way to represent an extremely
  57. dense matrix with entries in some form of a sequence, in a most
  58. sparse way.
  59. """
  60. def __new__(cls, rows, cols, lamda):
  61. rows, cols = _sympify(rows), _sympify(cols)
  62. cls._check_dim(rows)
  63. cls._check_dim(cols)
  64. lamda = sympify(lamda)
  65. if not isinstance(lamda, (FunctionClass, Lambda)):
  66. raise ValueError(
  67. "{} should be compatible with SymPy function classes."
  68. .format(lamda))
  69. if 2 not in lamda.nargs:
  70. raise ValueError(
  71. '{} should be able to accept 2 arguments.'.format(lamda))
  72. if not isinstance(lamda, Lambda):
  73. i, j = Dummy('i'), Dummy('j')
  74. lamda = Lambda((i, j), lamda(i, j))
  75. return super().__new__(cls, rows, cols, lamda)
  76. @property
  77. def shape(self):
  78. return self.args[0:2]
  79. @property
  80. def lamda(self):
  81. return self.args[2]
  82. def _entry(self, i, j, **kwargs):
  83. return self.lamda(i, j)
  84. def _eval_trace(self):
  85. from sympy.matrices.expressions.trace import Trace
  86. from sympy.concrete.summations import Sum
  87. return Trace(self).rewrite(Sum).doit()
  88. def as_real_imag(self):
  89. return (re(Matrix(self)), im(Matrix(self)))