Implementing Numpy and JAX substrates using exoplanet-core

.github/workflows/tests.yml

@@ -18,7 +18,7 @@ jobs:
       fail-fast: false
       matrix:
         python-version: ["3.7", "3.8"]
-        pymc-version: ["pymc3==3.9", "pymc3==3.10"]
+        pymc-version: ["pymc3<3.10", "pymc3<3.11"]
         include:
           - python-version: "3.8"
             pymc-version: "https://github.com/pymc-devs/pymc3/archive/master.zip"
@@ -55,6 +55,8 @@ jobs:
       - name: Run tests
         shell: bash -l {0}
         run: python -m pytest --cov=exoplanet -v tests
+        env:
+          THEANO_FLAGS: "floatX=float64,compute_test_value=off"
 
       - name: Get unique id
         id: unique-id

.gitignore

@@ -23,3 +23,6 @@ venv-test
 *.ipynb
 .envrc
 tmp
+var
+src/exoplanet/pymc/orbits/*
+src/exoplanet/pymc/light_curves/*

.gitmodules

@@ -1,3 +0,0 @@
-[submodule "src/exoplanet/theano_ops/lib/vendor/eigen"]
-	path = src/exoplanet/theano_ops/lib/vendor/eigen
-	url = https://gitlab.com/libeigen/eigen.git

setup.py

@@ -4,15 +4,16 @@
 # https://hynek.me/articles/sharing-your-labor-of-love-pypi-quick-and-dirty/
 
 import codecs
+import glob
 import os
 import re
-import sys
+from shutil import copyfile
 
-from setuptools import Extension, find_packages, setup
-from setuptools.command.build_ext import build_ext
+from setuptools import find_packages, setup
 
 # PROJECT SPECIFIC
 
+HERE = os.path.dirname(os.path.realpath(__file__))
 NAME = "exoplanet"
 PACKAGES = find_packages(where="src")
 META_PATH = os.path.join("src", "exoplanet", "__init__.py")
@@ -27,6 +28,7 @@
 ]
 SETUP_REQUIRES = ["setuptools>=40.6.0", "setuptools_scm"]
 INSTALL_REQUIRES = [
+    "exoplanet-core",
     "pybind11>=2.4",
     "numpy>=1.13.0",
     "pymc3>=3.5",
@@ -97,103 +99,32 @@
 
 # END PROJECT SPECIFIC
 
-# PYBIND11
-
-
-class get_pybind_include:
-    def __init__(self, user=False):
-        self.user = user
-
-    def __str__(self):
-        import pybind11
-
-        return pybind11.get_include(self.user)
-
-
-class get_numpy_include:
-    def __str__(self):
-        import numpy
-
-        return numpy.get_include()
-
-
-class custom_build_ext(build_ext):
-    c_opts = {"msvc": ["/EHsc"], "unix": []}
-    l_opts = {"msvc": [], "unix": []}
-
-    if sys.platform == "darwin":
-        darwin_opts = [
-            "-stdlib=libc++",
-            "-mmacosx-version-min=10.14",
-            "-march=native",
-        ]
-        c_opts["unix"] += darwin_opts
-        l_opts["unix"] += darwin_opts
-
-    def has_flag(self, flagname):
-        import tempfile
-
-        import setuptools
-
-        with tempfile.NamedTemporaryFile("w", suffix=".cpp") as f:
-            f.write("int main (int argc, char **argv) { return 0; }")
-            try:
-                self.compiler.compile([f.name], extra_postargs=[flagname])
-            except setuptools.distutils.errors.CompileError:
-                return False
-        return True
-
-    def cpp_flag(self):
-        flags = ["-std=c++17", "-std=c++14", "-std=c++11"]
-
-        for flag in flags:
-            if self.has_flag(flag):
-                return flag
-
-        raise RuntimeError(
-            "Unsupported compiler. At least C++11 support is needed."
-        )
-
-    def build_extensions(self):
-        ct = self.compiler.compiler_type
-        opts = self.c_opts.get(ct, [])
-        link_opts = self.l_opts.get(ct, [])
-        if ct == "unix":
-            opts.append(
-                '-DVERSION_INFO="%s"' % self.distribution.get_version()
-            )
-            opts.append(self.cpp_flag())
-            if self.has_flag("-fvisibility=hidden"):
-                opts.append("-fvisibility=hidden")
-        elif ct == "msvc":
-            opts.append(
-                '/DVERSION_INFO=\\"%s\\"' % self.distribution.get_version()
-            )
-        for ext in self.extensions:
-            ext.extra_compile_args = opts
-            ext.extra_link_args = link_opts
-        build_ext.build_extensions(self)
-
-
-include_dirs = [
-    "src/exoplanet/theano_ops/lib/include",
-    "src/exoplanet/theano_ops/lib/vendor/eigen",
-    get_numpy_include(),
-    get_pybind_include(),
-    get_pybind_include(user=True),
-]
-ext_modules = [
-    Extension(
-        "exoplanet.theano_ops.driver",
-        ["src/exoplanet/theano_ops/driver.cpp"],
-        include_dirs=include_dirs,
-        language="c++",
-    )
-]
+# COPY SUBSTRATES
 
-# END PYBIND11
 
-HERE = os.path.dirname(os.path.realpath(__file__))
+def setup_substrates(filelist):
+    import logging
+
+    logger = logging.getLogger(__name__)
+    for substrate in ["pymc"]:
+        for pattern in filelist:
+            for source in glob.glob(f"src/exoplanet/numpy/{pattern}"):
+                dest = source.replace(
+                    "src/exoplanet/numpy", f"src/exoplanet/{substrate}"
+                )
+                logger.info(f"Copying {source} -> {dest}")
+                os.makedirs(os.path.dirname(dest), exist_ok=True)
+                copyfile(source, dest)
+
+
+setup_substrates(
+    [
+        "orbits/*.py",
+        "light_curves/*.py",
+    ]
+)
+
+# END SUBSTRATES
 
 
 def read(*parts):
@@ -236,6 +167,4 @@ def find_meta(meta, meta_file=read(META_PATH)):
         classifiers=CLASSIFIERS,
         setup_requires=SETUP_REQUIRES,
         zip_safe=False,
-        ext_modules=ext_modules,
-        cmdclass={"build_ext": custom_build_ext},
     )

src/exoplanet/__init__.py

@@ -1,24 +1,10 @@
 # -*- coding: utf-8 -*-
 
-__all__ = [
-    "__version__",
-    "distributions",
-    "orbits",
-    "interp",
-    "get_dense_nuts_step",
-    "sample",
-    "optimize",
-]
+__all__ = ["__version__"]
 
-from . import distributions, interp, orbits
 from .citations import CITATIONS
-from .distributions import *  # NOQA
 from .estimators import *  # NOQA
 from .exoplanet_version import __version__
-from .light_curves import *  # NOQA
-from .optim import optimize
-from .sampling import get_dense_nuts_step, sample
-from .utils import *  # NOQA
 
 __bibtex__ = __citation__ = CITATIONS["exoplanet"][1]
 __uri__ = "https://docs.exoplanet.codes"

src/exoplanet/citations.py

@@ -1,38 +1,18 @@
 # -*- coding: utf-8 -*-
 
-__all__ = ["add_citations_to_model", "CITATIONS"]
+__all__ = ["format_citations", "CITATIONS"]
 
-import logging
 import textwrap
 
-import pymc3 as pm
 
-
-def add_citations_to_model(citations, model=None):
-    try:
-        model = pm.modelcontext(model)
-        if not hasattr(model, "__citations__"):
-            model.__citations__ = dict()
-        for k in citations:
-            model.__citations__[k] = CITATIONS[k]
-
-    except TypeError:
-        pass
-
-
-def get_citations_for_model(model=None, width=79):
-    """Get the citations for the components used an exoplanet PyMC3
+def format_citations(citations, width=79):
+    """Get the citations for the components used in an exoplanet model
 
     Returns: The acknowledgement text for exoplanet and its dependencies and a
     string containing the BibTeX entries for the citations in the
     acknowledgement.
 
     """
-    model = pm.modelcontext(model)
-    if not hasattr(model, "__citations__"):
-        logging.warning("no citations registered with model")
-        return "", ""
-
     cite = (
         list(CITATIONS["exoplanet"][0])
         + list(CITATIONS["pymc3"][0])
@@ -45,7 +25,7 @@ def get_citations_for_model(model=None, width=79):
         CITATIONS["theano"][1],
         CITATIONS["arviz"][1],
     ]
-    for k, v in model.__citations__.items():
+    for k, v in citations.items():
         cite += list(v[0])
         bib.append(v[1])