Spaces:
Sleeping
Sleeping
| """JSON-friendly handlers that drive the three.js demos with ``spinor_lib``. | |
| Every handler takes a dict of plain JSON values and returns a dict of plain | |
| JSON values. Real tensors become nested lists; complex tensors become nested | |
| lists of ``[re, im]`` pairs. The demos never re-implement any of the maths: | |
| all rotations, boosts, Bloch vectors, partial traces and gradient steps are | |
| computed here by the library and shipped to the browser as numbers. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| import uuid | |
| from collections import OrderedDict | |
| from typing import Callable, Dict | |
| import torch | |
| from spinor_lib import ( | |
| QuaternionRotation, | |
| Spinor, | |
| bloch_vector, | |
| concurrence, | |
| embed_operator, | |
| entanglement_entropy, | |
| expectation, | |
| fidelity, | |
| from_bloch_vector, | |
| generate_boost, | |
| generate_su2_rotation, | |
| inner_product, | |
| kron, | |
| pauli, | |
| purity, | |
| quaternion_to_su2, | |
| rapidity_from_velocity, | |
| reduced_density_matrix, | |
| rotate, | |
| sl2c_to_lorentz, | |
| su2_to_quaternion, | |
| su2_to_so3, | |
| tensor_product, | |
| weyl_current, | |
| ) | |
| REAL = torch.float64 | |
| COMPLEX = torch.complex128 | |
| HANDLERS: Dict[str, Callable[[dict], dict]] = {} | |
| def handler(name: str): | |
| def register(fn): | |
| HANDLERS[name] = fn | |
| return fn | |
| return register | |
| # ------------------------------------------------------------ serialisation | |
| def r2j(t: torch.Tensor): | |
| """Real tensor -> nested lists.""" | |
| return t.detach().cpu().tolist() | |
| def c2j(t: torch.Tensor): | |
| """Complex tensor -> nested lists of [re, im].""" | |
| return torch.stack([t.real, t.imag], dim=-1).detach().cpu().tolist() | |
| def _int(p: dict, key: str, default: int, lo: int, hi: int) -> int: | |
| """An integer parameter clamped to [lo, hi], so a public server cannot be | |
| asked for a million stars or a hundred-thousand-sample sweep.""" | |
| try: | |
| v = int(p.get(key, default)) | |
| except (TypeError, ValueError): | |
| v = default | |
| return max(lo, min(hi, v)) | |
| def _vec3(v, default): | |
| v = torch.as_tensor(v if v is not None else default, dtype=REAL) | |
| if v.shape != (3,): | |
| raise ValueError(f"expected a 3-vector, got shape {tuple(v.shape)}") | |
| if float(v.norm()) < 1e-9: | |
| raise ValueError("vector must be non-zero") | |
| return v | |
| def _state_from_params(p: dict) -> Spinor: | |
| """A single normalised spinor from either Bloch coordinates or components.""" | |
| if p.get("spinor") is not None: | |
| comps = torch.tensor(p["spinor"], dtype=REAL) # [2, 2] as [re, im] | |
| return Spinor(torch.complex(comps[:, 0], comps[:, 1])).normalize() | |
| return from_bloch_vector(_vec3(p.get("bloch"), [0, 0, 1]), dtype=COMPLEX) | |
| # ------------------------------------------- demo 1: Bloch sphere & double cover | |
| def rotation_sweep(p: dict) -> dict: | |
| """Rotate one state through angles 0..max_angle about an axis. | |
| Returns everything the browser needs to scrub the angle without further | |
| requests: the Bloch vector, the raw components, the overlap with the | |
| starting state (which reaches -1 at 2π and +1 at 4π), the unit quaternion | |
| and the SO(3) matrix of every sample. | |
| """ | |
| axis = _vec3(p.get("axis"), [0, 0, 1]) | |
| max_angle = float(p.get("max_angle", 4 * math.pi)) | |
| n = _int(p, "samples", 721, 2, 2001) | |
| psi0 = _state_from_params(p) | |
| angles = torch.linspace(0, max_angle, n, dtype=REAL) | |
| R = generate_su2_rotation(axis, angles) # [n, 2, 2] | |
| psi = rotate(psi0, R) # [n, 2] | |
| return { | |
| "angles": r2j(angles), | |
| "axis": r2j(axis / axis.norm()), | |
| "bloch": r2j(bloch_vector(psi)), | |
| "components": c2j(psi.data), | |
| "overlap": c2j(inner_product(psi0, psi)), | |
| "quaternion": r2j(su2_to_quaternion(R)), | |
| "so3": r2j(su2_to_so3(R)), | |
| "matrix": c2j(R), | |
| } | |
| def precession_sweep(p: dict) -> dict: | |
| """Larmor precession of a cloud of random spins about a field axis. | |
| The time evolution exp(-i B·σ t/2) of a spin-1/2 in a field is exactly an | |
| SU(2) rotation about B by angle ∝ t, so a whole batch is one broadcast | |
| call to ``generate_su2_rotation``. | |
| """ | |
| axis = _vec3(p.get("axis"), [0, 0, 1]) | |
| n_spins = _int(p, "n_spins", 150, 1, 500) | |
| n = _int(p, "samples", 241, 2, 721) | |
| max_angle = float(p.get("max_angle", 4 * math.pi)) | |
| g = torch.Generator().manual_seed(int(p.get("seed", 0))) | |
| psi0 = Spinor(torch.randn(n_spins, 2, dtype=COMPLEX, generator=g)).normalize() | |
| angles = torch.linspace(0, max_angle, n, dtype=REAL) | |
| R = generate_su2_rotation(axis, angles)[:, None] # [n, 1, 2, 2] | |
| psi = rotate(psi0, R) # [n, n_spins, 2] | |
| b = bloch_vector(psi) | |
| return { | |
| "angles": r2j(angles), | |
| "axis": r2j(axis / axis.norm()), | |
| "bloch": r2j(b), | |
| "mean_spin": r2j(b.mean(dim=1)), | |
| } | |
| # ------------------------------------------------------ demo 2: the belt trick | |
| def belt(p: dict) -> dict: | |
| """A belt twisted by ``turns`` full turns, deformed by a homotopy parameter. | |
| The belt's orientation along its length is a path R(s) in SO(3). Lifted to | |
| SU(2) the 2π twist is a path from I to -I (an open arc) while the 4π twist | |
| is a closed loop. Sliding the lifted path along great circles of S³ | |
| towards a fixed rotation gives a family of paths P_t(s); left-multiplying | |
| by P_t(0)⁻¹ pins the near end of the belt to the identity. For 4π the far | |
| end stays put for every t and the twist disappears; for 2π the far end is | |
| forced to tumble, which is exactly why one twist cannot be undone. | |
| ``slerp_su2`` is deliberately *not* used: it takes the shortest arc in | |
| SO(3) and so discards the sign of the SU(2) element, which here is the | |
| whole point. The interpolation is done on unit quaternions instead, | |
| through the library's ``su2_to_quaternion`` / ``quaternion_to_su2``. | |
| """ | |
| turns = float(p.get("turns", 2)) | |
| t = float(p.get("t", 0.0)) | |
| n = _int(p, "segments", 96, 4, 400) | |
| s = torch.linspace(0, 1, n, dtype=REAL) | |
| loop = generate_su2_rotation([0, 0, 1], 2 * math.pi * turns * s) # [n, 2, 2] | |
| anchor = generate_su2_rotation([1, 0, 0], math.pi, dtype=COMPLEX) # contract towards this | |
| q_loop, q_anchor = su2_to_quaternion(loop), su2_to_quaternion(anchor) | |
| lam = t * math.pi / 2 # the two are orthogonal in R⁴, so this is an exact slerp | |
| P = quaternion_to_su2(math.cos(lam) * q_loop + math.sin(lam) * q_anchor) # [n, 2, 2] | |
| R = P[0].conj().mT @ P # pin the near end: R(0) = I | |
| M = su2_to_so3(R) # [n, 3, 3] | |
| # Centreline: integrate the belt's local "along" direction (its z axis). | |
| tangent = M[:, :, 2] # M @ ẑ | |
| length = float(p.get("length", 3.0)) | |
| centre = torch.cumsum(tangent, dim=0) * (length / n) | |
| centre = centre - centre[0] | |
| end_q = su2_to_quaternion(R[-1]) | |
| return { | |
| "s": r2j(s), | |
| "so3": r2j(M), | |
| "centre": r2j(centre), | |
| "quaternion": r2j(su2_to_quaternion(R)), | |
| "lifted_path": r2j(su2_to_quaternion(P)), # the path in S³ before pinning | |
| "end_so3": r2j(M[-1]), | |
| "end_angle": float(2 * torch.arccos(end_q[0].abs().clamp(max=1))), | |
| } | |
| # --------------------------------------- demo 3: relativistic aberration | |
| def _star_field(n: int, seed: int): | |
| g = torch.Generator().manual_seed(seed) | |
| d = torch.randn(n, 3, dtype=REAL, generator=g) | |
| d = d / d.norm(dim=-1, keepdim=True) | |
| mag = torch.rand(n, dtype=REAL, generator=g) | |
| return d, mag | |
| def _sky_grid(n_lat=5, n_lon=12, samples=73): | |
| """Latitude and longitude circles of the sky as polylines, shape [k, samples, 3].""" | |
| u = torch.linspace(0, 2 * math.pi, samples, dtype=REAL) | |
| lines = [] | |
| for i in range(1, n_lat + 1): | |
| th = math.pi * i / (n_lat + 1) | |
| lines.append(torch.stack([math.sin(th) * torch.cos(u), math.sin(th) * torch.sin(u), | |
| torch.full_like(u, math.cos(th))], -1)) | |
| v = torch.linspace(0, math.pi, samples, dtype=REAL) | |
| for j in range(n_lon): | |
| ph = 2 * math.pi * j / n_lon | |
| lines.append(torch.stack([torch.sin(v) * math.cos(ph), torch.sin(v) * math.sin(ph), | |
| torch.cos(v)], -1)) | |
| return torch.stack(lines) | |
| def _aberrate(directions: torch.Tensor, A: torch.Tensor): | |
| """Transform sky directions (unit vectors pointing *at* the sources). | |
| A photon arriving from direction n̂ has null wave-vector k = (1, -n̂). The | |
| right-handed spinor with Bloch vector -n̂ has Weyl current exactly k, so | |
| boosting the spinor and reading its current back gives Λk, from which the | |
| new direction and the Doppler factor follow. | |
| """ | |
| psi = from_bloch_vector(-directions, dtype=COMPLEX) | |
| k = weyl_current(rotate(psi, A), "right") # [..., 4] = (D, -D n̂') | |
| doppler = k[..., 0] | |
| new_dir = -k[..., 1:] / doppler[..., None] | |
| return new_dir, doppler | |
| def aberration(p: dict) -> dict: | |
| """The celestial sphere seen by an observer moving at speed β. | |
| Boosts are elements of SL(2,C) acting on Weyl spinors, and SL(2,C) acts | |
| on the sky as a Möbius transformation: circles stay circles, stars bunch | |
| up ahead of you and blueshift there. | |
| """ | |
| beta = float(p.get("beta", 0.0)) | |
| beta = max(-0.999, min(0.999, beta)) | |
| direction = _vec3(p.get("direction"), [0, 0, 1]) | |
| n_stars = _int(p, "n_stars", 1500, 1, 5000) | |
| seed = int(p.get("seed", 1)) | |
| eta = rapidity_from_velocity(torch.tensor(beta, dtype=REAL)) | |
| # Λ(η) carries rest -> velocity β·n, i.e. maps the moving frame's | |
| # coordinates into ours; the observer's own view needs the inverse boost. | |
| A = generate_boost(direction, -eta, handedness="right", dtype=COMPLEX) | |
| stars, mag = _star_field(n_stars, seed) | |
| new_dir, doppler = _aberrate(stars, A) | |
| grid = _sky_grid() | |
| grid_new, _ = _aberrate(grid, A) | |
| return { | |
| "beta": beta, | |
| "gamma": float(torch.cosh(eta)), | |
| "rapidity": float(eta), | |
| "direction": r2j(direction / direction.norm()), | |
| "stars": r2j(stars), | |
| "magnitude": r2j(mag), | |
| "aberrated": r2j(new_dir), | |
| "doppler": r2j(doppler), | |
| "grid": r2j(grid), | |
| "grid_aberrated": r2j(grid_new), | |
| "sl2c": c2j(A), | |
| "lorentz": r2j(sl2c_to_lorentz(A, "right")), | |
| } | |
| # ---------------------------------------------- demo 4: two-qubit entanglement | |
| CNOT = torch.tensor( | |
| [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], dtype=COMPLEX | |
| ) | |
| def two_qubit(p: dict) -> dict: | |
| """cos θ|00⟩ + sin θ|11⟩ followed by local rotations on each qubit. | |
| The reduced density matrices come from a partial trace; their Bloch | |
| vectors shrink inside the ball as the concurrence grows. Local rotations | |
| move the arrows but cannot change concurrence, entropy or purity. | |
| """ | |
| theta = float(p.get("theta", 0.0)) | |
| up = Spinor(torch.tensor([1, 0], dtype=COMPLEX)) | |
| psi = tensor_product(up, up) | |
| ry = generate_su2_rotation([0, 1, 0], 2 * theta, dtype=COMPLEX) | |
| psi = rotate(psi, embed_operator(ry, (2, 2), target=0)) | |
| psi = rotate(psi, CNOT) | |
| for target in (0, 1): | |
| loc = p.get(f"local{target}") or {} | |
| angle = float(loc.get("angle", 0.0)) | |
| if abs(angle) > 1e-12: | |
| R = generate_su2_rotation(_vec3(loc.get("axis"), [0, 0, 1]), angle, dtype=COMPLEX) | |
| psi = rotate(psi, embed_operator(R, (2, 2), target=target)) | |
| rho0 = reduced_density_matrix(psi, (2, 2), keep=0) | |
| rho1 = reduced_density_matrix(psi, (2, 2), keep=1) | |
| sig = pauli(dtype=COMPLEX) | |
| corr = torch.stack( | |
| [torch.stack([expectation(psi, kron(sig[i], sig[j])) for j in range(3)]) for i in range(3)] | |
| ) | |
| return { | |
| "amplitudes": c2j(psi.data), | |
| "probabilities": r2j(psi.data.abs() ** 2), | |
| "bloch0": r2j(bloch_vector(rho0)), | |
| "bloch1": r2j(bloch_vector(rho1)), | |
| "purity0": float(purity(rho0)), | |
| "purity1": float(purity(rho1)), | |
| "concurrence": float(concurrence(psi)), | |
| "entropy_bits": float(entanglement_entropy(psi, (2, 2), keep=0, base=2)), | |
| "correlation": r2j(corr), | |
| "rho0": c2j(rho0), | |
| } | |
| # ------------------------------------------ demo 5: learning a rotation live | |
| class _Session: | |
| def __init__(self, p: dict): | |
| seed = int(p.get("seed", 0)) | |
| g = torch.Generator().manual_seed(seed) | |
| n = _int(p, "n_points", 64, 4, 512) | |
| noise = float(p.get("noise", 0.02)) | |
| lr = float(p.get("lr", 0.05)) | |
| axis = torch.randn(3, generator=g) | |
| angle = torch.rand(1, generator=g).item() * 1.6 * math.pi + 0.4 * math.pi | |
| self.target = generate_su2_rotation(axis, torch.tensor(angle)) | |
| self.x = Spinor(torch.randn(n, 2, dtype=torch.complex64, generator=g)).normalize() | |
| y = rotate(self.x, self.target) | |
| y = Spinor(y.data + noise * torch.randn(y.data.shape, dtype=y.dtype, generator=g)) | |
| self.y = y.normalize() | |
| self.model = QuaternionRotation() | |
| self.opt = torch.optim.Adam(self.model.parameters(), lr=lr) | |
| self.step = 0 | |
| with torch.no_grad(): | |
| self.loss = (1 - fidelity(self.model(self.x), self.y)).mean().item() | |
| def train(self, steps: int): | |
| for _ in range(steps): | |
| self.opt.zero_grad() | |
| loss = (1 - fidelity(self.model(self.x), self.y)).mean() | |
| loss.backward() | |
| self.opt.step() | |
| self.step += 1 | |
| self.loss = loss.item() | |
| def snapshot(self) -> dict: | |
| with torch.no_grad(): | |
| R = self.model.matrix | |
| pred = self.model(self.x) | |
| err = (su2_to_so3(R) - su2_to_so3(self.target)).abs().max().item() | |
| grad = self.model.quaternion.grad | |
| return { | |
| "step": self.step, | |
| "loss": self.loss, | |
| "quaternion": r2j(su2_to_quaternion(R)), | |
| "so3": r2j(su2_to_so3(R)), | |
| "predicted": r2j(bloch_vector(pred)), | |
| "so3_error": err, | |
| "grad_norm": float(grad.norm()) if grad is not None else 0.0, | |
| } | |
| _SESSIONS: "OrderedDict[str, _Session]" = OrderedDict() | |
| _MAX_SESSIONS = 16 | |
| def learn_reset(p: dict) -> dict: | |
| """Start a fresh ``QuaternionRotation`` fit against a hidden rotation.""" | |
| sid = uuid.uuid4().hex | |
| sess = _Session(p) | |
| _SESSIONS[sid] = sess | |
| while len(_SESSIONS) > _MAX_SESSIONS: | |
| _SESSIONS.popitem(last=False) | |
| with torch.no_grad(): | |
| return { | |
| "session": sid, | |
| "inputs": r2j(bloch_vector(sess.x)), | |
| "targets": r2j(bloch_vector(sess.y)), | |
| "target_quaternion": r2j(su2_to_quaternion(sess.target)), | |
| "target_so3": r2j(su2_to_so3(sess.target)), | |
| **sess.snapshot(), | |
| } | |
| def learn_step(p: dict) -> dict: | |
| """Run a few optimiser steps and report the new state.""" | |
| sess = _SESSIONS.get(str(p.get("session", ""))) | |
| if sess is None: | |
| raise KeyError("unknown or expired session; call learn_reset") | |
| sess.train(_int(p, "steps", 1, 1, 200)) | |
| return sess.snapshot() | |