File size: 3,455 Bytes
6cd35b4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
# Copyright (c) Microsoft Corporation.
import math
import torch
from torch import nn
from .spherical_armonics import SH as SH_analytic
class SphericalHarmonics(nn.Module):
"""
Spherical Harmonics locaiton encoder
"""
def __init__(self, legendre_polys: int = 10, harmonics_calculation="analytic"):
"""
legendre_polys: determines the number of legendre polynomials.
more polynomials lead more fine-grained resolutions
calculation of spherical harmonics:
analytic uses pre-computed equations. This is exact, but works only up to degree 50,
closed-form uses one equation but is computationally slower (especially for high degrees)
"""
super(SphericalHarmonics, self).__init__()
self.L, self.M = int(legendre_polys), int(legendre_polys)
self.embedding_dim = self.L * self.M
if harmonics_calculation == "closed-form":
self.SH = SH_closed_form
elif harmonics_calculation == "analytic":
self.SH = SH_analytic
def forward(self, lonlat):
lon, lat = lonlat[:, 0], lonlat[:, 1]
# convert degree to rad
phi = torch.deg2rad(lon + 180)
theta = torch.deg2rad(lat + 90)
"""
greater_than_50 = (lon > 50).any() or (lat > 50).any()
if greater_than_50:
SH = SH_closed_form
else:
SH = SH_analytic
"""
SH = self.SH
Y = []
for l in range(self.L):
for m in range(-l, l + 1):
y = SH(m, l, phi, theta)
if isinstance(y, float):
y = y * torch.ones_like(phi)
if y.isnan().any():
print(m, l, y)
Y.append(y)
return torch.stack(Y, dim=-1)
####################### Spherical Harmonics utilities ########################
# Code copied from https://github.com/BachiLi/redner/blob/master/pyredner/utils.py
# Code adapted from "Spherical Harmonic Lighting: The Gritty Details", Robin Green
# http://silviojemma.com/public/papers/lighting/spherical-harmonic-lighting.pdf
def associated_legendre_polynomial(l, m, x):
pmm = torch.ones_like(x)
if m > 0:
somx2 = torch.sqrt((1 - x) * (1 + x))
fact = 1.0
for i in range(1, m + 1):
pmm = pmm * (-fact) * somx2
fact += 2.0
if l == m:
return pmm
pmmp1 = x * (2.0 * m + 1.0) * pmm
if l == m + 1:
return pmmp1
pll = torch.zeros_like(x)
for ll in range(m + 2, l + 1):
pll = ((2.0 * ll - 1.0) * x * pmmp1 - (ll + m - 1.0) * pmm) / (ll - m)
pmm = pmmp1
pmmp1 = pll
return pll
def SH_renormalization(l, m):
return math.sqrt(
(2.0 * l + 1.0) * math.factorial(l - m) / (4 * math.pi * math.factorial(l + m))
)
def SH_closed_form(m, l, phi, theta):
if m == 0:
return SH_renormalization(l, m) * associated_legendre_polynomial(
l, m, torch.cos(theta)
)
elif m > 0:
return (
math.sqrt(2.0)
* SH_renormalization(l, m)
* torch.cos(m * phi)
* associated_legendre_polynomial(l, m, torch.cos(theta))
)
else:
return (
math.sqrt(2.0)
* SH_renormalization(l, -m)
* torch.sin(-m * phi)
* associated_legendre_polynomial(l, -m, torch.cos(theta))
)
|