Updates for JAX 0.7

.github/workflows/ci.yml

@@ -16,13 +16,15 @@ jobs:
 
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
       - run: "python -m pip install black"
 
       - name: lint
         run: "python -m black --line-length 100 --check ./src/tfc"
 
-      - run: "sudo apt-get update && sudo apt-get install -y swig gcc g++"
-      - run: python -m pip install wheel setuptools numpy pytest
+      - run: "sudo apt-get update && sudo apt-get install -y gcc g++"
+      - run: python -m pip install wheel setuptools numpy pytest pybind11 mypy
       - run: python setup.py bdist_wheel
       - run: pip install ./dist/*.whl
 

.github/workflows/publish_wheels.yml

@@ -12,13 +12,15 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Install libraries
-        run: "sudo apt-get update && sudo apt-get install -y swig gcc g++"
+        run: "sudo apt-get update && sudo apt-get install -y gcc g++"
 
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
 
       - name: Checkout dependencies
-        run: python -m pip install wheel setuptools numpy
+        run: python -m pip install wheel setuptools numpy pybind11 mypy
 
       - name: Create source distribution
         run: python setup.py sdist
@@ -48,6 +50,8 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
 
       - name: Install cibuildwheel
         run: python -m pip install cibuildwheel

.github/workflows/publish_wheels_test_pypi.yml

@@ -9,13 +9,15 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Install libraries
-        run: "sudo apt-get update && sudo apt-get install -y swig gcc g++"
+        run: "sudo apt-get update && sudo apt-get install -y gcc g++"
 
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
 
       - name: Checkout dependencies
-        run: python -m pip install wheel setuptools numpy
+        run: python -m pip install wheel setuptools numpy pybind11 mypy
 
       - name: Create source distribution
         run: python setup.py sdist
@@ -45,6 +47,8 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
 
       - name: Install cibuildwheel
         run: python -m pip install cibuildwheel

.github/workflows/run_most_examples.yml

@@ -13,9 +13,11 @@ jobs:
 
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
 
-      - run: "sudo apt-get update && sudo apt-get install -y swig gcc g++ graphviz"
-      - run: python -m pip install wheel setuptools numpy pytest
+      - run: "sudo apt-get update && sudo apt-get install -y gcc g++ graphviz"
+      - run: python -m pip install wheel setuptools numpy pytest pybind11 mypy
       - run: python setup.py bdist_wheel
       - run: pip install ./dist/*.whl
       - run: pip install -r ./requirements.txt

.readthedocs.yml

@@ -3,11 +3,13 @@ version: 2
 
 # Set the version of Python and other tools you might need
 build:
-  os: ubuntu-22.04
+  os: ubuntu-24.04
   tools:
-    python: "3.10"
+    python: "3.12"
   apt_packages:
     - graphviz
+    - cmake
+    - ninja-build
 
 # Build documentation in the docs/ directory with Sphinx
 sphinx:
@@ -20,3 +22,5 @@ formats:  all
 python:
   install:
     - requirements: docs/requirements.txt
+    - method: pip
+      path: .

MANIFEST.in

@@ -1,3 +1,5 @@
-include src/tfc/utils/BF/BF.h
-include src/tfc/utils/BF/numpy.i
+include src/tfc/utils/BF.h
+include src/tfc/utils/BF.cc
+include src/tfc/utils/BF_Py.cc
+include src/tfc/utils/CMakeLists.txt
 include src/tfc/py.typed

clang-format.yaml

@@ -0,0 +1,14 @@
+BasedOnStyle: LLVM
+IndentWidth: 4
+ColumnLimit: 120
+AllowShortIfStatementsOnASingleLine: false
+
+# Make method/function arguments line up if they are broken onto a new line
+BinPackArguments: false
+BinPackParameters: false
+
+# If we break up constructor initializers, do so before the comma.
+BreakConstructorInitializers: BeforeComma
+
+# Indent all namespaces
+NamespaceIndentation: All

docs/Doxyfile

@@ -794,7 +794,7 @@ WARN_LOGFILE           =
 # Note: If this tag is empty the current directory is searched.
 
 #INPUT                  = "../../../src/cxx" "../../../src/tfc" "../../../src/tfc/utils"
-INPUT                  = "../../../src/tfc/utils/BF"
+INPUT                  = "../../../src/tfc/utils"
 
 # This tag can be used to specify the character encoding of the source files
 # that doxygen parses. Internally doxygen uses the UTF-8 encoding. Doxygen uses

docs/requirements.txt

@@ -9,6 +9,8 @@ breathe
 exhale
 nbsphinx
 ipykernel
+pybind11
+mypy
 jax
 jaxlib
 jaxtyping

pyproject.toml

@@ -2,40 +2,38 @@
 requires = ["setuptools>=42", 
             "wheel", 
             "numpy>=2.1",
+            "pybind11~=3.0.0",
 ]
 build-backend = "setuptools.build_meta"
 
 [project]
 name = "tfc"
 version = "1.2.1"
-requires-python = ">=3.10"
+requires-python = ">=3.11,<3.14"
 readme = "README.md"
 dynamic = ["dependencies", "classifiers", "authors", "license", "description"]
 
 [tool.black]
 line-length = 100
-target-version = ['py310','py311']
+target-version = ['py311']
 
 [tool.cibuildwheel]
-before-build = "pip install setuptools wheel numpy"
-skip = "pp* *-musllinux*"
+before-build = "pip install setuptools wheel numpy pybind11 mypy"
+skip = "*-musllinux*"
 manylinux-x86_64-image = "manylinux2014"
 test-requires = ["pytest"]
 test-command = "pytest {package}/tests"
 test-skip = "*-macosx_arm64"
 
 [tool.cibuildwheel.linux]
-before-all = "yum install -y swig"
+before-all = "yum install -y ninja-build"
 archs = ["x86_64"]
 
 [tool.cibuildwheel.macos]
-before-all = "brew install swig"
 archs = ["arm64"]
 
 [tool.cibuildwheel.windows]
 archs = ["AMD64"]
 
 [[tool.cibuildwheel.overrides]]
 select = "*-musllinux*"
-before-all = "apk add swig"
-

requirements.txt

@@ -1,7 +1,7 @@
 numpy >= 1.25
 scipy >= 1.11
-jax ~= 0.6
-jaxlib ~= 0.6.0
+jax >= 0.6
+jaxlib >= 0.6
 jaxtyping
 annotated-types
 matplotlib
@@ -16,6 +16,7 @@ tqdm
 pandas
 openpyxl
 sympy
+pybind11
 
 # Optional
 # pdfCropMargins # Used to crop pdfs in PlotlyMakePlot

setup.py

@@ -1,33 +1,28 @@
 import sys
-from os import path, name
-import numpy
+import os
+from pathlib import Path
 from setuptools import setup, Extension, find_packages
-from setuptools.command.build_py import build_py as _build_py
+from setuptools.command.build_ext import build_ext
+from subprocess import check_call
 
 # Get long description
-this_directory = path.abspath(path.dirname(__file__))
-with open(path.join(this_directory, "README.md"), encoding="utf-8") as f:
+this_directory = os.path.abspath(os.path.dirname(__file__))
+with open(os.path.join(this_directory, "README.md"), encoding="utf-8") as f:
     long_description = f.read()
 long_description = long_description.replace(
     '<img src="https://github.com/leakec/tfc/blob/main/docs/Univariate_TFC_Animation.gif" width="600" height="467">',
     "",
     1,
 )
 
-# Get numpy directory
-try:
-    numpy_include = numpy.get_include()
-except AttributeError:
-    numpy_include = numpy.get_numpy_include()
-
 # Get version info
 version_dict = {}
 with open("src/tfc/version.py") as f:
     exec(f.read(), version_dict)
     version = version_dict["__version__"]
 
 # In the future, can add -DHAS_CUDA to this to enable GPU support
-if name == "nt":
+if os.name == "nt":
     # Windows compile flags
     cxxFlags = ["/O2", "/std:c++17", "/Wall", "/DWINDOWS_MSVC"]
 else:
@@ -38,22 +33,45 @@
 else:
     numpy_version = "numpy>=1.21.0"
 
-# Create basis function c++ extension
-BF = Extension(
-    "tfc.utils.BF._BF",
-    sources=["src/tfc/utils/BF/BF.i", "src/tfc/utils/BF/BF.cxx"],
-    include_dirs=["src/tfc/utils/BF", numpy_include],
-    swig_opts=["-c++", "-doxygen", "-O", "-olddefs"],
-    extra_compile_args=cxxFlags,
-    extra_link_args=cxxFlags,
-)
+class CMakeExtension(Extension):
+    def __init__(self, name, sourcedir=""):
+        super().__init__(name, sources=[])
+        self.sourcedir = str((Path(sourcedir) / "src" / "tfc" / "utils").absolute())
+
+
+class CMakeBuild(build_ext):
+    def build_extension(self, ext):
+        extdir = Path(self.get_ext_fullpath(ext.name)).parents[0].absolute()
+        bf_dir = extdir / "tfc" / "utils"
+
+        import pybind11
+        dark = Path(pybind11.__file__).parents[0]
+        pybind11_dir = dark / "share" / "cmake" / "pybind11"
+
+
+        cfg = "Debug" if self.debug else "Release"
+        cmake_args = [
+            f"-DCMAKE_BUILD_TYPE={cfg}",
+            f"-DCMAKE_INSTALL_PREFIX={bf_dir}",
+            f"-Dpybind11_DIR={pybind11_dir}",
+            f"-DPython_EXECUTABLE={sys.executable}"
+        ]
+
+        # Optional: use Ninja if available
+        if "CMAKE_GENERATOR" not in os.environ:
+            cmake_args += ["-G", "Ninja"]
+
+        build_temp = Path(self.build_temp)
+        build_temp.mkdir(parents=True, exist_ok=True)
+
+        # Run CMake configuration
+        check_call(["cmake", ext.sourcedir] + cmake_args, cwd=build_temp)
 
+        # Run CMake build
+        check_call(["cmake", "--build", ".", "--config", cfg], cwd=build_temp)
 
-# Custom build options to include swig Python files
-class build_py(_build_py):
-    def run(self):
-        self.run_command("build_ext")
-        super(build_py, self).run()
+        # Run CMake install
+        check_call(["cmake", "--install", "."], cwd=build_temp)
 
 
 # Setup
@@ -70,13 +88,14 @@ def run(self):
     packages=find_packages("src"),
     package_dir={"": "src"},
     package_data={"": ["src/tfc/py.typed"]},
-    python_requires=">=3.10",
+    python_requires=">=3.11",
     include_package_data=True,
-    ext_modules=[BF],
+    ext_modules=[CMakeExtension("BF")],
+    cmdclass={"build_ext": CMakeBuild},
     install_requires=[
         numpy_version,
-        "jax ~= 0.6.0",
-        "jaxlib ~= 0.6.0",
+        "jax >= 0.6.0",
+        "jaxlib >= 0.6.0",
         "jaxtyping",
         "annotated-types",
         "matplotlib",
@@ -87,14 +106,10 @@ def run(self):
     ],
     classifiers=[
         "Development Status :: 4 - Beta",
-        "License :: OSI Approved :: MIT License",
         "Natural Language :: English",
         "Programming Language :: C++",
         "Programming Language :: Python :: 3 :: Only",
         "Topic :: Scientific/Engineering",
         "Topic :: Education",
     ],
-    cmdclass={
-        "build_py": build_py,
-    },
 )