-
Notifications
You must be signed in to change notification settings - Fork 0
/
tensor-contraction.js
86 lines (70 loc) · 2.45 KB
/
tensor-contraction.js
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
var indicesPermutations = require('indices-permutations')
var multiDimArrayIndex = require('multidim-array-index')
/**
* Computes tensor contraction
*
* @params {Function} addition
* @params {Array} indicesPair
* @params {Array} tensorDim
* @params {Array} tensorData
* @returns {Array} contractedTensorData
*/
function tensorContraction (addition, indicesPair, tensorDim, tensorData) {
// Sort indices pair, otherwise algorithm gets unnecessary complicated.
indicesPair.sort()
var p0 = indicesPair[0]
var p1 = indicesPair[1]
var dim0 = tensorDim[p0]
var dim1 = tensorDim[p1]
if (dim0 !== dim1) {
throw new TypeError('Contraction indices does not have the same dimension: ' +
p0 + '-th index = ' + dim0 + ' but ' + p1 + '-th index = ' + dim1 + '.')
}
function varyingTensorDim (result, element, index) {
if ((index !== p0) && (index !== p1)) {
result.push(element)
}
return result
}
function copyArray (result, element) {
result.push(element)
return result
}
function sumOverVarying (tensorData) {
return function (result, varyingCombination) {
var firstCombination = varyingCombination.reduce(copyArray, [])
firstCombination.splice(p0, 0, 0)
firstCombination.splice(p1, 0, 0)
var firstIndex = multiDimArrayIndex(tensorDim, firstCombination)
var element = tensorData[firstIndex]
for (var i = 1; i < dim0; i++) {
var combination = varyingCombination.reduce(copyArray, [])
combination.splice(p0, 0, i)
combination.splice(p1, 0, i)
var index = multiDimArrayIndex(tensorDim, combination)
element = addition(element, tensorData[index])
}
result.push(element)
return result
}
}
// If given tensor has order 2, the contracted tensor will be a scalar
// so it makes sense to return an element, not an array.
// Furthermore, varyingTensorDim will be an empty array so generic algorithm
// will not even be triggered. Then it will be simply computed the trace.
if (tensorDim.length === 2) {
var trace = tensorData[0]
for (var i = 1; i < dim0; i++) {
var combination = [i, i]
var index = multiDimArrayIndex(tensorDim, combination)
trace = addition(trace, tensorData[index])
}
return trace
} else {
return tensorDim
.reduce(varyingTensorDim, [])
.reduce(indicesPermutations, [])
.reduce(sumOverVarying(tensorData), [])
}
}
module.exports = tensorContraction