Merge pull request #17 from vdutor/devel

Updates

.github/workflows/quality-checks.yaml

@@ -7,7 +7,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.7, 3.8, 3.9]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
       fail-fast: false
 
     name: Python-${{ matrix.python-version }}

README.md

@@ -70,26 +70,26 @@ assert out.numpy().shape == (101, num_harmonics)
 
 The fundamental systems up to dimensino 20 are precomputed and stored in `spherical_harmonics/fundamental_system`. For each dimension we precompute the first amount of spherical harmonics. This means that in each dimension we support a varying number of maximum degree (`max_degree`) and number of spherical harmonics:
 
-|   Dimension |   Max Degree |   Number Harmonics |
-|------------:|-------------:|-------------------:|
-|           3 |           34 |               1156 |
-|           4 |           14 |               1015 |
-|           5 |           10 |               1210 |
-|           6 |            8 |               1254 |
-|           7 |            7 |               1386 |
-|           8 |            6 |               1122 |
-|           9 |            6 |               1782 |
-|          10 |            6 |               2717 |
-|          11 |            5 |               1287 |
-|          12 |            5 |               1729 |
-|          13 |            5 |               2275 |
-|          14 |            5 |               2940 |
-|          15 |            5 |               3740 |
-|          16 |            4 |                952 |
-|          17 |            4 |               1122 |
-|          18 |            4 |               1311 |
-|          19 |            4 |               1520 |
-|          20 |            4 |               1750 |
+| Dimension | Max Degree | Number Harmonics |
+|----------:|-----------:|-----------------:|
+|         3 |         34 |             1156 |
+|         4 |         20 |             2870 |
+|         5 |         10 |            16170 |
+|         6 |          8 |             1254 |
+|         7 |          7 |             1386 |
+|         8 |          6 |             1122 |
+|         9 |          6 |             1782 |
+|        10 |          6 |             2717 |
+|        11 |          5 |             1287 |
+|        12 |          5 |             1729 |
+|        13 |          5 |             2275 |
+|        14 |          5 |             2940 |
+|        15 |          5 |             3740 |
+|        16 |          4 |              952 |
+|        17 |          4 |             1122 |
+|        18 |          4 |             1311 |
+|        19 |          4 |             1520 |
+|        20 |          4 |             1750 |
 
 To precompute a larger fundamental system for a dimension run the following script
 ```

pyproject.toml

@@ -1,14 +1,33 @@
+[build-system]
+requires = ["flit_core >=3.2,<4"]
+build-backend = "flit_core.buildapi"
+
+[project]
+name = "SphericalHarmonics"
+version = "0.0.2a2"
+authors = [{ name = "Vincent Dutordoir", email = "vd309@cam.ac.uk" }]
+description = "Python Implementation of Spherical harmonics in dimension >= 3"
+readme = "README.md"
+dependencies = [
+    "numpy",
+    "scipy",
+]
+
+[project.urls]
+Source = "https://github.com/vdutor/SphericalHarmonics"
+
+[tool.flit.module]
+name = "spherical_harmonics"
+
 [tool.mypy]
 ignore_missing_imports = true
 strict_optional = false
 allow_redefinition = true
 
-
 [tool.black]
 line-length = 88
 target-version = ['py37']
 
-
 [tool.isort]
 profile = "black"
 skip_glob = [
@@ -17,4 +36,4 @@ skip_glob = [
 known_third_party = [
      "lab",
      "plum",
-]
+]

src/spherical_harmonics/__init__.py

@@ -14,4 +14,6 @@
 
 # flake8: noqa
 
+"""Zonal Spherical Harmonics in d Dimensions in TensorFlow, PyTorch and Jax"""
+
 from .spherical_harmonics import SphericalHarmonics

src/spherical_harmonics/fundamental_set.py

@@ -13,7 +13,9 @@
 # limitations under the License.
 
 import argparse
+import warnings
 from pathlib import Path
+from typing import Optional
 
 import numpy as np
 from pkg_resources import resource_filename
@@ -48,43 +50,67 @@ class FundamentalSystemCache:
     harmonics in an arbitrary dimension"""
 
     def __init__(
-        self, dimension: int, load_dir="fundamental_system", only_use_cache: bool = True
+        self,
+        dimension: int,
+        load_dir="fundamental_system",
+        only_use_cache: bool = True,
+        strict_loading: bool = True,
     ):
         self.file_name = Path(
             resource_filename(__name__, f"{load_dir}/fs_{dimension}D.npz")
         )
         self.dimension = dimension
         self.only_use_cache = only_use_cache
+        self.strict_loading = strict_loading
 
         if self.file_name.exists():
             with np.load(self.file_name) as data:
                 self.cache = {k: v for (k, v) in data.items()}
-        elif only_use_cache:
+        elif only_use_cache and self.strict_loading:
             raise ValueError(
                 f"Fundamental system for dimension {dimension} has not been precomputed."
                 "Terminating computations. Precompute set by running `fundamental_set.py`"
             )
         else:
             self.cache = {}
 
+    @property
+    def max_computed_degree(self) -> int:
+        max_degree = -1
+        while True:
+            max_degree = max_degree + 1
+            key = self.cache_key(max_degree)
+            if key not in self.cache:
+                break
+        return max_degree
+
     def cache_key(self, degree: int) -> str:
         """Return the key used in the cache"""
         return f"degree_{degree}"
 
-    def load(self, degree: int) -> np.ndarray:
+    def load(self, degree: int) -> Optional[np.ndarray]:
         """Load or calculate the set for given degree"""
         key = self.cache_key(degree)
         if key not in self.cache:
-            if self.only_use_cache:
+            if self.only_use_cache and self.strict_loading:
                 raise ValueError(
                     f"Fundamental system for dimension {self.dimension} and degree "
                     f"{degree} has not been precomputed. Terminating "
                     "computations. Precompute set by running `fundamental_set.py`"
                 )
-            else:
+            elif not self.only_use_cache:
                 print("WARNING: Cache miss - calculating  system")
                 self.cache[key] = self.calculate(self.dimension, degree)
 
+        if key not in self.cache and not self.strict_loading:
+            warnings.warn(
+                f"Fundamental system for dimension {self.dimension} has been "
+                f"computed up to degree {self.max_computed_degree}. This means you will not "
+                "be able to evaluate individual spherical harmonics for larger degrees.",
+                RuntimeWarning,
+            )
+            return None
+
         return self.cache[key]
 
     def regenerate_and_save_cache(self, max_degrees: int) -> None:

src/spherical_harmonics/gegenbauer_polynomial.py

@@ -17,6 +17,7 @@
 import lab as B
 import numpy as np
 from beartype.typing import List, Tuple, Union
+from lab import dispatch
 from scipy.special import gegenbauer as scipy_gegenbauer
 from scipy.special import loggamma
 
@@ -141,10 +142,22 @@ class GegenbauerScipyCoefficients:
     def __init__(self, n: int, alpha: float):
         self.n = n
         self.alpha = alpha
-        C = scipy_gegenbauer(n, alpha)
-        self._at_1 = C(1.0)
-        self.coefficients = list(C.coefficients)
+        self.C = scipy_gegenbauer(n, alpha)
+        self._at_1 = self.C(1.0)
+        self.coefficients = list(self.C.coefficients)
 
+    @dispatch
+    def __call__(self, x: B.NPNumeric) -> B.Numeric:
+        if self.n < 0:
+            return B.zeros(x)
+        elif self.n == 0:
+            return B.ones(x)
+        elif self.n == 1:
+            return 2 * self.alpha * x
+
+        return self.C(x)
+
+    @dispatch  # type: ignore[no-redef]
     def __call__(self, x: B.Numeric) -> B.Numeric:
         """x: [...], return [...]"""
         if self.n < 0:

src/spherical_harmonics/spherical_harmonics.py

@@ -37,6 +37,7 @@ def __init__(
         dimension: int,
         degrees: Union[int, List[int]],
         debug: bool = False,
+        allow_uncomputed_levels: bool = False,
     ):
         """
         :param dimension: if d = dimension, then
@@ -46,6 +47,8 @@ def __init__(
             if integer all levels (or degrees) up to `degrees` are used.
         highest degree of polynomial
             in the collection (exclusive)
+        :param allow_uncomputed_levels: if True, allow levels without the precomputed
+            fundamental system.
         :param debug: print debug messages.
         """
         assert (
@@ -56,7 +59,9 @@ def __init__(
         if isinstance(degrees, int):
             degrees = list(range(degrees))
 
-        self.fundamental_system = FundamentalSystemCache(dimension)
+        self.fundamental_system = FundamentalSystemCache(
+            dimension, strict_loading=not allow_uncomputed_levels
+        )
         self.harmonic_levels = [
             SphericalHarmonicsLevel(dimension, degree, self.fundamental_system)
             for degree in degrees
@@ -134,20 +139,36 @@ def __init__(self, dimension: int, degree: int, fundamental_system=None):
         self.surface_area_sphere = surface_area_sphere(dimension)
         # normalising constant
         c = self.alpha / (degree + self.alpha)
-        VtV = np.dot(self.V, self.V.T)
-        self.A = c * scipy_gegenbauer(self.degree, self.alpha)(VtV)
-        self.L = np.linalg.cholesky(self.A)  # [M, M]
-        # Cholesky inverse corresponds to the weights you get from Gram-Schmidt
-        self.L_inv = np.linalg.solve(self.L, np.eye(len(self.L)))
+
+        if self.is_level_computed:
+            VtV = np.dot(self.V, self.V.T)
+            self.A = c * scipy_gegenbauer(self.degree, self.alpha)(VtV)
+            self.L = np.linalg.cholesky(self.A)  # [M, M]
+            # Cholesky inverse corresponds to the weights you get from Gram-Schmidt
+            self.L_inv = np.linalg.solve(self.L, np.eye(len(self.L)))
         self.gegenbauer = Gegenbauer(self.degree, self.alpha)
 
+    @property
+    def is_level_computed(self) -> bool:
+        """
+        Whether the level has the fundamental system computed.
+        """
+        return self.V is not None
+
     def __call__(self, X: B.Numeric) -> B.Numeric:
         r"""
         :param X: M normalised (i.e. unit) D-dimensional vector, [N, D]
 
         :return: `X` evaluated at the M spherical harmonics in the set.
             [\phi_m(x_i)], shape [M, N]
         """
+        if not self.is_level_computed:
+            raise ValueError(
+                f"Fundamental system for dimension {self.dimension} and degree "
+                f"{self.degree} has not been precomputed. Terminating "
+                "computations. Precompute set by running `fundamental_set.py`"
+            )
+
         VXT = B.matmul(
             B.cast(B.dtype(X), from_numpy(X, self.V)), X, tr_b=True
         )  # [M, N]