xagm 0.1.1__tar.gz

xagm-0.1.1/PKG-INFO

@@ -0,0 +1,18 @@
+Metadata-Version: 2.4
+Name: xagm
+Version: 0.1.1
+Description-Content-Type: text/markdown
+Requires-Dist: jax
+Requires-Dist: jaxlib
+Requires-Dist: jaxtyping
+
+XAGM is a Riemannian Differentiable Geometry engine which stands for Accelerated Autodiff Geometry Multi-dimensional. It deals exclusively in Riemannian SPD metrics, and it is MANDATORY the metrics are Symmetric Positive Definite (SPD) for it to work.
+It offers a vast array of functions, with 4 modules to call upon, them being metrics, linear, vectors, and calc. Vectors deal mainly with linear algebra adjacent functions with respect to the metric tensor. Speaking of the metric tensor, XAGM allows you to use fwdmet to create a pullback metric. 
+
+The crown jewels of XAGM would be christoffel(), geoexp_solver(), geolog_solver(), and geodist(), with geoexp_solver consistently performing at sub millisecond speeds, and geolog_solver being in the comfortable range of 2-20ms each run depending on how many steps are given to the solver. 
+
+XAGM has been benchmarked (quite unofficially so you are free to do your own runtime checks) and observed to outperform basically every other geometry application in numpy and the dominating Geometry powerhouses. You are highly encouraged, however, to confirm that yourself too. 
+
+XAGM is a bit hard to use at first since it expects a decent background in maths for most of the functions and a clear understanding of how to use JAX native functions like vmap and jit along with static_argnums and static_argnames, but, overall, if you behave nicely and pass clean arrays into it, it will reward you. Documentation on this project will be coming soon! (or never at all. No in between.)
+
+

xagm-0.1.1/README.md

@@ -0,0 +1,10 @@
+XAGM is a Riemannian Differentiable Geometry engine which stands for Accelerated Autodiff Geometry Multi-dimensional. It deals exclusively in Riemannian SPD metrics, and it is MANDATORY the metrics are Symmetric Positive Definite (SPD) for it to work.
+It offers a vast array of functions, with 4 modules to call upon, them being metrics, linear, vectors, and calc. Vectors deal mainly with linear algebra adjacent functions with respect to the metric tensor. Speaking of the metric tensor, XAGM allows you to use fwdmet to create a pullback metric. 
+
+The crown jewels of XAGM would be christoffel(), geoexp_solver(), geolog_solver(), and geodist(), with geoexp_solver consistently performing at sub millisecond speeds, and geolog_solver being in the comfortable range of 2-20ms each run depending on how many steps are given to the solver. 
+
+XAGM has been benchmarked (quite unofficially so you are free to do your own runtime checks) and observed to outperform basically every other geometry application in numpy and the dominating Geometry powerhouses. You are highly encouraged, however, to confirm that yourself too. 
+
+XAGM is a bit hard to use at first since it expects a decent background in maths for most of the functions and a clear understanding of how to use JAX native functions like vmap and jit along with static_argnums and static_argnames, but, overall, if you behave nicely and pass clean arrays into it, it will reward you. Documentation on this project will be coming soon! (or never at all. No in between.)
+
+

xagm-0.1.1/pyproject.toml

@@ -0,0 +1,18 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "xagm"
+version = "0.1.1"
+dependencies = [
+    "jax", "jaxlib", "jaxtyping"
+]
+readme = "README.md"
+# Add any dependencies you need here, e.g., ["jax", "jaxlib"]
+
+[tool.setuptools.packages.find]
+# This tells setuptools to look for your code inside the 'src' folder
+where = ["src"]
+
+[tool.setuptools]

xagm-0.1.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

xagm-0.1.1/src/xagm/basis/__init__.py

@@ -0,0 +1,4 @@
+from .linear import (grid, line, ang)
+
+from .metrics import (euclid, iprod, norm, fwdmet, revmet, metinv, 
+metinterp)

xagm-0.1.1/src/xagm/basis/linear.py

@@ -0,0 +1,33 @@
+import geoutils as us
+import jax 
+import jax.numpy as jnp
+from geoutils import Vector, Matrix, Scalar, Tensor, JAXArray
+from basis import metrics as mtc
+
+def grid(idx: JAXArray, dimens: tuple):
+    fg = jnp.unravel_index(idx, dimens)
+    g = fg[::-1]
+    ng = jnp.stack(g, axis=-1)
+
+    return ng
+
+static_argnums = (2,)
+def line(p1: Vector, p2: Vector, segs: int) -> Matrix:
+    t = jnp.linspace(0, 1, segs)[:, jnp.newaxis]
+
+    l = p1 + (t * (p2 - p1))
+
+    return l
+
+
+
+def ang(g: Matrix, u: Vector, v: Vector) -> Scalar:
+
+    numerator = mtc.iprod(g, u, v)
+    den1 = mtc.norm(g, u)
+    den2 = mtc.norm(g, v)
+
+    angle = us.div(numerator, (den1 * den2))
+    safe_cos = jnp.clip(angle, -1.0 + 1e-8, 1.0 - 1e-8)
+    
+    return jnp.arccos(safe_cos)

xagm-0.1.1/src/xagm/basis/metrics.py

@@ -0,0 +1,67 @@
+import geoutils as us
+import jax 
+import jax.numpy as jnp
+from geoutils import Vector, Matrix, Scalar, Tensor, JAXArray
+
+
+def euclid(x: Vector) -> Matrix:
+    return jnp.eye(x.shape[-1])
+
+def iprod(g: Matrix, u: Vector|Matrix, v: Vector|Matrix) -> Vector:
+    return jnp.einsum('...i, ...ij, ...j -> ...', u, g, v)
+
+def norm(g: Matrix, u: Vector) -> Scalar: 
+    return jnp.sqrt(jnp.maximum(iprod(g, u, u), 1e-16))
+
+
+static_argnums = (0,)
+def fwdmet(f, v: Vector) -> Matrix:
+    J = jax.jacfwd(f)(v)
+    nJ = J.reshape(-1, v.shape[-1])
+    return jnp.einsum('ai, aj -> ij', nJ, nJ)
+
+static_argnums = (0,)
+def revmet(f, v: Vector) -> Matrix:
+    J = jax.jacrev(f)(v)
+    nJ = J.reshape(-1, v.shape[-1])
+    return jnp.einsum('ai, aj -> ij', nJ, nJ)
+
+def metinv(g: Matrix) -> Matrix:
+    vals, vecs = jnp.linalg.eigh(g)
+    inv_vals = us.div(1.0, jnp.maximum(vals, 1e-12))
+    met = jnp.einsum('ik, k, jk -> ij', vecs, inv_vals, vecs)
+
+    return met
+
+def metinterp(g0: Matrix, v0: Vector,
+ g1: Matrix, v1: Vector, 
+ target: Vector) -> Matrix:
+    
+    vals0, vecs0 = jnp.linalg.eigh(g0)
+    logvals0 = jnp.log(jnp.maximum(vals0, 1e-7))
+    lg0 = jnp.einsum('ik, k, jk -> ij', vecs0, logvals0, vecs0)
+
+    vals1, vecs1 = jnp.linalg.eigh(g1)
+    logvals1 = jnp.log(jnp.maximum(vals1, 1e-7))
+    lg1 = jnp.einsum('ik, k, jk -> ij', vecs1, logvals1, vecs1)
+
+    d = v1 - v0
+    p = target - v0
+
+    t = us.div(jnp.dot(p, d),jnp.dot(d, d))
+    t = jnp.clip(t, 0.0, 1.0)
+
+    interp = (1.0 - t) * lg0 + (t * lg1)
+
+    intvals, intvecs = jnp.linalg.eigh(interp)
+
+    ival = jnp.exp(intvals)
+
+    ig = jnp.einsum('ik, k, jk -> ij', intvecs, ival, intvecs)
+
+    return ig
+
+
+ 
+
+

xagm-0.1.1/src/xagm/geoutils.py

@@ -0,0 +1,27 @@
+import jax
+from jax import config
+
+import os
+os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'
+config.update("jax_enable_x64", True)
+
+import jax
+import jax.numpy as jnp
+from jaxtyping import Float64, Array
+
+@jax.jit
+def div(a, b):
+    safe = b != 0
+    den = jnp.where(safe, b, 1.0)
+
+    return jnp.where(safe, a/den, 0.0)
+
+
+eps = 1e-15
+
+
+type Scalar = Float64[Array, ""] # 0D [1]
+type Vector = Float64[Array, "N"] # 1D [1, 2]
+type Matrix = Float64[Array, "M N"] # 2D [[1, 2], [2, 1]]
+type Tensor = Float64[Array, "*batch M N O"] # 3-D+
+type JAXArray = Float64[Array]

xagm-0.1.1/src/xagm/manifolds/__init__.py

@@ -0,0 +1,7 @@
+from .calc import (christoffel, geoexp_solver, geolog_solver, geodist)
+
+from .vectors import (nrml,
+ scalproj, 
+vectproj, 
+ rejvect,
+ unitize)

xagm-0.1.1/src/xagm/manifolds/calc.py

@@ -0,0 +1,103 @@
+import geoutils as us
+import jax 
+import jax.numpy as jnp
+from geoutils import Vector, Matrix, Scalar, Tensor, JAXArray
+
+from basis import metrics as mtc
+
+def christoffel(func, x: Vector) -> Matrix:
+    
+    g = mtc.fwdmet(func, x)
+    ginv = mtc.metinv(g)
+    mtc_func = lambda v: mtc.fwdmet(func, v)
+
+    __,dg_raw = jax.vmap(lambda v: jax.jvp(mtc_func, (x,), (v,)))(jnp.eye(x.shape[0]))
+
+    dg = jnp.moveaxis(dg_raw, 0, -1)
+
+    term1 = jnp.transpose(dg, axes=[1, 2, 0])
+    term2 = jnp.transpose(dg, axes=[0, 1, 2])
+    term3 = jnp.transpose(dg, axes=[2, 0, 1])
+
+    contract1 = 0.5 * ginv
+    contract2 = term1 + term2 - term3
+    gamma = jnp.einsum('kl, lij -> kij', contract1, contract2)
+
+    return gamma
+
+import diffrax
+
+
+def geoexp_term(t, state, args) -> Vector:
+    dim = state.shape[0] // 3
+
+    x = state[:dim] 
+    v = state[dim:2*dim]
+    y = state[2*dim:]
+
+    func = args['func']
+
+    gamma = christoffel(func, x)
+
+    v_dot = -jnp.einsum('kij, i, j -> k', gamma, v, v)
+
+    dvecdt = -jnp.einsum('kij, i, j -> k', gamma, v, y) 
+
+    return jnp.concatenate([v, v_dot, dvecdt])
+
+
+def geoexp_solver(p: Vector, v: Vector, mapped_func, vt: Vector, steps = 4096) -> Vector:
+
+    state = jnp.concatenate([p, v, vt])
+
+    solution = diffrax.diffeqsolve(
+        terms = diffrax.ODETerm(geoexp_term),
+        solver = diffrax.Tsit5(),
+        t0=0,
+        t1=1,
+        dt0=1e-2,
+        y0=state,
+        args = {'func': mapped_func},
+        stepsize_controller = diffrax.PIDController(rtol=1e-8, atol=1e-10),
+        saveat=diffrax.SaveAt(t1=True),
+        adjoint = diffrax.DirectAdjoint(),
+        max_steps = steps,
+        throw = False
+    )
+
+    result = solution.ys[0]
+
+    dim = p.shape[0]
+    final_pos = result[:dim]
+    final_vel = result[dim:2*dim]
+    transported_v = result[2*dim:]
+
+    return final_pos, final_vel, transported_v
+
+def geolog_solver(p: Vector, q: Vector, mapped_func, steps: int) -> Vector:
+    
+    def shoot(v_guess):
+        pos, _, _ = geoexp_solver(p, v_guess, mapped_func, jnp.zeros_like(p))
+        return pos
+
+    v = q-p
+    J = jax.jacobian(shoot)(v)
+
+    def bodyfun(i, v):
+        error = shoot(v) - q
+        #J = jax.jacobian(shoot)(v)
+        delta = jnp.linalg.solve(J, error)
+        return v - delta
+
+    final_v = jax.lax.fori_loop(0, steps, bodyfun, v)
+    return final_v
+
+
+def geodist(p: Vector, q: Vector, mapped_func, steps: int) -> Scalar:
+    v = geolog_solver(p, q, mapped_func, steps)
+    g = mtc.fwdmet(mapped_func, p)
+    dist = mtc.norm(g, v)
+    return dist
+
+
+

xagm-0.1.1/src/xagm/manifolds/vectors.py

@@ -0,0 +1,56 @@
+
+import geoutils as us
+import jax 
+import jax.numpy as jnp
+from geoutils import Vector, Matrix, Scalar, Tensor, JAXArray
+
+from basis import metrics as mtc
+
+def nrml(g: Matrix, basis: Matrix) -> Matrix:
+
+    nvals, vecs = jnp.linalg.eigh(g)
+    vals = jnp.maximum(nvals, 0.0)
+
+    L = jnp.sqrt(vals)[:, None] * vecs.T 
+    bflat = basis @ L.T 
+
+    Q, R = jnp.linalg.qr(bflat.T) 
+    linvt = us.div(vecs, jnp.sqrt(vals))
+
+    ortho = Q.T @ linvt.T 
+    det = jnp.linalg.det(ortho @ L.T) > 0 
+    check = jnp.where(det, 1.0, -1.0) 
+    
+    northo = ortho.at[0, :].multiply(check) 
+
+    return northo
+
+#dot product territory
+
+def scalproj(g: Matrix, a: Vector, b: Vector) -> Scalar: 
+    
+    norm = mtc.norm(g, b)
+    prod = us.div(mtc.iprod(g, a, b), norm)
+
+    return prod
+
+
+def vectproj(g: Matrix, a: Vector, b: Vector) -> Vector:
+
+    term = mtc.iprod(g, b, b)
+    prod = us.div(mtc.iprod(g, a, b), term)
+    proj = prod * b
+
+    return proj
+
+
+def rejvect(g: Matrix, a: Vector, b: Vector) -> Vector:
+
+    proj = vectproj(g, a, b)
+    reject = a - proj
+
+    return reject
+
+
+def unitize(g: Matrix, u: Vector) -> Vector: 
+    return us.div(u, mtc.norm(g, u))

xagm-0.1.1/src/xagm.egg-info/PKG-INFO

@@ -0,0 +1,18 @@
+Metadata-Version: 2.4
+Name: xagm
+Version: 0.1.1
+Description-Content-Type: text/markdown
+Requires-Dist: jax
+Requires-Dist: jaxlib
+Requires-Dist: jaxtyping
+
+XAGM is a Riemannian Differentiable Geometry engine which stands for Accelerated Autodiff Geometry Multi-dimensional. It deals exclusively in Riemannian SPD metrics, and it is MANDATORY the metrics are Symmetric Positive Definite (SPD) for it to work.
+It offers a vast array of functions, with 4 modules to call upon, them being metrics, linear, vectors, and calc. Vectors deal mainly with linear algebra adjacent functions with respect to the metric tensor. Speaking of the metric tensor, XAGM allows you to use fwdmet to create a pullback metric. 
+
+The crown jewels of XAGM would be christoffel(), geoexp_solver(), geolog_solver(), and geodist(), with geoexp_solver consistently performing at sub millisecond speeds, and geolog_solver being in the comfortable range of 2-20ms each run depending on how many steps are given to the solver. 
+
+XAGM has been benchmarked (quite unofficially so you are free to do your own runtime checks) and observed to outperform basically every other geometry application in numpy and the dominating Geometry powerhouses. You are highly encouraged, however, to confirm that yourself too. 
+
+XAGM is a bit hard to use at first since it expects a decent background in maths for most of the functions and a clear understanding of how to use JAX native functions like vmap and jit along with static_argnums and static_argnames, but, overall, if you behave nicely and pass clean arrays into it, it will reward you. Documentation on this project will be coming soon! (or never at all. No in between.)
+
+

xagm-0.1.1/src/xagm.egg-info/SOURCES.txt

@@ -0,0 +1,15 @@
+README.md
+pyproject.toml
+src/xagm/__init__.py
+src/xagm/geoutils.py
+src/xagm.egg-info/PKG-INFO
+src/xagm.egg-info/SOURCES.txt
+src/xagm.egg-info/dependency_links.txt
+src/xagm.egg-info/requires.txt
+src/xagm.egg-info/top_level.txt
+src/xagm/basis/__init__.py
+src/xagm/basis/linear.py
+src/xagm/basis/metrics.py
+src/xagm/manifolds/__init__.py
+src/xagm/manifolds/calc.py
+src/xagm/manifolds/vectors.py

xagm-0.1.1/src/xagm.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

xagm-0.1.1/src/xagm.egg-info/requires.txt

@@ -0,0 +1,3 @@
+jax
+jaxlib
+jaxtyping

xagm-0.1.1/src/xagm.egg-info/top_level.txt

@@ -0,0 +1 @@
+xagm