使用 python 在四阶张量中切换索引

switch indices in 4th order tensor with python

我想复制以下操作:

import numpy as np
a = np.eye(3)
A = np.einsum("ij,kl->ikjl", a,a)

我有以下内容:

def IndexPermutation(stringWithTensorIndices, Tensor):

    # Function to split the string with the indices
    def split(word): 
       return [char for char in word]
    
    # Gets indices from string               
    a, b, c, d = split(stringWithTensorIndices) 
    
    store=np.zeros((3,3,3,3))

    # Performs the tensor index permutations
    for i in [0,1,2]:
        for j in [0,1,2]:
            for k in [0,1,2]:
                for l in [0,1,2]:
                    store[i,j,k,l]=Tensor[vars()[a],vars()[b],vars()[c],vars()[d]]
                    
    return store

我想删除 vars() 功能。目前,它只接受“ij”“kl”然后排列索引(相当于 einsum(“ijkl -> ijkl 的某种组合”)。

任何人都可以帮助我吗?

我知道您想要做的是提供一些输入 stringWithTensorIndices,其值如 "ijkl""jlki" 以确定如何实现求和。那么循环的逻辑就是

  • 对于输入为“jlki”的情况,a="j"
  • 等等 vars()[a] = vars()["j"]
  • 这是循环变量的值j
  • 所以你得到store[i,j,k,l] = Tensor[j,l,k,i]

这里有几点要说。

首先,您的 split 函数没有为您做任何事情。 a,b,c,d = "jkli" 将正确解析为 a="j" 等,因为字符串本身可以被视为 Python 集合。无论哪种方式,不正确的用户输入都不安全。最佳做法可能是捕获 ValueError 并提供信息异常。在实践中,我认为无论如何我们都会采取不同的方法。

现在,重要的部分。在 heavily-nested 循环中的任何一点,您都会有一些索引 [i, j, k, l] (这些是您的循环变量)。输入有点像从这个索引集合中进行选择的键。

例如“jlki”应该挑出indices[1], indices[3], indices[2], indices[0].

我们需要的是辅助函数:

tensor_indices = "iljk"
lookup = {"i": 0, "j": 1, "k": 2, "l": 3}
mapped_indices = [lookup[x] for x in tensor_indices]


def get_indices(loop_indices, mapped_indices):
    # loop_indices is [i, j, k, l], changes every iteration
    # mapped indices is [0, 3, 1, 2], never changes

    for_tensor = tuple(loop_indices[x] for x in mapped_indices)

    return for_tensor  # (i, l, j, k), as was required

作为元组而不是列表返回很重要,因为如果您使用列表 a[[0,2,1,0]] 索引 multi-dimensional numpy 数组 a,这与索引它是完全不同的作为 a[0,2,1,0]。前者沿数组的第一维提取切片,后者提取特定值。

幸运的是,Python 通常会尝试将元组视为 comma-separated 变量列表,因此 a[t] 其中 t = (0,2,1,0) 等同于 a[0,2,1,0] .

有了这个,我想我们可以转向一个完整的解决方案:

import numpy as np
from typing import List, Tuple
# type annotations are useful, and mean you don't have to use
# names like stringWithTensorIndices to keep track


def get_indices(
    loop_indices: List[int], mapped_indices: List[int]
) -> Tuple[int, int, int, int]:
    """
    get_indices
    
    This function takes a set of loop variables, e.g. i, j, k, l
    and uses a set of mapped indices (e.g. [0, 1, 1, 2] for 'ijjk')
    and returns loop variables chosen according to the mapping
    (i.e. (i, j, j, k) in this example).

    """
    for_tensor = tuple(loop_indices[x] for x in mapped_indices)

    return for_tensor  # (i, l, j, k), as was required


def index_permutation(tensor_indices: str, tensor: np.ndarray) -> np.ndarray:
    """
    index_permutation

    Permutes indices of a tensor

    Inputs:
        tensor_indices - e.g. "ijkl"
        tensor - the tensor to operate on

    Output:
        new_tensor - the resulting tensor

    """
    lookup = {"i": 0, "j": 1, "k": 2, "l": 3}
    # for e.g. 'iljk' we want 0, 3, 1, 2:
    mapped_indices = [lookup[x] for x in tensor_indices]

    dim = 3
    shape = (dim, dim, dim, dim)
    new_tensor = np.zeros(shape)

    for i in range(dim):
        for j in range(dim):
            for k in range(dim):
                for l in range(dim):
                    indices = [i, j, k, l]
                    reordered = get_indices(indices, mapped_indices)
                    new_tensor[i,j,k,l]=Tensor[reordered]

    return store

杂项

我已经在此处实施了一些 best-practices:

  • D.R.Y。 (不要重复你自己) - 如果你想在未来改变行为,很多东西要同时改变,所以 [0,1,2] 例如只是被存储为一个你引用四次的列表
  • range(N)enumeratezip 这样的内置迭代器很方便
  • docstrings,评论 - 即使 no-one else
  • 以后你也会感谢你
  • 类型注释 - 即使您没有使用 mypy 检查类型,它们也有助于提高可读性 尤其是 在您有一堆具有相似名称的变量的情况下。在最新的 python 中你甚至不需要 List,普通的 list[int] 就可以了。
  • 标准 python 约定 - snake_case 用于变量和函数名称,CamelCase 用于 class 名称

我还没有做的一件事是试图消除那个严重嵌套的循环。 numpy 有很多方法可以做到这一点(这是它擅长的),但在我看来,这与练习的重点无关。