Source code for tfga.tfga

"""Provides classes and operations for performing geometric algebra
with TensorFlow.

The `GeometricAlgebra` class is used to construct the algebra given a metric.
It exposes methods for operating on `tf.Tensor` instances where their last
axis is interpreted as blades of the algebra.
"""
from typing import List, Any, Union, Optional
import numbers
import tensorflow as tf
import numpy as np

from .cayley import get_cayley_tensor, blades_from_bases
from .blades import (
    BladeKind, get_blade_of_kind_indices, get_blade_indices_from_names,
    get_blade_repr, invert_blade_indices
)
from .mv_ops import mv_multiply, mv_reversion, mv_grade_automorphism, mv_conv1d
from .mv import MultiVector


[docs]class GeometricAlgebra: """Class used for performing geometric algebra operations on `tf.Tensor` instances. Exposes methods for operating on `tf.Tensor` instances where their last axis is interpreted as blades of the algebra. Holds the metric and other quantities derived from it. """ def __init__(self, metric: List[float]): """Creates a GeometricAlgebra object given a metric. The algebra will have as many basis vectors as there are elements in the metric. Args: metric: Metric as a list. Specifies what basis vectors square to """ self._metric = tf.convert_to_tensor(metric, dtype=tf.float32) self._num_bases = len(metric) self._bases = list(map(str, range(self._num_bases))) self._blades, self._blade_degrees = blades_from_bases(self._bases) self._blade_degrees = tf.convert_to_tensor(self._blade_degrees) self._num_blades = len(self._blades) self._max_degree = tf.reduce_max(self._blade_degrees) # [Blades, Blades, Blades] self._cayley, self._cayley_inner, self._cayley_outer = tf.convert_to_tensor( get_cayley_tensor(self.metric, self._bases, self._blades), dtype=tf.float32 ) self._blade_mvs = tf.eye(self._num_blades) self._basis_mvs = self._blade_mvs[1:1+self._num_bases] # Find the dual by looking at the anti-diagonal in the Cayley tensor. self._dual_blade_indices = [] self._dual_blade_signs = [] for blade_index in range(self._num_blades): dual_index = self.num_blades - blade_index - 1 anti_diag = self._cayley[blade_index, dual_index] dual_sign = tf.gather(anti_diag, tf.where( anti_diag != 0.0)[..., 0])[..., 0] self._dual_blade_indices.append(dual_index) self._dual_blade_signs.append(dual_sign) self._dual_blade_indices = tf.convert_to_tensor( self._dual_blade_indices, dtype=tf.int64) self._dual_blade_signs = tf.convert_to_tensor( self._dual_blade_signs, dtype=tf.float32)
[docs] def print(self, *args, **kwargs): """Same as the default `print` function but formats `tf.Tensor` instances that have as many elements on their last axis as the algebra has blades using `mv_repr()`. """ def _is_mv(arg): return isinstance(arg, tf.Tensor) and arg.shape.ndims > 0 and arg.shape[-1] == self.num_blades new_args = [self.mv_repr(arg) if _is_mv(arg) else arg for arg in args] print(*new_args, **kwargs)
@property def metric(self) -> tf.Tensor: """Metric list which contains the number that each basis vector in the algebra squares to (ie. the diagonal of the metric tensor). """ return self._metric @property def cayley(self) -> tf.Tensor: """`MxMxM` tensor where `M` is the number of basis blades in the algebra. Used for calculating the geometric product: `a_i, b_j, cayley_ijk -> c_k` """ return self._cayley @property def cayley_inner(self) -> tf.Tensor: """Analagous to cayley but for inner product.""" return self._cayley_inner @property def cayley_outer(self) -> tf.Tensor: """Analagous to cayley but for outer product.""" return self._cayley_outer @property def blades(self) -> List[str]: """List of all blade names. Blades are all possible independent combinations of basis vectors. Basis vectors are named starting from `"0"` and counting up. The scalar blade is the empty string `""`. Example - Bases: `["0", "1", "2"]` - Blades: `["", "0", "1", "2", "01", "02", "12", "012"]` """ return self._blades @property def blade_mvs(self) -> tf.Tensor: """List of all blade tensors in the algebra.""" return self._blade_mvs @property def dual_blade_indices(self) -> tf.Tensor: """Indices of the dual blades for each blade.""" return self._dual_blade_indices @property def dual_blade_signs(self) -> tf.Tensor: """Signs of the dual blades for each blade.""" return self._dual_blade_signs @property def num_blades(self) -> int: """Total number of blades in the algebra.""" return self._num_blades @property def blade_degrees(self) -> tf.Tensor: """List of blade-degree for each blade in the algebra.""" return self._blade_degrees @property def max_degree(self) -> int: """Highest blade degree in the algebra.""" return self._max_degree @property def basis_mvs(self) -> tf.Tensor: """List of basis vectors as tf.Tensor.""" return self._basis_mvs
[docs] def get_kind_blade_indices(self, kind: BladeKind, invert: bool = False) -> tf.Tensor: """Find all indices of blades of a given kind in the algebra. Args: kind: kind of blade to give indices for invert: whether to return all blades not of the kind Returns: indices of blades of a given kind in the algebra """ return get_blade_of_kind_indices(self.blade_degrees, kind, self.max_degree, invert=invert)
[docs] def get_blade_indices_of_degree(self, degree: int) -> tf.Tensor: """Find all indices of blades of the given degree. Args: degree: degree to return blades for Returns: indices of blades with the given degree in the algebra """ return tf.gather(tf.range(self.num_blades), tf.where(self.blade_degrees == degree)[..., 0])
[docs] def is_pure(self, tensor: tf.Tensor, blade_indices: tf.Tensor) -> bool: """Returns whether the given tensor is purely of the given blades and has no non-zero values for blades not in the given blades. Args: tensor: tensor to check purity for blade_indices: blade indices to check purity for Returns: Whether the tensor is purely of the given blades and has no non-zero values for blades not in the given blades """ tensor = tf.convert_to_tensor(tensor, dtype_hint=tf.float32) blade_indices = tf.convert_to_tensor( blade_indices, dtype_hint=tf.int64) inverted_blade_indices = invert_blade_indices( self.num_blades, blade_indices) return tf.reduce_all(tf.gather( tensor, inverted_blade_indices, axis=-1 ) == 0)
[docs] def is_pure_kind(self, tensor: tf.Tensor, kind: BladeKind) -> bool: """Returns whether the given tensor is purely of a given kind and has no non-zero values for blades not of the kind. Args: tensor: tensor to check purity for kind: kind of blade to check purity for Returns: Whether the tensor is purely of a given kind and has no non-zero values for blades not of the kind """ tensor = tf.convert_to_tensor(tensor, dtype_hint=tf.float32) inverted_kind_indices = self.get_kind_blade_indices(kind, invert=True) return tf.reduce_all(tf.gather( tensor, inverted_kind_indices, axis=-1 ) == 0)
[docs] def from_tensor(self, tensor: tf.Tensor, blade_indices: tf.Tensor) -> tf.Tensor: """Creates a geometric algebra tf.Tensor from a tf.Tensor and blade indices. The blade indices have to align with the last axis of the tensor. Args: tensor: tf.Tensor to take as values for the geometric algebra tensor blade_indices: Blade indices corresponding to the tensor. Can be obtained from blade names eg. using get_kind_blade_indices() or as indices from the blades list property. Returns: Geometric algebra tf.Tensor from tensor and blade indices """ blade_indices = tf.cast( tf.convert_to_tensor(blade_indices, dtype_hint=tf.int64), dtype=tf.int64 ) tensor = tf.convert_to_tensor(tensor, dtype_hint=tf.float32) # Put last axis on first axis so scatter_nd becomes easier. # Later undo the transposition again. t = tf.concat([[tensor.shape.ndims - 1], tf.range(0, tensor.shape.ndims - 1)], axis=0) t_inv = tf.concat([tf.range(1, tensor.shape.ndims), [0]], axis=0) tensor = tf.transpose(tensor, t) shape = tf.concat([ tf.convert_to_tensor([self.num_blades], dtype=tf.int64), tf.shape(tensor, tf.int64)[1:] ], axis=0) tensor = tf.scatter_nd( tf.expand_dims(blade_indices, axis=-1), tensor, shape ) return tf.transpose(tensor, t_inv)
[docs] def from_tensor_with_kind(self, tensor: tf.Tensor, kind: BladeKind) -> tf.Tensor: """Creates a geometric algebra tf.Tensor from a tf.Tensor and a kind. The kind's blade indices have to align with the last axis of the tensor. Args: tensor: tf.Tensor to take as values for the geometric algebra tensor kind: Kind corresponding to the tensor Returns: Geometric algebra tf.Tensor from tensor and kind """ # Put last axis on first axis so scatter_nd becomes easier. # Later undo the transposition again. tensor = tf.convert_to_tensor(tensor, dtype_hint=tf.float32) kind_indices = self.get_kind_blade_indices(kind) return self.from_tensor(tensor, kind_indices)
[docs] def from_scalar(self, scalar: numbers.Number) -> tf.Tensor: """Creates a geometric algebra tf.Tensor with scalar elements. Args: scalar: Elements to be used as scalars Returns: Geometric algebra tf.Tensor from scalars """ return self.from_tensor_with_kind(tf.expand_dims(scalar, axis=-1), BladeKind.SCALAR)
[docs] def e(self, *blades: List[str]) -> tf.Tensor: """Returns a geometric algebra tf.Tensor with the given blades set to 1. Args: blades: list of blade names, can be unnormalized Returns: tf.Tensor with blades set to 1 """ blade_signs, blade_indices = get_blade_indices_from_names( blades, self.blades) blade_indices = tf.convert_to_tensor(blade_indices) # Don't allow duplicate indices tf.Assert( blade_indices.shape[0] == tf.unique(blade_indices)[0].shape[0], [blades] ) x = ( tf.expand_dims(blade_signs, axis=-1) * tf.gather(self.blade_mvs, blade_indices) ) # a, b -> b return tf.reduce_sum(x, axis=-2)
def __getattr__(self, name: str) -> tf.Tensor: """Returns basis blade tensors if name was a basis.""" if name.startswith("e") and (name[1:] == "" or int(name[1:]) >= 0): return self.e(name[1:]) raise AttributeError
[docs] def dual(self, tensor: tf.Tensor) -> tf.Tensor: """Returns the dual of the geometric algebra tensor. Args: tensor: Geometric algebra tensor to return dual for Returns: Dual of the geometric algebra tensor """ tensor = tf.convert_to_tensor(tensor, dtype_hint=tf.float32) return self.dual_blade_signs * tf.gather(tensor, self.dual_blade_indices, axis=-1)
[docs] def grade_automorphism(self, tensor: tf.Tensor) -> tf.Tensor: """Returns the geometric algebra tensor with odd grades negated. See https://en.wikipedia.org/wiki/Paravector#Grade_automorphism. Args: tensor: Geometric algebra tensor to return grade automorphism for Returns: Geometric algebra tensor with odd grades negated """ tensor = tf.convert_to_tensor(tensor, dtype_hint=tf.float32) return mv_grade_automorphism(tensor, self.blade_degrees)
[docs] def reversion(self, tensor: tf.Tensor) -> tf.Tensor: """Returns the grade-reversed geometric algebra tensor. See https://en.wikipedia.org/wiki/Paravector#Reversion_conjugation. Args: tensor: Geometric algebra tensor to return grade-reversion for Returns: Grade-reversed geometric algebra tensor """ tensor = tf.convert_to_tensor(tensor, dtype_hint=tf.float32) return mv_reversion(tensor, self.blade_degrees)
[docs] def conjugation(self, tensor: tf.Tensor) -> tf.Tensor: """Combines reversion and grade automorphism. See https://en.wikipedia.org/wiki/Paravector#Clifford_conjugation. Args: tensor: Geometric algebra tensor to return conjugate for Returns: Geometric algebra tensor after `reversion()` and `grade_automorphism()` """ tensor = tf.convert_to_tensor(tensor, dtype_hint=tf.float32) return self.grade_automorphism(self.reversion(tensor))
[docs] def simple_inverse(self, a: tf.Tensor) -> tf.Tensor: """Returns the inverted geometric algebra tensor `X^-1` such that `X * X^-1 = 1`. Only works for elements that square to scalars. Faster than the general inverse. Args: a: Geometric algebra tensor to return inverse for Returns: inverted geometric algebra tensor """ a = tf.convert_to_tensor(a, dtype_hint=tf.float32) rev_a = self.reversion(a) divisor = self.geom_prod(a, rev_a) if not self.is_pure_kind(divisor, BladeKind.SCALAR): raise Exception( "Can't invert multi-vector (inversion divisor V ~V not scalar: %s)." % divisor) # Divide by scalar part return rev_a / divisor[..., :1]
[docs] def reg_prod(self, a: tf.Tensor, b: tf.Tensor) -> tf.Tensor: """Returns the regressive product of two geometric algebra tensors. Args: a: Geometric algebra tensor on the left hand side of the regressive product b: Geometric algebra tensor on the right hand side of the regressive product Returns: regressive product of a and b """ a = tf.convert_to_tensor(a, dtype_hint=tf.float32) b = tf.convert_to_tensor(b, dtype_hint=tf.float32) return self.dual(self.ext_prod(self.dual(a), self.dual(b)))
[docs] def ext_prod(self, a: tf.Tensor, b: tf.Tensor) -> tf.Tensor: """Returns the exterior product of two geometric algebra tensors. Args: a: Geometric algebra tensor on the left hand side of the exterior product b: Geometric algebra tensor on the right hand side of the exterior product Returns: exterior product of a and b """ a = tf.convert_to_tensor(a, dtype_hint=tf.float32) b = tf.convert_to_tensor(b, dtype_hint=tf.float32) return mv_multiply(a, b, self._cayley_outer)
[docs] def geom_prod(self, a: tf.Tensor, b: tf.Tensor) -> tf.Tensor: """Returns the geometric product of two geometric algebra tensors. Args: a: Geometric algebra tensor on the left hand side of the geometric product b: Geometric algebra tensor on the right hand side of the geometric product Returns: geometric product of a and b """ a = tf.convert_to_tensor(a, dtype_hint=tf.float32) b = tf.convert_to_tensor(b, dtype_hint=tf.float32) a = tf.convert_to_tensor(a) b = tf.convert_to_tensor(b) return mv_multiply(a, b, self._cayley)
[docs] def inner_prod(self, a: tf.Tensor, b: tf.Tensor) -> tf.Tensor: """Returns the inner product of two geometric algebra tensors. Args: a: Geometric algebra tensor on the left hand side of the inner product b: Geometric algebra tensor on the right hand side of the inner product Returns: inner product of a and b """ a = tf.convert_to_tensor(a, dtype_hint=tf.float32) b = tf.convert_to_tensor(b, dtype_hint=tf.float32) return mv_multiply(a, b, self._cayley_inner)
[docs] def geom_conv1d(self, a: tf.Tensor, k: tf.Tensor, stride: int, padding: str, dilations: Union[int, None] = None) -> tf.Tensor: """Returns the 1D convolution of a sequence with a geometric algebra tensor kernel. The convolution is performed using the geometric product. Args: a: Input geometric algebra tensor of shape [..., Length, ChannelsIn, Blades] k: Geometric algebra tensor for the convolution kernel of shape [KernelSize, ChannelsIn, ChannelsOut, Blades] stride: Stride to use for the convolution padding: "SAME" (zero-pad input length so output length == input length / stride) or "VALID" (no padding) Returns: Geometric algbra tensor of shape [..., OutputLength, ChannelsOut, Blades] representing `a` convolved with `k` """ a = tf.convert_to_tensor(a, dtype_hint=tf.float32) k = tf.convert_to_tensor(k, dtype_hint=tf.float32) return mv_conv1d(a, k, self._cayley, stride=stride, padding=padding)
[docs] def mv_repr(self, a: tf.Tensor) -> str: """Returns a string representation for the given geometric algebra tensor. Args: a: Geometric algebra tensor to return the representation for Returns: string representation for `a` """ a = tf.convert_to_tensor(a, dtype_hint=tf.float32) if len(a.shape) == 1: return "MultiVector[%s]" % " + ".join( "%.2f*%s" % (value, get_blade_repr(blade_name)) for value, blade_name in zip(a, self.blades) if value != 0 ) else: return "MultiVector[batch_shape=%s]" % a.shape[:-1]
[docs] def approx_exp(self, a: tf.Tensor, order: int = 50) -> tf.Tensor: """Returns an approximation of the exponential using a centered taylor series. Args: a: Geometric algebra tensor to return exponential for order: order of the approximation Returns: Approximation of `exp(a)` """ a = tf.convert_to_tensor(a, dtype_hint=tf.float32) v = self.from_scalar(1.0) result = self.from_scalar(1.0) for i in range(1, order + 1): v = self.geom_prod(a, v) i_factorial = tf.exp(tf.math.lgamma(i + 1.0)) result += v / i_factorial return result
[docs] def exp(self, a: tf.Tensor, square_scalar_tolerance: Union[float, None] = 1e-4) -> tf.Tensor: """Returns the exponential of the passed geometric algebra tensor. Only works for multivectors that square to scalars. Args: a: Geometric algebra tensor to return exponential for square_scalar_tolerance: Tolerance to use for the square scalar check or None if the check should be skipped Returns: `exp(a)` """ # See https://www.euclideanspace.com/maths/algebra/clifford/algebra/functions/exponent/index.htm # for an explanation of how to exponentiate multivectors. self_sq = self.geom_prod(a, a) if square_scalar_tolerance is not None: tf.Assert(tf.reduce_all( tf.abs(self_sq[..., 1:]) < square_scalar_tolerance ), [self_sq]) scalar_self_sq = self_sq[..., :1] # "Complex" square root (argument can be negative) s_sqrt = tf.sign(scalar_self_sq) * tf.sqrt(tf.abs(scalar_self_sq)) # Square to +1: cosh(sqrt(||a||)) + a / sqrt(||a||) sinh(sqrt(||a||)) # Square to -1: cos(sqrt(||a||)) + a / sqrt(||a||) sin(sqrt(||a||)) # TODO: Does this work for values other than 1 too? eg. square to +0.5? # TODO: Find a solution that doesnt require calculating all possibilities # first. non_zero_result = tf.where( scalar_self_sq < 0, (self.from_tensor(tf.cos(s_sqrt), [0]) + a / s_sqrt * tf.sin(s_sqrt)), (self.from_tensor(tf.cosh(s_sqrt), [0]) + a / s_sqrt * tf.sinh(s_sqrt)) ) return tf.where(scalar_self_sq == 0, self.from_scalar(1.0) + a, non_zero_result)
[docs] def approx_log(self, a: tf.Tensor, order: int = 50) -> tf.Tensor: """Returns an approximation of the natural logarithm using a centered taylor series. Only converges for multivectors where `||mv - 1|| < 1`. Args: a: Geometric algebra tensor to return logarithm for order: order of the approximation Returns: Approximation of `log(a)` """ a = tf.convert_to_tensor(a, dtype_hint=tf.float32) result = self.from_scalar(0.0) a_minus_one = a - self.from_scalar(1.0) v = None for i in range(1, order + 1): v = a_minus_one if v is None else v * a_minus_one result += (((-1.0) ** i) / i) * v return -result
[docs] def int_pow(self, a: tf.Tensor, n: int) -> tf.Tensor: """Returns the geometric algebra tensor to the power of an integer using repeated multiplication. Args: a: Geometric algebra tensor to raise n: integer power to raise the multivector to Returns: `a` to the power of `n` """ a = tf.convert_to_tensor(a, dtype_hint=tf.float32) if not isinstance(n, int): raise Exception("n must be an integer.") if n < 0: raise Exception("Can't raise to negative powers.") if n == 0: # TODO: more efficient (ones only in scalar) return tf.ones_like(a) * self.e("") result = a for i in range(n - 1): result = self.geom_prod(result, a) return result
[docs] def keep_blades(self, a: tf.Tensor, blade_indices: List[int]) -> tf.Tensor: """Takes a geometric algebra tensor and returns it with only the given blade_indices as non-zeros. Args: a: Geometric algebra tensor to copy blade_indices: Indices for blades to keep Returns: `a` with only `blade_indices` components as non-zeros """ a = tf.convert_to_tensor(a, dtype_hint=tf.float32) blade_indices = tf.cast( tf.convert_to_tensor(blade_indices, dtype_hint=tf.int64), dtype=tf.int64 ) blade_values = tf.gather(a, blade_indices, axis=-1) return self.from_tensor(blade_values, blade_indices)
[docs] def keep_blades_with_name(self, a: tf.Tensor, blade_names: Union[List[str], str]) -> tf.Tensor: """Takes a geometric algebra tensor and returns it with only the given blades as non-zeros. Args: a: Geometric algebra tensor to copy blade_names: Blades to keep Returns: `a` with only `blade_names` components as non-zeros """ if isinstance(blade_names, str): blade_names = [blade_names] _, blade_indices = get_blade_indices_from_names( blade_names, self.blades) return self.keep_blades(a, blade_indices)
[docs] def select_blades(self, a: tf.Tensor, blade_indices: List[int]) -> tf.Tensor: """Takes a geometric algebra tensor and returns a `tf.Tensor` with the blades in blade_indices on the last axis. Args: a: Geometric algebra tensor to copy blade_indices: Indices for blades to select Returns: `tf.Tensor` based on `a` with `blade_indices` on last axis. """ a = tf.convert_to_tensor(a, dtype_hint=tf.float32) blade_indices = tf.cast( tf.convert_to_tensor(blade_indices, dtype_hint=tf.int64), dtype=tf.int64 ) result = tf.gather(a, blade_indices, axis=-1) return result
[docs] def select_blades_with_name(self, a: tf.Tensor, blade_names: Union[List[str], str]) -> tf.Tensor: """Takes a geometric algebra tensor and returns a `tf.Tensor` with the blades in blade_names on the last axis. Args: a: Geometric algebra tensor to copy blade_names: Blades to keep Returns: `tf.Tensor` based on `a` with `blade_names` on last axis. """ a = tf.convert_to_tensor(a, dtype_hint=tf.float32) is_single_blade = isinstance(blade_names, str) if is_single_blade: blade_names = [blade_names] blade_signs, blade_indices = get_blade_indices_from_names( blade_names, self.blades) result = blade_signs * self.select_blades(a, blade_indices) if is_single_blade: return result[..., 0] return result
[docs] def inverse(self, a: tf.Tensor) -> tf.Tensor: """Returns the inverted geometric algebra tensor `X^-1` such that `X * X^-1 = 1`. Using Shirokov's inverse algorithm that works in arbitrary dimensions, see https://arxiv.org/abs/2005.04015 Theorem 4. Args: a: Geometric algebra tensor to return inverse for Returns: inverted geometric algebra tensor """ a = tf.convert_to_tensor(a, dtype_hint=tf.float32) n = 2 ** ((len(self.metric) + 1) // 2) u = a for k in range(1, n): c = n / k * self.keep_blades_with_name(u, "") u_minus_c = u - c u = self.geom_prod(a, u_minus_c) if not self.is_pure_kind(u, BladeKind.SCALAR): raise Exception( "Can't invert multi-vector (det U not scalar: %s)." % u) # adj / det return u_minus_c / u[..., :1]
def __call__(self, a: tf.Tensor) -> MultiVector: """Creates a `MultiVector` from a geometric algebra tensor. Mainly used as a wrapper for the algebra's functions for convenience. Args: a: Geometric algebra tensor to return `MultiVector` for Returns: `MultiVector` for `a` """ return MultiVector(tf.convert_to_tensor(a), self)