From 3214fcfeb5a0ff385f5cb75d9046c51e4152d9bb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=C3=A1=C5=A1=20Kulh=C3=A1nek?= <jonas.kulhanek@live.com>
Date: Thu, 22 Jun 2023 04:40:22 +0200
Subject: [PATCH] Fix spherical harmonics #2081 (#2112)

---
 nerfstudio/utils/math.py |  4 ++--
 tests/utils/test_math.py | 16 ++++++++++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)
 create mode 100644 tests/utils/test_math.py

diff --git a/nerfstudio/utils/math.py b/nerfstudio/utils/math.py
index 866d280f8..788a0520d 100644
--- a/nerfstudio/utils/math.py
+++ b/nerfstudio/utils/math.py
@@ -80,12 +80,12 @@ def components_from_spherical_harmonics(
         components[..., 16] = 2.5033429417967046 * x * y * (xx - yy)
         components[..., 17] = 1.7701307697799304 * y * z * (3 * xx - yy)
         components[..., 18] = 0.9461746957575601 * x * y * (7 * zz - 1)
-        components[..., 19] = 0.6690465435572892 * y * (7 * zz - 3)
+        components[..., 19] = 0.6690465435572892 * y * z * (7 * zz - 3)
         components[..., 20] = 0.10578554691520431 * (35 * zz * zz - 30 * zz + 3)
         components[..., 21] = 0.6690465435572892 * x * z * (7 * zz - 3)
         components[..., 22] = 0.47308734787878004 * (xx - yy) * (7 * zz - 1)
         components[..., 23] = 1.7701307697799304 * x * z * (xx - 3 * yy)
-        components[..., 24] = 0.4425326924449826 * (xx * (xx - 3 * yy) - yy * (3 * xx - yy))
+        components[..., 24] = 0.6258357354491761 * (xx * (xx - 3 * yy) - yy * (3 * xx - yy))
 
     return components
 
diff --git a/tests/utils/test_math.py b/tests/utils/test_math.py
new file mode 100644
index 000000000..c0e1edc43
--- /dev/null
+++ b/tests/utils/test_math.py
@@ -0,0 +1,16 @@
+import pytest
+import torch
+
+from nerfstudio.utils.math import components_from_spherical_harmonics
+
+
+@pytest.mark.parametrize("components", list(range(1, 5 + 1)))
+def test_spherical_harmonics(components):
+    torch.manual_seed(0)
+    N = 1000000
+
+    dx = torch.normal(0, 1, size=(N, 3))
+    dx = dx / torch.linalg.norm(dx, dim=-1, keepdim=True)
+    sh = components_from_spherical_harmonics(components, dx)
+    matrix = (sh.T @ sh) / N * 4 * torch.pi
+    torch.testing.assert_close(matrix, torch.eye(components**2), rtol=0, atol=1e-2)