Package: jaxga
JAXGA - JAX Geometric AlgebraGitHub|DocsJAXGA is a Geometric Algebra package on top ofJAX. It can handle high dimensional algebras by storing only the non-zero basis blade coefficients.
It makes use of JAX's just-in-time (JIT) compilation by first precomputing blade indices and signs and then JITting the function doing the actual calculations.InstallationInstall using pip:pip install jaxgaRequirements:Python 3jaxUsageUnlike most other Geometric Algebra packages, it is not necessary to pre-specify an algebra. JAXGA can either be used with theMultiVectorclass
or by usinglower-level functionswhich is useful for example when using JAX's jit or automatic differentaition.The MultiVector class provides operator overloading and is constructed with an array of values and their corresponding basis blades. The basis blades are encoded as tuples,
for example the multivector2 e_1 + 4 e_23would have the values[2, 4]and the basis blade tuple((1,), (2, 3)).MultiVector exampleimportjax.numpyasjnpfromjaxga.mvimportMultiVectora=MultiVector(values=2*jnp.ones([1],dtype=jnp.float32),indices=((1,),))b=MultiVector(values=4*jnp.ones([2],dtype=jnp.float32),indices=((2,3),))c=a*bprint(c)Output:Multivector(8.0 e_{1, 2, 3})The lower-level functions also deal with values and blades. Functions are provided that take the blades and return a function that does the actual calculation.
The returned function is JITted and can also be automatically differentiated with JAX. Furthermore, some operations like the geometric product take a signature
function that takes a basis vector index and returns their square.Lower-level function exampleimportjax.numpyasjnpfromjaxga.signaturesimportpositive_signaturefromjaxga.ops.multiplyimportget_mv_multiplya_values=2*jnp.ones([1],dtype=jnp.float32)a_indices=((1,),)b_values=4*jnp.ones([1],dtype=jnp.float32)b_indices=((2,3),)mv_multiply,c_indices=get_mv_multiply(a_indices,b_indices,positive_signature)c_values=mv_multiply(a_values,b_values)print("C indices:",c_indices,"C values:",c_values)Output:C indices: ((1, 2, 3),) C values: [8.]Some notesBoth the MultiVector and lower-level function approaches support batches: the axes after the first one (which indexes the basis blades) are treated as batch indices.The MultiVector class can also take a signature in its constructor (default is square to 1 for all basis vectors). Doing operations with
MultiVectors with different signatures is undefined.Thejaxga.signaturessubmodule contains a few predefined signature functions.get_mv_multiplyand similar functions cache their result by their inputs.Theflaxmodulesprovidesflax(a popular neural network library for jax) modules with
Geometric Algebra operations.Because we don't deal with a specific algebra, thedualneeds an input that specifies the dimensionality of the space in which we want to find the dual element.BenchmarksN-d vector * N-d vector, batch size 100, N=range(1, 10), CPUJaxGA stores only the non-zero basis blade coefficients.TFGAandCliffordon
the other hand store all GA elements as full multivectors including all zeros. As a result, JaxGA does better than these for high dimensional algebras.Below is a benchmark of the geometric product of two vectors with increasing dimensionality from 1 to 9. 100 vectors are multiplied at a time.JAXGA (CPU)tfga (CPU)cliffordN-d vector * N-d vector, batch size 100, N=range(1, 50, 5), CPUBelow is a benchmark for higher dimensions that TFGA and Clifford could not handle. Note that the X axis isn't sorted naturally.