# Source code for pde.tools.spherical

r"""
Module collecting functions for handling spherical geometry

The coordinate systems use the following convention for polar coordinates
:math:(r, \phi), where :math:r is the radial coordinate and :math:\phi is
the polar angle:

.. math::
\begin{cases}
x = r \cos(\phi) &\\
y = r \sin(\phi) &
\end{cases}
\text{for} \; r \in [0, \infty] \;
\text{and} \; \phi \in [0, 2\pi)

Similarly, for spherical coordinates :math:(r, \theta, \phi), where :math:r
is the radial coordinate, :math:\theta is the azimuthal angle, and
:math:\phi is the polar angle, we use

.. math::
\begin{cases}
x = r \sin(\theta) \cos(\phi) &\\
y = r \sin(\theta) \sin(\phi) &\\
z = r \cos(\theta)
\end{cases}
\text{for} \; r \in [0, \infty], \;
\theta \in [0, \pi], \; \text{and} \;
\phi \in [0, 2\pi)

The module also provides functions for handling spherical harmonics.
These spherical harmonics are described by the degree :math:l and the order
:math:m or, alternatively, by the mode :math:k. The relation between these
values is

.. math::
k = l(l + 1) + m

and

.. math::
l &= \text{floor}(\sqrt{k}) \\
m &= k - l(l + 1)

We will use these indices interchangeably, although the mode :math:k is
preferred internally. Note that we also consider axisymmetric spherical
harmonics, where the order is always zero and the degree :math:l and the mode
:math:k are thus identical.

.. autosummary::
:nosignatures:

spherical_index_k
spherical_index_lm
spherical_index_count
spherical_index_count_optimal
spherical_harmonic_symmetric
spherical_harmonic_real
spherical_harmonic_real_k

.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
"""

import itertools
from typing import Callable, Optional, Tuple, TypeVar

import numpy as np
from scipy.special import sph_harm

from .cache import cached_method
from .numba import jit

π = float(np.pi)

TNumArr = TypeVar("TNumArr", float, np.ndarray)

[docs]def radius_from_volume(volume: TNumArr, dim: int) -> TNumArr:
"""Return the radius of a sphere with a given volume

Args:
volume (float or :class:~numpy.ndarray): Volume of the sphere
dim (int): Dimension of the space

Returns:
float or :class:~numpy.ndarray: Radius of the sphere
"""
if dim == 1:
return volume / 2  # type: ignore
elif dim == 2:
return np.sqrt(volume / π)  # type: ignore
elif dim == 3:
return (3 * volume / (4 * π)) ** (1 / 3)  # type: ignore
else:
raise NotImplementedError(f"Cannot calculate the radius in {dim} dimensions")

[docs]def make_radius_from_volume_compiled(dim: int) -> Callable[[TNumArr], TNumArr]:
"""Return a function calculating the radius of a sphere with a given volume

Args:
dim (int): Dimension of the space

Returns:
function: A function that takes a volume and returns the radius
"""
if dim == 1:

return volume / 2

elif dim == 2:

return np.sqrt(volume / π)

elif dim == 3:

return (3 * volume / (4 * π)) ** (1 / 3)

else:
raise NotImplementedError(f"Cannot calculate the radius in {dim} dimensions")
return jit(radius_from_volume)  # type: ignore

[docs]def volume_from_radius(radius: TNumArr, dim: int) -> TNumArr:
"""Return the volume of a sphere with a given radius

Args:
radius (float or :class:~numpy.ndarray): Radius of the sphere
dim (int): Dimension of the space

Returns:
float or :class:~numpy.ndarray: Volume of the sphere
"""
if dim == 1:
return 2 * radius  # type: ignore
elif dim == 2:
return π * radius ** 2  # type: ignore
elif dim == 3:
return 4 * π / 3 * radius ** 3  # type: ignore
else:
raise NotImplementedError(f"Cannot calculate the volume in {dim} dimensions")

[docs]def make_volume_from_radius_compiled(dim: int) -> Callable[[TNumArr], TNumArr]:
"""Return a function calculating the volume of a sphere with a given radius

Args:
dim (int): Dimension of the space

Returns:
function: A function that takes a radius and returns the volume
"""
if dim == 1:

return 2 * radius  # type: ignore

elif dim == 2:

return π * radius ** 2  # type: ignore

elif dim == 3:

return 4 * π / 3 * radius ** 3  # type: ignore

else:
raise NotImplementedError(f"Cannot calculate the volume in {dim} dimensions")
return jit(volume_from_radius)  # type: ignore

[docs]def surface_from_radius(radius: TNumArr, dim: int) -> TNumArr:
"""Return the surface area of a sphere with a given radius

Args:
radius (float or :class:~numpy.ndarray): Radius of the sphere
dim (int): Dimension of the space

Returns:
float or :class:~numpy.ndarray: Surface area of the sphere
"""
if dim == 1:
else:
return 2
elif dim == 2:
return 2 * π * radius  # type: ignore
elif dim == 3:
return 4 * π * radius ** 2  # type: ignore
else:
raise NotImplementedError(
f"Cannot calculate the surface area in {dim} dimensions"
)

[docs]def radius_from_surface(surface: TNumArr, dim: int) -> TNumArr:
"""Return the radius of a sphere with a given surface area

Args:
surface (float or :class:~numpy.ndarray): Surface area of the sphere
dim (int): Dimension of the space

Returns:
float or :class:~numpy.ndarray: Radius of the sphere
"""
if dim == 1:
raise RuntimeError("Cannot calculate radius of 1-d sphere from surface")
elif dim == 2:
return surface / (2 * π)  # type: ignore
elif dim == 3:
return np.sqrt(surface / (4 * π))  # type: ignore
else:
raise NotImplementedError(f"Cannot calculate the radius in {dim} dimensions")

[docs]def make_surface_from_radius_compiled(dim: int) -> Callable[[TNumArr], TNumArr]:
"""Return a function calculating the surface area of a sphere

Args:
dim (int): Dimension of the space

Returns:
function: A function that takes a radius and returns the surface area
"""
import numba as nb

if dim == 1:

if nb.config.DISABLE_JIT:
# jitting is disabled => return generic python function
else:
return 2

else:
# jitting is enabled => return specific compiled functions
@nb.generated_jit(nopython=True)
return lambda radius: 2  # type: ignore
else:
return lambda radius: np.full(radius.shape, 2)  # type: ignore

elif dim == 2:

@jit
return 2 * π * radius  # type: ignore

elif dim == 3:

@jit
return 4 * π * radius ** 2  # type: ignore

else:
raise NotImplementedError(
f"Cannot calculate the surface area in {dim} dimensions"
)

[docs]def points_cartesian_to_spherical(points: np.ndarray) -> np.ndarray:
"""Convert points from Cartesian to spherical coordinates

Args:
points (:class:~numpy.ndarray): Points in Cartesian coordinates

Returns:
:class:~numpy.ndarray: Points (r, θ, φ) in spherical coordinates
"""
points = np.atleast_1d(points)
assert points.shape[-1] == 3, "Points must have 3 coordinates"

ps_spherical = np.empty(points.shape)
# calculate radius in [0, infinity]
ps_spherical[..., 0] = np.linalg.norm(points, axis=-1)
# calculate θ in [0, pi]
ps_spherical[..., 1] = np.arccos(points[..., 2] / ps_spherical[..., 0])
# calculate φ in [0, 2 * pi]
ps_spherical[..., 2] = np.arctan2(points[..., 1], points[..., 0]) % (2 * π)
return ps_spherical

[docs]def points_spherical_to_cartesian(points: np.ndarray) -> np.ndarray:
"""Convert points from spherical to Cartesian coordinates

Args:
points (:class:~numpy.ndarray):
Points in spherical coordinates (r, θ, φ)

Returns:
:class:~numpy.ndarray: Points in Cartesian coordinates
"""
points = np.atleast_1d(points)
assert points.shape[-1] == 3, "Points must have 3 coordinates"

sin_θ = np.sin(points[..., 1])
ps_cartesian = np.empty(points.shape)
ps_cartesian[..., 0] = points[..., 0] * np.cos(points[..., 2]) * sin_θ
ps_cartesian[..., 1] = points[..., 0] * np.sin(points[..., 2]) * sin_θ
ps_cartesian[..., 2] = points[..., 0] * np.cos(points[..., 1])
return ps_cartesian

[docs]def haversine_distance(point1: np.ndarray, point2: np.ndarray) -> np.ndarray:
"""Calculate the haversine-based distance between two points on the surface
of a sphere. Should be more accurate than the arc cosine strategy.
See, for example: https://en.wikipedia.org/wiki/Haversine_formula

Licensed under MIT License (see copy in root of this project)

Args:
point1 (:class:~numpy.ndarray):
First point(s) on the sphere (given in Cartesian coordinates)
point2 (:class:~numpy.ndarray): Second point on the sphere
Second point(s) on the sphere (given in Cartesian coordinates)

Returns:
:class:~numpy.ndarray: The distances between the points
"""
# note that latitude φ is θ and longitude λ is φ in our notation
coords = points_cartesian_to_spherical(point1)
r1, φ1, λ1 = coords[..., 0], coords[..., 1], coords[..., 2]
coords = points_cartesian_to_spherical(point2)
r2, φ2, λ2 = coords[..., 0], coords[..., 1], coords[..., 2]

# check whether both points lie on the same sphere
assert np.allclose(r1, r2)

# we rewrite the standard Haversine slightly as long/lat is not the same as
# spherical coordinates - φ differs by π/4
factor = (1 - np.cos(λ2 - λ1)) / 2
arg = (1 - np.cos(φ2 - φ1)) / 2 + np.sin(φ1) * np.sin(φ2) * factor
return 2 * r1 * np.arcsin(np.sqrt(arg))  # type: ignore

[docs]def get_spherical_polygon_area(vertices: np.ndarray, radius: float = 1) -> float:
"""Calculate the surface area of a polygon on the surface of a sphere.
Based on equation provided here:
http://mathworld.wolfram.com/LHuiliersTheorem.html
Decompose into triangles, calculate excess for each

Licensed under MIT License (see copy in root of this project)

Args:
vertices (:class:~numpy.ndarray): List of vertices (using Cartesian
coordinates) that describe the corners of the polygon. The vertices
need to be oriented.
"""
# have to convert to unit sphere before applying the formula
spherical_coordinates = points_cartesian_to_spherical(vertices)
spherical_coordinates[..., 0] = 1.0
vertices = points_spherical_to_cartesian(spherical_coordinates)

n = vertices.shape
# point we start from
root_point = vertices
totalexcess = 0

# loop from 1 to n-2, with point 2 to n-1 as other vertex of triangle
# this could definitely be written more nicely
b_point = vertices
root_b_dist = haversine_distance(root_point, b_point)
for i in np.arange(1, n - 1):
a_point = b_point
b_point = vertices[i + 1]
root_a_dist = root_b_dist
root_b_dist = haversine_distance(root_point, b_point)
a_b_dist = haversine_distance(a_point, b_point)
s = (root_a_dist + root_b_dist + a_b_dist) / 2.0
arg = (
np.tan(0.5 * s)
* np.tan(0.5 * (s - root_a_dist))
* np.tan(0.5 * (s - root_b_dist))
* np.tan(0.5 * (s - a_b_dist))
)
totalexcess += 4 * np.arctan(np.sqrt(arg))

[docs]class PointsOnSphere:
"""class representing points on an n-dimensional unit sphere"""

def __init__(self, points):
"""
Args:
points (:class:~numpy.ndarray):
The list of points on the unit sphere
"""
self.points = np.asarray(points, dtype=np.double)
# normalize vectors to force them onto the unit-sphere
self.points /= np.linalg.norm(self.points, axis=1)[:, np.newaxis]
self.dim = self.points.shape[-1]

[docs]    @classmethod
def make_uniform(cls, dim: int, num_points: Optional[int] = None):
"""create uniformly distributed points on a sphere

Args:
dim (int): The dimension of space
num_points (int, optional): The number of points to generate. Note
that for one-dimensional spheres (intervals), only exactly two
points can be generated
"""
if dim == 1:
# just have two directions in 2d
if num_points is None:
num_points = 2
if num_points != 2:
raise ValueError("Can only place 2 points in 1d")
points = [[-1], ]

elif dim == 2:
if num_points is None:
num_points = 8
# distribute points evenly around the circle
φs = np.linspace(0, 2 * π, num_points, endpoint=False)
points = np.c_[np.cos(φs), np.sin(φs)]

elif dim == 3:
# Distribute points on the unit sphere using a sunflower spiral
# (inspired by https://stackoverflow.com/a/44164075/932593)
if num_points is None:
num_points = 18
indices = np.arange(0, num_points) + 0.5
φ = np.arccos(1 - 2 * indices / num_points)
θ = π * (1 + 5 ** 0.5) * indices

# convert to Cartesian coordinates
points = np.c_[np.cos(θ) * np.sin(φ), np.sin(θ) * np.sin(φ), np.cos(φ)]

elif num_points is None:
# use vertices of hypercube in n dimensions
points = [
p  # type: ignore
for p in itertools.product([-1, 0, 1], repeat=dim)
if any(c != 0 for c in p)
]

else:
raise NotImplementedError()

# normalize vectors
return cls(points)

[docs]    @cached_method()
def get_area_weights(self, balance_axes: bool = True):
"""return the weight of each point associated with the unit cell size

Args:
balance_axes (bool): Flag determining whether the weights should be
chosen such that the weighted average of all points is the
zero vector

Returns:
:class:~numpy.ndarray: The weight associated with each point
"""
from scipy import spatial

points_flat = self.points.reshape(-1, self.dim)
if self.dim == 1:
assert points_flat.shape == (2, 1)
weights = np.array([0.5, 0.5])

elif self.dim == 2:
# get angles
φ = np.arctan2(points_flat[:, 1], points_flat[:, 0])
idx = np.argsort(φ)
s0 = φ[idx] + 2 * π - φ[idx[-1]]
sizes = np.r_[s0, np.diff(φ[idx]), s0]
weights = (sizes[1:] + sizes[:-1]) / 2
weights /= 2 * π

elif self.dim == 3:
# calculate weights using spherical voronoi construction
voronoi = spatial.SphericalVoronoi(points_flat)
voronoi.sort_vertices_of_regions()

weight_vals = [
get_spherical_polygon_area(voronoi.vertices[ix])
for ix in voronoi.regions
]
weights = np.array(weight_vals, dtype=np.double)
weights /= surface_from_radius(1, dim=self.dim)

else:
raise NotImplementedError()

if balance_axes:
weights /= weights.sum()  # normalize weights
# adjust weights such that all distances are weighted equally, i.e.,
# the weighted sum of all shell vectors should vanish. Additionally,
# the sum of all weights needs to be one. To satisfy these
# constraints simultaneously, the weights are adjusted minimally
# (in a least square sense).
matrix = np.c_[points_flat, np.ones(len(points_flat))]
vector = -weights @ matrix + np.r_[np.zeros(self.dim), 1]
weights += np.linalg.lstsq(matrix.T, vector, rcond=None)

return weights.reshape(self.points.shape[:-1])

[docs]    def get_distance_matrix(self):
"""calculate the (spherical) distances between each point

Returns:
:class:~numpy.ndarray: the distance of each point to each other
"""
from scipy import spatial

if self.dim == 1:
raise ValueError("Distances can only be calculated for dim >= 2")

elif self.dim == 2:
# use arc length on unit circle to calculate distances
def metric(a, b):
return np.arccos(a @ b)

elif self.dim == 3:
# calculate distances on sphere using haversine definition
metric = haversine_distance

else:
raise NotImplementedError()

# determine the distances between all points
dists = spatial.distance.pdist(self.points, metric)
return spatial.distance.squareform(dists)

[docs]    def get_mean_separation(self) -> float:
"""float: calculates the mean distance to the nearest neighbor"""
if len(self.points) < 1:
return float("nan")

dists_sorted = np.sort(self.get_distance_matrix(), axis=1)
return float(dists_sorted[:, 1].mean())

[docs]    def write_to_xyz(self, path: str, comment: str = "", symbol: str = "S"):
"""write the point coordinates to a xyz file

Args:
path (str): location of the file where data is written
comment (str, optional): comment that is written to the second line
symbol (str, optional): denotes the symbol used for the atoms
"""
with open(path, "w") as fp:
fp.write("%d\n" % len(self.points))
fp.write(comment + "\n")
for point in self.points:
point_str = " ".join(("%.12g" % v for v in point))
line = "%s %s\n" % (symbol, point_str)
fp.write(line)

[docs]def spherical_index_k(degree: int, order: int = 0) -> int:
"""returns the mode k from the degree degree and order order

Args:
degree (int): Degree of the spherical harmonics
order (int): Order of the spherical harmonics

Raises:
ValueError: if order < -degree or order > degree

Returns:
int: a combined index k
"""
if not -degree <= order <= degree:
raise ValueError("order must lie between -degree and degree")
return degree * (degree + 1) + order

[docs]def spherical_index_lm(k: int) -> Tuple[int, int]:
"""returns the degree l and the order m from the mode k

Args:
k (int): The combined index for the spherical harmonics

Returns:
tuple: The degree l and order m of the spherical harmonics
assoicated with the combined index
"""
degree = int(np.floor(np.sqrt(k)))
return degree, k - degree * (degree + 1)

[docs]def spherical_index_count(l: int) -> int:
"""return the number of modes for all indices <= l

The returned value is one less than the maximal mode k required.

Args:
l (int): Maximal degree of the spherical harmonics

Returns:
int: The number of modes
"""
return 1 + 2 * l + l * l

[docs]def spherical_index_count_optimal(k_count: int) -> bool:
"""checks whether the modes captures all orders for maximal degree

Args:
k_count (int): The number of modes considered
"""
is_square = bool(int(np.sqrt(k_count) + 0.5) ** 2 == k_count)
return is_square

[docs]def spherical_harmonic_symmetric(degree: int, θ: float) -> float:
r"""axisymmetric spherical harmonics with degree degree, so m=0.

Args:
degree (int): Degree of the spherical harmonics
θ (float): Azimuthal angle at which the spherical harmonics is
evaluated (in :math:[0, \pi])

Returns:
float: The value of the spherical harmonics
"""
# note that the definition of sph_harm has a different convention for the
# usage of the variables φ and θ and we thus have to swap the args
return np.real(sph_harm(0.0, degree, 0.0, θ))  # type: ignore

[docs]def spherical_harmonic_real(degree: int, order: int, θ: float, φ: float) -> float:
r"""real spherical harmonics of degree l and order m

Args:
degree (int): Degree :math:l of the spherical harmonics
order (int): Order :math:m of the spherical harmonics
θ (float): Azimuthal angle (in :math:[0, \pi]) at which the
spherical harmonics is evaluated.
φ (float): Polar angle (in :math:[0, 2\pi]) at which the spherical
harmonics is evaluated.

Returns:
float: The value of the spherical harmonics
"""
# note that the definition of sph_harm has a different convention for the
# usage of the variables φ and θ and we thus have to swap the args
# Moreover, the scipy functions expect first the order and then the degree
if order > 0:
term1 = sph_harm(order, degree, φ, θ)
term2 = (-1) ** order * sph_harm(-order, degree, φ, θ)
return np.real((term1 + term2) / np.sqrt(2))  # type: ignore

elif order == 0:
return np.real(sph_harm(0, degree, φ, θ))  # type: ignore

else:  # order < 0
term1 = sph_harm(-order, degree, φ, θ)
term2 = (-1) ** order * sph_harm(order, degree, φ, θ)
return np.real((term1 - term2) / (complex(0, np.sqrt(2))))  # type: ignore

[docs]def spherical_harmonic_real_k(k: int, θ: float, φ: float) -> float:
r"""real spherical harmonics described by mode k

Args:
k (int): Combined index determining the degree and order of the
spherical harmonics
θ (float): Azimuthal angle (in :math:[0, \pi]) at which the
spherical harmonics is evaluated.
φ (float): Polar angle (in :math:[0, 2\pi]) at which the spherical
harmonics is evaluated.

Returns:
float: The value of the spherical harmonics
"""
return spherical_harmonic_real(*spherical_index_lm(k), θ=θ, φ=φ)