1# -*- coding: utf-8 -*-
2
3"""
4Tensors
5"""
6
7from mathics.version import __version__  # noqa used in loading to check consistency.
8
9from mathics.builtin.base import Builtin, BinaryOperator
10from mathics.core.expression import Expression, Integer, Integer0, String, SymbolTrue, SymbolFalse
11from mathics.core.rules import Pattern
12
13from mathics.builtin.lists import get_part
14
15
16class ArrayQ(Builtin):
17    """
18    <dl>
19    <dt>'ArrayQ[$expr$]'
20        <dd>tests whether $expr$ is a full array.
21    <dt>'ArrayQ[$expr$, $pattern$]'
22        <dd>also tests whether the array depth of $expr$ matches $pattern$.
23    <dt>'ArrayQ[$expr$, $pattern$, $test$]'</dt>
24        <dd>furthermore tests whether $test$ yields 'True' for all elements of $expr$.
25        'ArrayQ[$expr$]' is equivalent to 'ArrayQ[$expr$, _, True&]'.
26    </dl>
27
28    >> ArrayQ[a]
29     = False
30    >> ArrayQ[{a}]
31     = True
32    >> ArrayQ[{{{a}},{{b,c}}}]
33     = False
34    >> ArrayQ[{{a, b}, {c, d}}, 2, SymbolQ]
35     = True
36    """
37
38    rules = {
39        "ArrayQ[expr_]": "ArrayQ[expr, _, True&]",
40        "ArrayQ[expr_, pattern_]": "ArrayQ[expr, pattern, True&]",
41    }
42
43    def apply(self, expr, pattern, test, evaluation):
44        "ArrayQ[expr_, pattern_, test_]"
45
46        pattern = Pattern.create(pattern)
47
48        dims = [len(expr.get_leaves())]  # to ensure an atom is not an array
49
50        def check(level, expr):
51            if not expr.has_form("List", None):
52                test_expr = Expression(test, expr)
53                if test_expr.evaluate(evaluation) != SymbolTrue:
54                    return False
55                level_dim = None
56            else:
57                level_dim = len(expr.leaves)
58
59            if len(dims) > level:
60                if dims[level] != level_dim:
61                    return False
62            else:
63                dims.append(level_dim)
64            if level_dim is not None:
65                for leaf in expr.leaves:
66                    if not check(level + 1, leaf):
67                        return False
68            return True
69
70        if not check(0, expr):
71            return SymbolFalse
72
73        depth = len(dims) - 1  # None doesn't count
74        if not pattern.does_match(Integer(depth), evaluation):
75            return SymbolFalse
76        return SymbolTrue
77
78
79class VectorQ(Builtin):
80    """
81    <dl>
82    <dt>'VectorQ[$v$]'
83        <dd>returns 'True' if $v$ is a list of elements which are
84        not themselves lists.
85    <dt>'VectorQ[$v$, $f$]'
86        <dd>returns 'True' if $v$ is a vector and '$f$[$x$]' returns
87        'True' for each element $x$ of $v$.
88    </dl>
89
90    >> VectorQ[{a, b, c}]
91     = True
92    """
93
94    rules = {
95        "VectorQ[expr_]": "ArrayQ[expr, 1]",
96        "VectorQ[expr_, test_]": "ArrayQ[expr, 1, test]",
97    }
98
99
100class MatrixQ(Builtin):
101    """
102    <dl>
103    <dt>'MatrixQ[$m$]'
104        <dd>returns 'True' if $m$ is a list of equal-length lists.
105    <dt>'MatrixQ[$m$, $f$]'
106        <dd>only returns 'True' if '$f$[$x$]' returns 'True' for each
107        element $x$ of the matrix $m$.
108    </dl>
109
110    >> MatrixQ[{{1, 3}, {4.0, 3/2}}, NumberQ]
111     = True
112    """
113
114    rules = {
115        "MatrixQ[expr_]": "ArrayQ[expr, 2]",
116        "MatrixQ[expr_, test_]": "ArrayQ[expr, 2, test]",
117    }
118
119
120def get_dimensions(expr, head=None):
121    if expr.is_atom():
122        return []
123    else:
124        if head is not None and not expr.head.sameQ(head):
125            return []
126        sub_dim = None
127        sub = []
128        for leaf in expr.leaves:
129            sub = get_dimensions(leaf, expr.head)
130            if sub_dim is None:
131                sub_dim = sub
132            else:
133                if sub_dim != sub:
134                    sub = []
135                    break
136        return [len(expr.leaves)] + sub
137
138
139class Dimensions(Builtin):
140    """
141    <dl>
142    <dt>'Dimensions[$expr$]'
143        <dd>returns a list of the dimensions of the expression $expr$.
144    </dl>
145
146    A vector of length 3:
147    >> Dimensions[{a, b, c}]
148     = {3}
149
150    A 3x2 matrix:
151    >> Dimensions[{{a, b}, {c, d}, {e, f}}]
152     = {3, 2}
153
154    Ragged arrays are not taken into account:
155    >> Dimensions[{{a, b}, {b, c}, {c, d, e}}]
156     = {3}
157
158    The expression can have any head:
159    >> Dimensions[f[f[a, b, c]]]
160     = {1, 3}
161
162    #> Dimensions[{}]
163     = {0}
164    #> Dimensions[{{}}]
165     = {1, 0}
166    """
167
168    def apply(self, expr, evaluation):
169        "Dimensions[expr_]"
170
171        return Expression("List", *[Integer(dim) for dim in get_dimensions(expr)])
172
173
174class ArrayDepth(Builtin):
175    """
176    <dl>
177    <dt>'ArrayDepth[$a$]'
178        <dd>returns the depth of the non-ragged array $a$, defined as
179        'Length[Dimensions[$a$]]'.
180    </dl>
181
182    >> ArrayDepth[{{a,b},{c,d}}]
183     = 2
184    >> ArrayDepth[x]
185     = 0
186    """
187
188    rules = {
189        "ArrayDepth[list_]": "Length[Dimensions[list]]",
190    }
191
192
193class Dot(BinaryOperator):
194    """
195    <dl>
196    <dt>'Dot[$x$, $y$]'
197    <dt>'$x$ . $y$'
198        <dd>computes the vector dot product or matrix product $x$ . $y$.
199    </dl>
200
201    Scalar product of vectors:
202    >> {a, b, c} . {x, y, z}
203     = a x + b y + c z
204    Product of matrices and vectors:
205    >> {{a, b}, {c, d}} . {x, y}
206     = {a x + b y, c x + d y}
207    Matrix product:
208    >> {{a, b}, {c, d}} . {{r, s}, {t, u}}
209     = {{a r + b t, a s + b u}, {c r + d t, c s + d u}}
210    >> a . b
211     = a . b
212    """
213
214    operator = "."
215    precedence = 490
216    attributes = ("Flat", "OneIdentity")
217
218    rules = {
219        "Dot[a_List, b_List]": "Inner[Times, a, b, Plus]",
220    }
221
222
223class Inner(Builtin):
224    """
225    <dl>
226    <dt>'Inner[$f$, $x$, $y$, $g$]'
227        <dd>computes a generalised inner product of $x$ and $y$, using
228        a multiplication function $f$ and an addition function $g$.
229    </dl>
230
231    >> Inner[f, {a, b}, {x, y}, g]
232     = g[f[a, x], f[b, y]]
233
234    'Inner' can be used to compute a dot product:
235    >> Inner[Times, {a, b}, {c, d}, Plus] == {a, b} . {c, d}
236     = True
237
238    The inner product of two boolean matrices:
239    >> Inner[And, {{False, False}, {False, True}}, {{True, False}, {True, True}}, Or]
240     = {{False, False}, {True, True}}
241
242    Inner works with tensors of any depth:
243    >> Inner[f, {{{a, b}}, {{c, d}}}, {{1}, {2}}, g]
244     = {{{g[f[a, 1], f[b, 2]]}}, {{g[f[c, 1], f[d, 2]]}}}
245
246
247    ## Issue #670
248    #> A = {{ b ^ ( -1 / 2), 0}, {a * b ^ ( -1 / 2 ), b ^ ( 1 / 2 )}}
249     = {{1 / Sqrt[b], 0}, {a / Sqrt[b], Sqrt[b]}}
250    #> A . Inverse[A]
251     = {{1, 0}, {0, 1}}
252    #> A
253     = {{1 / Sqrt[b], 0}, {a / Sqrt[b], Sqrt[b]}}
254    """
255
256    rules = {
257        "Inner[f_, list1_, list2_]": "Inner[f, list1, list2, Plus]",
258    }
259
260    messages = {
261        "incom": (
262            "Length `1` of dimension `2` in `3` is incommensurate with "
263            "length `4` of dimension 1 in `5."
264        ),
265    }
266
267    def apply(self, f, list1, list2, g, evaluation):
268        "Inner[f_, list1_, list2_, g_]"
269
270        m = get_dimensions(list1)
271        n = get_dimensions(list2)
272
273        if not m or not n:
274            evaluation.message("Inner", "normal")
275            return
276        if list1.get_head() != list2.get_head():
277            evaluation.message("Inner", "heads", list1.get_head(), list2.get_head())
278            return
279        if m[-1] != n[0]:
280            evaluation.message("Inner", "incom", m[-1], len(m), list1, n[0], list2)
281            return
282
283        head = list1.get_head()
284        inner_dim = n[0]
285
286        def rec(i_cur, j_cur, i_rest, j_rest):
287            evaluation.check_stopped()
288            if i_rest:
289                leaves = []
290                for i in range(1, i_rest[0] + 1):
291                    leaves.append(rec(i_cur + [i], j_cur, i_rest[1:], j_rest))
292                return Expression(head, *leaves)
293            elif j_rest:
294                leaves = []
295                for j in range(1, j_rest[0] + 1):
296                    leaves.append(rec(i_cur, j_cur + [j], i_rest, j_rest[1:]))
297                return Expression(head, *leaves)
298            else:
299
300                def summand(i):
301                    part1 = get_part(list1, i_cur + [i])
302                    part2 = get_part(list2, [i] + j_cur)
303                    return Expression(f, part1, part2)
304
305                part = Expression(g, *[summand(i) for i in range(1, inner_dim + 1)])
306                # cur_expr.leaves.append(part)
307                return part
308
309        return rec([], [], m[:-1], n[1:])
310
311
312class Outer(Builtin):
313    """
314    <dl>
315    <dt>'Outer[$f$, $x$, $y$]'
316        <dd>computes a generalised outer product of $x$ and $y$, using
317        the function $f$ in place of multiplication.
318    </dl>
319
320    >> Outer[f, {a, b}, {1, 2, 3}]
321     = {{f[a, 1], f[a, 2], f[a, 3]}, {f[b, 1], f[b, 2], f[b, 3]}}
322
323    Outer product of two matrices:
324    >> Outer[Times, {{a, b}, {c, d}}, {{1, 2}, {3, 4}}]
325     = {{{{a, 2 a}, {3 a, 4 a}}, {{b, 2 b}, {3 b, 4 b}}}, {{{c, 2 c}, {3 c, 4 c}}, {{d, 2 d}, {3 d, 4 d}}}}
326
327    'Outer' of multiple lists:
328    >> Outer[f, {a, b}, {x, y, z}, {1, 2}]
329     = {{{f[a, x, 1], f[a, x, 2]}, {f[a, y, 1], f[a, y, 2]}, {f[a, z, 1], f[a, z, 2]}}, {{f[b, x, 1], f[b, x, 2]}, {f[b, y, 1], f[b, y, 2]}, {f[b, z, 1], f[b, z, 2]}}}
330
331    Arrays can be ragged:
332    >> Outer[Times, {{1, 2}}, {{a, b}, {c, d, e}}]
333     = {{{{a, b}, {c, d, e}}, {{2 a, 2 b}, {2 c, 2 d, 2 e}}}}
334
335    Word combinations:
336    >> Outer[StringJoin, {"", "re", "un"}, {"cover", "draw", "wind"}, {"", "ing", "s"}] // InputForm
337     = {{{"cover", "covering", "covers"}, {"draw", "drawing", "draws"}, {"wind", "winding", "winds"}}, {{"recover", "recovering", "recovers"}, {"redraw", "redrawing", "redraws"}, {"rewind", "rewinding", "rewinds"}}, {{"uncover", "uncovering", "uncovers"}, {"undraw", "undrawing", "undraws"}, {"unwind", "unwinding", "unwinds"}}}
338
339    Compositions of trigonometric functions:
340    >> trigs = Outer[Composition, {Sin, Cos, Tan}, {ArcSin, ArcCos, ArcTan}]
341     = {{Composition[Sin, ArcSin], Composition[Sin, ArcCos], Composition[Sin, ArcTan]}, {Composition[Cos, ArcSin], Composition[Cos, ArcCos], Composition[Cos, ArcTan]}, {Composition[Tan, ArcSin], Composition[Tan, ArcCos], Composition[Tan, ArcTan]}}
342    Evaluate at 0:
343    >> Map[#[0] &, trigs, {2}]
344     = {{0, 1, 0}, {1, 0, 1}, {0, ComplexInfinity, 0}}
345    """
346
347    def apply(self, f, lists, evaluation):
348        "Outer[f_, lists__]"
349
350        lists = lists.get_sequence()
351        head = None
352        for list in lists:
353            if list.is_atom():
354                evaluation.message("Outer", "normal")
355                return
356            if head is None:
357                head = list.head
358            elif not list.head.sameQ(head):
359                evaluation.message("Outer", "heads", head, list.head)
360                return
361
362        def rec(item, rest_lists, current):
363            evaluation.check_stopped()
364            if item.is_atom() or not item.head.sameQ(head):
365                if rest_lists:
366                    return rec(rest_lists[0], rest_lists[1:], current + [item])
367                else:
368                    return Expression(f, *(current + [item]))
369            else:
370                leaves = []
371                for leaf in item.leaves:
372                    leaves.append(rec(leaf, rest_lists, current))
373                return Expression(head, *leaves)
374
375        return rec(lists[0], lists[1:], [])
376
377
378class Transpose(Builtin):
379    """
380    <dl>
381    <dt>'Tranpose[$m$]'
382        <dd>transposes rows and columns in the matrix $m$.
383    </dl>
384
385    >> Transpose[{{1, 2, 3}, {4, 5, 6}}]
386     = {{1, 4}, {2, 5}, {3, 6}}
387    >> MatrixForm[%]
388     = 1   4
389     .
390     . 2   5
391     .
392     . 3   6
393
394    #> Transpose[x]
395     = Transpose[x]
396    """
397
398    def apply(self, m, evaluation):
399        "Transpose[m_?MatrixQ]"
400
401        result = []
402        for row_index, row in enumerate(m.leaves):
403            for col_index, item in enumerate(row.leaves):
404                if row_index == 0:
405                    result.append([item])
406                else:
407                    result[col_index].append(item)
408        return Expression("List", *[Expression("List", *row) for row in result])
409
410
411class DiagonalMatrix(Builtin):
412    """
413    <dl>
414    <dt>'DiagonalMatrix[$list$]'
415        <dd>gives a matrix with the values in $list$ on its diagonal and zeroes elsewhere.
416    </dl>
417
418    >> DiagonalMatrix[{1, 2, 3}]
419     = {{1, 0, 0}, {0, 2, 0}, {0, 0, 3}}
420    >> MatrixForm[%]
421     = 1   0   0
422     .
423     . 0   2   0
424     .
425     . 0   0   3
426
427    #> DiagonalMatrix[a + b]
428     = DiagonalMatrix[a + b]
429    """
430
431    def apply(self, list, evaluation):
432        "DiagonalMatrix[list_List]"
433
434        result = []
435        n = len(list.leaves)
436        for index, item in enumerate(list.leaves):
437            row = [Integer0] * n
438            row[index] = item
439            result.append(Expression("List", *row))
440        return Expression("List", *result)
441
442
443class IdentityMatrix(Builtin):
444    """
445    <dl>
446    <dt>'IdentityMatrix[$n$]'
447        <dd>gives the identity matrix with $n$ rows and columns.
448    </dl>
449
450    >> IdentityMatrix[3]
451     = {{1, 0, 0}, {0, 1, 0}, {0, 0, 1}}
452    """
453
454    rules = {
455        "IdentityMatrix[n_Integer]": "DiagonalMatrix[Table[1, {n}]]",
456    }
457
458
459def get_default_distance(p):
460    if all(q.is_numeric() for q in p):
461        return "SquaredEuclideanDistance"
462    elif all(q.get_head_name() == "System`List" for q in p):
463        dimensions = [get_dimensions(q) for q in p]
464        if len(dimensions) < 1:
465            return None
466        d0 = dimensions[0]
467        if not all(d == d0 for d in dimensions[1:]):
468            return None
469        if len(dimensions[0]) == 1:  # vectors?
470
471            def is_boolean(x):
472                return x.get_head_name() == "System`Symbol" and x in (
473                    SymbolTrue,
474                    SymbolFalse,
475                )
476
477            if all(all(is_boolean(e) for e in q.leaves) for q in p):
478                return "JaccardDissimilarity"
479        return "SquaredEuclideanDistance"
480    elif all(isinstance(q, String) for q in p):
481        return "EditDistance"
482    else:
483        from mathics.builtin.graphics import expression_to_color
484
485        if all(expression_to_color(q) is not None for q in p):
486            return "ColorDistance"
487
488        return None
489