Build a reflection matrix in Pytorch

build_reflection_matrix.py

#!/usr/bin/env python
# -*- coding: UTF8 -*-

# Author: Guillaume Bouvier -- guillaume.bouvier@pasteur.fr
# https://research.pasteur.fr/en/member/guillaume-bouvier/
# 2020-10-01 09:45:36 (UTC+0200)

import torch


def build_reflection_matrix(abc, device):
    # See: https://en.wikipedia.org/w/index.php?title=Transformation_matrix&oldid=976277111#Reflection
    a, b, c = abc
    A = torch.tensor([[1. - 2 * a**2, -2 * a * b, -2 * a * c],
                     [-2 * a * b, 1. - 2 * b**2, -2 * b * c],
                     [-2 * a * c, -2 * b * c, 1. - 2 * c**2]], device=device)
    return A