From 7e94fef19e61d185ab405c32642a0598bc3179f2 Mon Sep 17 00:00:00 2001
From: Aaron Fenyes <aaron.fenyes@fareycircles.ooo>
Date: Sat, 6 Jul 2024 21:32:43 -0700
Subject: [PATCH] Improve random vector generator

---
 engine-proto/gram-test/Engine.jl | 27 +++++++++++++++++++++------
 1 file changed, 21 insertions(+), 6 deletions(-)

diff --git a/engine-proto/gram-test/Engine.jl b/engine-proto/gram-test/Engine.jl
index 44c5285..a0a59be 100644
--- a/engine-proto/gram-test/Engine.jl
+++ b/engine-proto/gram-test/Engine.jl
@@ -8,14 +8,29 @@ export rand_on_shell, Q, DescentHistory, realize_gram
 
 # === guessing ===
 
-##[TO DO] write a test to confirm that the outputs are on the correct shells
-##[TO DO] this can fail at generating spacelike vectors
-function rand_on_shell(rng::AbstractRNG, shells::Array{T}) where T <: Number
-  space_parts = randn(rng, T, 4, size(shells, 1))
-  space_self_prods = [dot(x, x) for x in eachcol(space_parts)]
-  return [space_parts; sqrt.(space_self_prods .- shells)']
+sconh(t, u) = 0.5*(exp(t) + u*exp(-t))
+
+function rand_on_sphere(rng::AbstractRNG, ::Type{T}, n) where T
+  out = randn(rng, T, n)
+  tries_left = 2
+  while dot(out, out) < 1e-6 && tries_left > 0
+    out = randn(rng, T, n)
+    tries_left -= 1
+  end
+  normalize(out)
 end
 
+##[TO DO] write a test to confirm that the outputs are on the correct shells
+function rand_on_shell(rng::AbstractRNG, shell::T) where T <: Number
+  space_part = rand_on_sphere(rng, T, 4)
+  rapidity = randn(rng, T)
+  sig = sign(shell)
+  [sconh(rapidity, sig)*space_part; sconh(rapidity, -sig)]
+end
+
+rand_on_shell(rng::AbstractRNG, shells::Array{T}) where T <: Number =
+  hcat([rand_on_shell(rng, sh) for sh in shells]...)
+
 rand_on_shell(shells::Array{<:Number}) = rand_on_shell(Random.default_rng(), shells)
 
 # === Gram matrix realization ===