Merge pull request sxs-collaboration#6412 from nikwit/ylm-phys-spec

Add pybindings for Spherepack

src/NumericalAlgorithms/SphericalHarmonics/Python/Bindings.cpp

@@ -3,6 +3,7 @@
 
 #include <pybind11/pybind11.h>
 
+#include "NumericalAlgorithms/SphericalHarmonics/Python/Spherepack.hpp"
 #include "NumericalAlgorithms/SphericalHarmonics/Python/Strahlkorper.hpp"
 #include "NumericalAlgorithms/SphericalHarmonics/Python/StrahlkorperFunctions.hpp"
 #include "Utilities/ErrorHandling/SegfaultHandler.hpp"
@@ -15,4 +16,5 @@ PYBIND11_MODULE(_Pybindings, m) {  // NOLINT
   py::module_::import("spectre.DataStructures.Tensor");
   ylm::py_bindings::bind_strahlkorper(m);
   ylm::py_bindings::bind_strahlkorper_functions(m);
+  ylm::py_bindings::bind_spherepack(m);
 }

src/NumericalAlgorithms/SphericalHarmonics/Python/CMakeLists.txt

@@ -8,6 +8,7 @@ spectre_python_add_module(
   LIBRARY_NAME ${LIBRARY}
   SOURCES
   Bindings.cpp
+  Spherepack.cpp
   Strahlkorper.cpp
   StrahlkorperFunctions.cpp
   PYTHON_FILES
@@ -18,6 +19,7 @@ spectre_python_headers(
   ${LIBRARY}
   INCLUDE_DIRECTORY ${CMAKE_SOURCE_DIR}/src
   HEADERS
+  Spherepack.hpp
   Strahlkorper.hpp
   StrahlkorperFunctions.hpp
   )

src/NumericalAlgorithms/SphericalHarmonics/Python/Spherepack.cpp

@@ -0,0 +1,65 @@
+// Distributed under the MIT License.
+// See LICENSE.txt for details.
+
+#include "NumericalAlgorithms/SphericalHarmonics/Python/Spherepack.hpp"
+
+#include <array>
+#include <cstddef>
+#include <pybind11/operators.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include "DataStructures/DataVector.hpp"
+#include "DataStructures/ModalVector.hpp"
+#include "DataStructures/Tensor/Tensor.hpp"
+#include "NumericalAlgorithms/SphericalHarmonics/Spherepack.hpp"
+#include "NumericalAlgorithms/SphericalHarmonics/SpherepackIterator.hpp"
+
+namespace py = pybind11;
+
+namespace ylm::py_bindings {
+void bind_spherepack(pybind11::module& m) {  // NOLINT
+  py::class_<ylm::Spherepack>(m, "Spherepack")
+      .def(py::init<size_t, size_t>(), py::arg("l_max"), py::arg("m_max"))
+      .def_property_readonly("l_max", &ylm::Spherepack::l_max)
+      .def_property_readonly("m_max", &ylm::Spherepack::m_max)
+      .def_property_readonly("physical_size",
+                             [](const ylm::Spherepack& spherepack) {
+                               return spherepack.physical_size();
+                             })
+      .def_property_readonly("spectral_size",
+                             [](const ylm::Spherepack& spherepack) {
+                               return spherepack.spectral_size();
+                             })
+      .def_property_readonly("theta_phi_points",
+                             &ylm::Spherepack::theta_phi_points)
+      .def("phys_to_spec",
+           static_cast<DataVector (ylm::Spherepack::*)(
+               const DataVector& collocation_values,
+               const size_t physical_stride, const size_t physical_offset)
+                           const>(&ylm::Spherepack::phys_to_spec),
+           py::arg("collocation_values"), py::arg("physical_stride") = 1,
+           py::arg("physical_offset") = 0)
+      .def("spec_to_phys",
+           static_cast<DataVector (ylm::Spherepack::*)(
+               const DataVector& spectral_coefs, const size_t spectral_stride,
+               const size_t spectral_offset) const>(
+               &ylm::Spherepack::spec_to_phys),
+           py::arg("spectral_coefs"), py::arg("spectral_stride") = 1,
+           py::arg("spectral_offset") = 0)
+      // NOLINTNEXTLINE(misc-redundant-expression)
+      .def(py::self == py::self)
+      // NOLINTNEXTLINE(misc-redundant-expression)
+      .def(py::self != py::self);
+
+  py::class_<ylm::SpherepackIterator>(m, "SpherepackIterator")
+      .def(py::init<size_t, size_t, size_t>(), py::arg("l_max"),
+           py::arg("m_max"), py::arg("physical_stride") = 1)
+      .def("__call__", &SpherepackIterator::operator())
+      .def(
+          "set",
+          [](ylm::SpherepackIterator& spherepack, const size_t l_input,
+             const int m_input) { spherepack.set(l_input, m_input); },
+          py::arg("l"), py::arg("m"));
+}
+}  // namespace ylm::py_bindings

src/NumericalAlgorithms/SphericalHarmonics/Python/Spherepack.hpp

@@ -0,0 +1,11 @@
+// Distributed under the MIT License.
+// See LICENSE.txt for details.
+
+#pragma once
+
+#include <pybind11/pybind11.h>
+
+namespace ylm::py_bindings {
+// NOLINTNEXTLINE(google-runtime-references)
+void bind_spherepack(pybind11::module& m);
+}  // namespace ylm::py_bindings

tests/Unit/NumericalAlgorithms/SphericalHarmonics/Python/CMakeLists.txt

@@ -2,7 +2,13 @@
 # See LICENSE.txt for details.
 
 spectre_add_python_bindings_test(
-  "Unit.Ylm.Python"
+  "Unit.Ylm.Python.Strahlkorper"
   Test_Strahlkorper.py
   "Unit;Python"
   PySphericalHarmonics)
+
+  spectre_add_python_bindings_test(
+  "Unit.Ylm.Python.Spherepack"
+  Test_Spherepack.py
+  "Unit;Python"
+  PySphericalHarmonics)

tests/Unit/NumericalAlgorithms/SphericalHarmonics/Python/Test_Spherepack.py

@@ -0,0 +1,89 @@
+# Distributed under the MIT License.
+# See LICENSE.txt for details.
+
+
+import random
+import unittest
+
+import numpy as np
+import numpy.testing as npt
+from scipy.special import sph_harm
+
+from spectre.SphericalHarmonics import Spherepack, SpherepackIterator
+
+
+def convert_coef_to_spherepack(complex_coef, l, m):
+    """
+    converts spherical harmonic coefficients to spherepack coefficients,
+    see https://spectre-code.org/classylm_1_1Spherepack.html
+    """
+    sign = 1.0 if m % 2 == 0 else -1.0
+    part = np.real(complex_coef) if m >= 0 else np.imag(complex_coef)
+    return sign * np.sqrt(2.0 / np.pi) * part
+
+
+class TestStrahlkorper(unittest.TestCase):
+    def test_sizes(self):
+        for l in range(2, 6):
+            for m in range(2, l + 1):
+                spherepack = Spherepack(l, m)
+                self.assertEqual(
+                    spherepack.physical_size, (l + 1) * (2 * m + 1)
+                )
+                self.assertEqual(
+                    spherepack.spectral_size, 2 * (l + 1) * (m + 1)
+                )
+
+    def test_spec_to_phys_and_back(self):
+        l_max = random.randint(2, 12)
+        m_max = random.randint(2, l_max)
+        spherepack = Spherepack(l_max, m_max)
+        thetas, phis = spherepack.theta_phi_points
+        iterator = SpherepackIterator(l_max, m_max)
+        spherepack_coefs = np.zeros(spherepack.spectral_size)
+        scipy_collocation_values = np.zeros(
+            np.shape(thetas), dtype=np.complex128
+        )
+
+        for l in range(0, l_max + 1):
+            for m in range(0, min(l + 1, m_max + 1)):
+                coef = random.uniform(-1.0, 1.0)
+                if m > 0:
+                    coef += random.uniform(-1.0, 1.0) * 1j
+                    iterator.set(l, -m)
+                    spherepack_coefs[iterator()] = convert_coef_to_spherepack(
+                        coef, l, -m
+                    )
+                    # we are checking for a real field, so for the
+                    # conjugate: a*(l,m) = (-1)^m a(l,m)
+                    sign = 1.0 if m % 2 == 0 else -1.0
+                    scipy_collocation_values += (
+                        sign * coef.conjugate() * sph_harm(-m, l, phis, thetas)
+                    )
+
+                iterator.set(l, m)
+                spherepack_coefs[iterator()] = convert_coef_to_spherepack(
+                    coef, l, m
+                )
+                scipy_collocation_values += coef * sph_harm(m, l, phis, thetas)
+
+        npt.assert_array_almost_equal(
+            np.zeros_like(thetas), np.imag(scipy_collocation_values)
+        )
+        collocation_values = np.asarray(
+            spherepack.spec_to_phys(spherepack_coefs)
+        )
+        npt.assert_array_almost_equal(
+            collocation_values, np.real(scipy_collocation_values)
+        )
+
+        recovered_coefs = np.asarray(
+            spherepack.phys_to_spec(collocation_values)
+        )
+        npt.assert_array_almost_equal(
+            recovered_coefs, np.asarray(spherepack_coefs)
+        )
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)