From 731b2ab2143913817109f0703000c25e2e623f09 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Fri, 12 Jan 2024 17:12:23 -0500 Subject: [PATCH 01/50] require dev version of twobody --- setup.cfg | 4 ++-- thejoker/_astropy_init.py | 37 ------------------------------------- 2 files changed, 2 insertions(+), 39 deletions(-) diff --git a/setup.cfg b/setup.cfg index 7ae9ca34..3a3020b4 100644 --- a/setup.cfg +++ b/setup.cfg @@ -17,8 +17,8 @@ packages = find: python_requires = >=3.7 install_requires = numpy>=1.20,<1.22 - astropy<6.0 - twobody>=0.8.3 + astropy + twobody @ git+https://github.com/adrn/twobody scipy h5py schwimmbad>=0.3.1 diff --git a/thejoker/_astropy_init.py b/thejoker/_astropy_init.py index 2dffe8fd..058ecc61 100644 --- a/thejoker/_astropy_init.py +++ b/thejoker/_astropy_init.py @@ -13,40 +13,3 @@ from .version import version as __version__ except ImportError: __version__ = '' - - -if not _ASTROPY_SETUP_: # noqa - import os - from warnings import warn - from astropy.config.configuration import ( - update_default_config, - ConfigurationDefaultMissingError, - ConfigurationDefaultMissingWarning) - - # Create the test function for self test - from astropy.tests.runner import TestRunner - test = TestRunner.make_test_runner_in(os.path.dirname(__file__)) - test.__test__ = False - __all__ += ['test'] - - # add these here so we only need to cleanup the namespace at the end - config_dir = None - - if not os.environ.get('ASTROPY_SKIP_CONFIG_UPDATE', False): - config_dir = os.path.dirname(__file__) - config_template = os.path.join(config_dir, __package__ + ".cfg") - if os.path.isfile(config_template): - try: - update_default_config( - __package__, config_dir, version=__version__) - except TypeError as orig_error: - try: - update_default_config(__package__, config_dir) - except ConfigurationDefaultMissingError as e: - wmsg = (e.args[0] + - " Cannot install default profile. If you are " - "importing from source, this is expected.") - warn(ConfigurationDefaultMissingWarning(wmsg)) - del e - except Exception: - raise orig_error From e667453fb451a2ffab9b43f2a42362ba5d3a1d85 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Sun, 14 Jan 2024 10:53:51 -0500 Subject: [PATCH 02/50] fix build --- pyproject.toml | 2 +- setup.cfg | 4 ++-- tox.ini | 9 ++++----- 3 files changed, 7 insertions(+), 8 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ce7ac684..f567329e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ requires = ["setuptools", "wheel", "extension-helpers", "scipy", - "oldest-supported-numpy", + "numpy", "cython", "twobody"] diff --git a/setup.cfg b/setup.cfg index 3a3020b4..abc0c2fb 100644 --- a/setup.cfg +++ b/setup.cfg @@ -16,9 +16,9 @@ zip_safe = False packages = find: python_requires = >=3.7 install_requires = - numpy>=1.20,<1.22 + numpy astropy - twobody @ git+https://github.com/adrn/twobody + twobody>=0.9 scipy h5py schwimmbad>=0.3.1 diff --git a/tox.ini b/tox.ini index 7d87cc33..10094d5a 100644 --- a/tox.ini +++ b/tox.ini @@ -1,6 +1,6 @@ [tox] envlist = - py{38,39,310}-test{,-oldestdeps,-devdeps}{,-cov} + py{39,310,311}-test{,-oldestdeps,-devdeps}{,-cov} build_docs requires = setuptools >= 30.3.0 @@ -13,7 +13,7 @@ indexserver = setenv = MPLBACKEND=agg # Pass through the following environment variables which may be needed for the CI -passenv = HOME WINDIR LC_ALL LC_CTYPE CC CI TRAVIS +passenv = HOME, WINDIR, LC_ALL, LC_CTYPE, CC, CI, TRAVIS # Run the tests in a temporary directory to make sure that we don't import # this package from the source tree @@ -37,12 +37,11 @@ deps = # The oldestdeps factor is intended to be used to install the oldest # versions of all dependencies that have a minimum version. - oldestdeps: numpy==1.19.* + oldestdeps: numpy==1.24.* oldestdeps: matplotlib==3.4.* oldestdeps: scipy==1.6.* - oldestdeps: astropy==4.2.* + oldestdeps: astropy==5.2.* - devdeps: :NIGHTLY:numpy devdeps: git+https://github.com/astropy/astropy.git#egg=astropy # The following indicates which extras_require from setup.cfg will be installed From da06f90202bc644bcba05274fdaf4d3dde7e617f Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Sun, 14 Jan 2024 10:58:16 -0500 Subject: [PATCH 03/50] update versions --- .github/workflows/packaging.yml | 63 +++++++++++++++++---------------- .github/workflows/tests.yml | 29 +++++++-------- 2 files changed, 44 insertions(+), 48 deletions(-) diff --git a/.github/workflows/packaging.yml b/.github/workflows/packaging.yml index 97350853..d378d843 100644 --- a/.github/workflows/packaging.yml +++ b/.github/workflows/packaging.yml @@ -8,52 +8,52 @@ on: - main env: - CIBW_BUILD: "cp38-* cp39-* cp310-*" + CIBW_BUILD: "cp39-* cp310-* cp311-*" CIBW_SKIP: "*-win32 *musllinux* *i686*" CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014 jobs: - build_wheels: - name: Build ${{ matrix.python-version }} wheels on ${{ matrix.os }} - runs-on: ${{ matrix.os }} - strategy: - fail-fast: false - matrix: - os: [ubuntu-latest, macos-latest] # , windows-latest] + # build_wheels: + # name: Build ${{ matrix.python-version }} wheels on ${{ matrix.os }} + # runs-on: ${{ matrix.os }} + # strategy: + # fail-fast: false + # matrix: + # os: [ubuntu-latest, macos-latest] # , windows-latest] - steps: - - uses: actions/checkout@v2 - with: - fetch-depth: 0 + # steps: + # - uses: actions/checkout@v4 + # with: + # fetch-depth: 0 - - uses: actions/setup-python@v3 - name: Install Python - with: - # Note: cibuildwheel builds for many Python versions beyond this one - python-version: "3.9" + # - uses: actions/setup-python@v3 + # name: Install Python + # with: + # # Note: cibuildwheel builds for many Python versions beyond this one + # python-version: "3.9" - - name: Install MSVC / Visual C++ - if: runner.os == 'Windows' - uses: ilammy/msvc-dev-cmd@v1 + # - name: Install MSVC / Visual C++ + # if: runner.os == 'Windows' + # uses: ilammy/msvc-dev-cmd@v1 - - name: Build wheels - run: | - python -m pip install cibuildwheel - python -m cibuildwheel --output-dir wheelhouse + # - name: Build wheels + # run: | + # python -m pip install cibuildwheel + # python -m cibuildwheel --output-dir wheelhouse - - uses: actions/upload-artifact@v2 - with: - path: ./wheelhouse/*.whl + # - uses: actions/upload-artifact@v2 + # with: + # path: ./wheelhouse/*.whl build_sdist: name: Build source distribution runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 with: fetch-depth: 0 - - uses: actions/setup-python@v2 + - uses: actions/setup-python@v5 name: Install Python with: python-version: "3.9" @@ -68,11 +68,12 @@ jobs: path: dist/*.tar.gz upload_pypi: - needs: [build_wheels, build_sdist] + # needs: [build_wheels, build_sdist] + needs: [build_sdist] runs-on: ubuntu-latest if: github.event_name == 'release' && github.event.action == 'published' steps: - - uses: actions/download-artifact@v2 + - uses: actions/download-artifact@v4 with: name: artifact path: dist diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 0ee68bfb..01141c47 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -16,10 +16,10 @@ jobs: matrix: include: - - name: Python 3.9 with minimal dependencies and coverage + - name: Python 3.11 with minimal dependencies and coverage os: ubuntu-latest - python: 3.9 - toxenv: py39-test-cov + python: "3.11" + toxenv: py311-test-cov - name: Python 3.10 os: ubuntu-latest @@ -33,30 +33,25 @@ jobs: toxposargs: --durations=50 || true # override exit code # Mac: - - name: Python 3.9 standard tests (macOS) + - name: Python 3.11 standard tests (macOS) os: macos-latest - python: 3.9 - toxenv: py39-test + python: "3.11" + toxenv: py311-test # Older Python versions: - - name: Python 3.8 + - name: Python 3.9 with oldest supported version of all dependencies os: ubuntu-latest - python: 3.8 - toxenv: py38-test - - - name: Python 3.7 with oldest supported version of all dependencies - os: ubuntu-latest - python: 3.7 - toxenv: py37-test-oldestdeps + python: 3.9 + toxenv: py39-test-oldestdeps steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 with: fetch-depth: 0 - name: Set up Python ${{ matrix.python-version }} on ${{ matrix.os }} if: "!startsWith(matrix.os, 'windows')" - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python }} @@ -90,7 +85,7 @@ jobs: with: file: ./coverage.xml # optional - - uses: actions/upload-artifact@v2 + - uses: actions/upload-artifact@v4 if: ${{ failure() }} with: path: ./result_images From 6441c023b3332f9e5979429970cf12365d4af00e Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 15 Jan 2024 18:08:36 -0500 Subject: [PATCH 04/50] update packaging numpy version --- pyproject.toml | 121 +++++++++++++++++++++++++++++++++++--- setup.cfg | 100 ------------------------------- setup.py | 81 ------------------------- thejoker/__init__.py | 41 +++++++++---- thejoker/_astropy_init.py | 15 ----- 5 files changed, 143 insertions(+), 215 deletions(-) delete mode 100644 setup.cfg delete mode 100644 setup.py delete mode 100644 thejoker/_astropy_init.py diff --git a/pyproject.toml b/pyproject.toml index f567329e..5fb438ac 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,12 +1,115 @@ [build-system] +requires = [ + "setuptools>=64", + "setuptools_scm>=8", + "wheel", + "numpy<1.22", + "cython", + "twobody" +] +build-backend = 'setuptools.build_meta' -requires = ["setuptools", - "setuptools_scm", - "wheel", - "extension-helpers", - "scipy", - "numpy", - "cython", - "twobody"] +[project] +name = "thejoker" +authors = [{name = "Adrian Price-Whelan", email = "adrianmpw@gmail.com"}] +description = "A custom Monte Carlo sampler for the two-body problem." +readme = "README.rst" +requires-python = ">=3.9" +license.file = "LICENSE" +dynamic = ["version"] +dependencies = [ + "astropy", + "numpy<1.22", + "twobody>=0.9", + "scipy", + "h5py", + "schwimmbad>=0.3.1", + "pymc3", + "pymc_ext", + "exoplanet>=0.2.2", + "tables" +] -build-backend = 'setuptools.build_meta' +[project.urls] +Homepage = "https://github.com/adrn/thejoker" +"Bug Tracker" = "https://github.com/adrn/thejoker/issues" +Discussions = "https://github.com/adrn/thejoker/discussions" +Changelog = "https://github.com/adrn/thejoker/releases" + +[project.optional-dependencies] +test = [ + "pytest", + "pytest-astropy", + "ipython", + "jupyter_client", + "corner", + "astroML", + "pyyaml" +] +docs = [ + "sphinx-astropy", + "nbsphinx", + "nbconvert", + "nbformat", + "ipykernel" +] + +[tool.setuptools_scm] +version_file = "thejoker/_version.py" + +[tool.pytest.ini_options] +testpaths = ["thejoker", "docs"] +doctest_plus = "enabled" +text_file_format = "rst" +addopts = [ + "--doctest-rst", "-ra", "--showlocals", "--strict-markers", "--strict-config" +] +xfail_strict = true +filterwarnings = [ + "error", + "ignore:unclosed file:ResourceWarning", + "ignore:unclosed + "PLR2004", # Magic value used in comparison + "ISC001", # Conflicts with formatter +] \ No newline at end of file diff --git a/setup.cfg b/setup.cfg deleted file mode 100644 index abc0c2fb..00000000 --- a/setup.cfg +++ /dev/null @@ -1,100 +0,0 @@ -[metadata] -name = thejoker -author = Adrian Price-Whelan -author_email = adrianmpw@gmail.com -license = MIT -license_file = LICENSE -url = https://github.com/adrn/thejoker -description = A custom Monte Carlo sampler for the two-body problem. -long_description = file: README.rst -long_description_content_type = text/x-rst -edit_on_github = False -github_project = adrn/thejoker - -[options] -zip_safe = False -packages = find: -python_requires = >=3.7 -install_requires = - numpy - astropy - twobody>=0.9 - scipy - h5py - schwimmbad>=0.3.1 - aesara_theano_fallback - pymc3>=3.7 - pymc3_ext - exoplanet>=0.2.2 - tables - -[options.entry_points] - -[options.extras_require] -# Must be checked against requirements-dev.txt -test = - pytest - pytest-astropy - ipython - jupyter_client - corner - astroML - pyyaml -docs = - sphinx-astropy - ipython - jupyter_client - corner - nbsphinx - nbconvert - nbformat - ipykernel - astroML - pyyaml - pytest - -[options.package_data] -* = *.c -thejoker.src = fast_likelihood.pyx -thejoker.tests = coveragerc - -[tool:pytest] -testpaths = "thejoker" "docs" -astropy_header = true -doctest_plus = enabled -text_file_format = rst -addopts = --doctest-rst -norecursedirs = _build _static examples tmp* - -[coverage:run] -omit = - thejoker/_astropy_init* - thejoker/conftest.py - thejoker/*setup_package* - thejoker/tests/* - thejoker/*/tests/* - */thejoker/_astropy_init* - */thejoker/conftest.py - */thejoker/*setup_package* - */thejoker/tests/* - */thejoker/*/tests/* - -[coverage:report] -exclude_lines = - # Have to re-enable the standard pragma - pragma: no cover - # Don't complain about packages we have installed - except ImportError - # Don't complain if tests don't hit assertions - raise AssertionError - raise NotImplementedError - # Don't complain about script hooks - def main\(.*\): - # Ignore branches that don't pertain to this version of Python - pragma: py{ignore_python_version} - # Don't complain about IPython completion helper - def _ipython_key_completions_ - -[flake8] -exclude = extern,sphinx,*parsetab.py,astropy_helpers,ah_bootstrap.py,conftest.py,docs/conf.py,setup.py -max-line-length = 80 diff --git a/setup.py b/setup.py deleted file mode 100644 index 1fa76f76..00000000 --- a/setup.py +++ /dev/null @@ -1,81 +0,0 @@ -#!/usr/bin/env python -# Licensed under a 3-clause BSD style license - see LICENSE.rst - -# NOTE: The configuration for the package, including the name, version, and -# other information are set in the setup.cfg file. - -import os -import sys - -from setuptools import setup - -from extension_helpers import get_extensions - - -# First provide helpful messages if contributors try and run legacy commands -# for tests or docs. - -TEST_HELP = """ -Note: running tests is no longer done using 'python setup.py test'. Instead -you will need to run: - - tox -e test - -If you don't already have tox installed, you can install it with: - - pip install tox - -If you only want to run part of the test suite, you can also use pytest -directly with:: - - pip install -e .[test] - pytest - -For more information, see: - - http://docs.astropy.org/en/latest/development/testguide.html#running-tests -""" - -if 'test' in sys.argv: - print(TEST_HELP) - sys.exit(1) - -DOCS_HELP = """ -Note: building the documentation is no longer done using -'python setup.py build_docs'. Instead you will need to run: - - tox -e build_docs - -If you don't already have tox installed, you can install it with: - - pip install tox - -You can also build the documentation with Sphinx directly using:: - - pip install -e .[docs] - cd docs - make html - -For more information, see: - - http://docs.astropy.org/en/latest/install.html#builddocs -""" - -if 'build_docs' in sys.argv or 'build_sphinx' in sys.argv: - print(DOCS_HELP) - sys.exit(1) - -VERSION_TEMPLATE = """ -# Note that we need to fall back to the hard-coded version if either -# setuptools_scm can't be imported or setuptools_scm can't determine the -# version, so we catch the generic 'Exception'. -try: - from setuptools_scm import get_version - version = get_version(root='..', relative_to=__file__) -except Exception: - version = '{version}' -""".lstrip() - -setup(use_scm_version={'write_to': os.path.join('thejoker', 'version.py'), - 'write_to_template': VERSION_TEMPLATE}, - ext_modules=get_extensions()) diff --git a/thejoker/__init__.py b/thejoker/__init__.py index 93cdd8cb..67880d00 100644 --- a/thejoker/__init__.py +++ b/thejoker/__init__.py @@ -1,14 +1,35 @@ -from ._astropy_init import * # noqa +from ._version import version as __version__ +from .data import RVData +from .plot import plot_phase_fold, plot_rv_curves +from .prior import JokerPrior +from .samples import JokerSamples +from .samples_analysis import ( + MAP_sample, + is_P_Kmodal, + is_P_unimodal, + max_phase_gap, + periods_spanned, + phase_coverage, + phase_coverage_per_period, +) +from .thejoker import TheJoker - -# For egg_info test builds to pass, put package imports here. -if not _ASTROPY_SETUP_: # noqa - from .thejoker import TheJoker # noqa - from .data import RVData # noqa - from .samples import JokerSamples # noqa - from .prior import JokerPrior # noqa - from .plot import plot_rv_curves, plot_phase_fold # noqa - from .samples_analysis import * # noqa +__all__ = [ + "__version__", + "TheJoker", + "RVData", + "JokerSamples", + "JokerPrior", + "plot_rv_curves", + "plot_phase_fold", + "MAP_sample", + "is_P_unimodal", + "is_P_Kmodal", + "max_phase_gap", + "phase_coverage", + "periods_spanned", + "phase_coverage_per_period", +] __bibtex__ = __citation__ = """@ARTICLE{thejoker, diff --git a/thejoker/_astropy_init.py b/thejoker/_astropy_init.py deleted file mode 100644 index 058ecc61..00000000 --- a/thejoker/_astropy_init.py +++ /dev/null @@ -1,15 +0,0 @@ -# Licensed under a 3-clause BSD style license - see LICENSE.rst - -__all__ = ['__version__'] - -# this indicates whether or not we are in the package's setup.py -try: - _ASTROPY_SETUP_ -except NameError: - import builtins - builtins._ASTROPY_SETUP_ = False - -try: - from .version import version as __version__ -except ImportError: - __version__ = '' From a1ac12d35d843b3a37b1e748fe330a8ad98a84e7 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 15 Jan 2024 18:35:19 -0500 Subject: [PATCH 05/50] restrict python version --- .gitignore | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index bcaf1835..66917297 100644 --- a/.gitignore +++ b/.gitignore @@ -10,7 +10,7 @@ __pycache__ *.c # Other generated files -*/version.py +*/_version.py */cython_version.py htmlcov .coverage diff --git a/pyproject.toml b/pyproject.toml index 5fb438ac..e1e08e71 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,7 @@ name = "thejoker" authors = [{name = "Adrian Price-Whelan", email = "adrianmpw@gmail.com"}] description = "A custom Monte Carlo sampler for the two-body problem." readme = "README.rst" -requires-python = ">=3.9" +requires-python = ">=3.7,<3.11" license.file = "LICENSE" dynamic = ["version"] dependencies = [ From c6d103eac7c3daf94ec3c7fb72c0812c11c9b78e Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 15 Jan 2024 18:36:11 -0500 Subject: [PATCH 06/50] remove 3.11 --- .github/workflows/packaging.yml | 2 +- .github/workflows/tests.yml | 12 ++++++------ 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/.github/workflows/packaging.yml b/.github/workflows/packaging.yml index d378d843..af0b6aff 100644 --- a/.github/workflows/packaging.yml +++ b/.github/workflows/packaging.yml @@ -8,7 +8,7 @@ on: - main env: - CIBW_BUILD: "cp39-* cp310-* cp311-*" + CIBW_BUILD: "cp38-* cp39-* cp310-*" CIBW_SKIP: "*-win32 *musllinux* *i686*" CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014 diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 01141c47..6a97bb88 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -16,10 +16,10 @@ jobs: matrix: include: - - name: Python 3.11 with minimal dependencies and coverage + - name: Python 3.10 with minimal dependencies and coverage os: ubuntu-latest - python: "3.11" - toxenv: py311-test-cov + python: "3.10" + toxenv: py310-test-cov - name: Python 3.10 os: ubuntu-latest @@ -33,10 +33,10 @@ jobs: toxposargs: --durations=50 || true # override exit code # Mac: - - name: Python 3.11 standard tests (macOS) + - name: Python 3.10 standard tests (macOS) os: macos-latest - python: "3.11" - toxenv: py311-test + python: "3.10" + toxenv: py310-test # Older Python versions: - name: Python 3.9 with oldest supported version of all dependencies From 7fbf20e8a93155d21f065902efb0a894dc8853e0 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 15 Jan 2024 19:04:18 -0500 Subject: [PATCH 07/50] oops - oldest numpy issue --- tox.ini | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tox.ini b/tox.ini index 10094d5a..7b1c8f7e 100644 --- a/tox.ini +++ b/tox.ini @@ -37,7 +37,7 @@ deps = # The oldestdeps factor is intended to be used to install the oldest # versions of all dependencies that have a minimum version. - oldestdeps: numpy==1.24.* + oldestdeps: numpy==1.20.* oldestdeps: matplotlib==3.4.* oldestdeps: scipy==1.6.* oldestdeps: astropy==5.2.* From cc7cf6ae39821aee57ab5cb377e92f466b509275 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 15 Jan 2024 21:19:49 -0500 Subject: [PATCH 08/50] oof - pinning versions --- pyproject.toml | 19 ++++++++++++++----- thejoker/prior.py | 2 +- thejoker/src/setup_package.py | 3 +++ tox.ini | 4 ++-- 4 files changed, 20 insertions(+), 8 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e1e08e71..bc8692c9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,9 @@ requires = [ "setuptools>=64", "setuptools_scm>=8", "wheel", - "numpy<1.22", + "extension-helpers==1.*", + "numpy", + "scipy", "cython", "twobody" ] @@ -14,12 +16,12 @@ name = "thejoker" authors = [{name = "Adrian Price-Whelan", email = "adrianmpw@gmail.com"}] description = "A custom Monte Carlo sampler for the two-body problem." readme = "README.rst" -requires-python = ">=3.7,<3.11" +requires-python = ">=3.7,<3.10" license.file = "LICENSE" dynamic = ["version"] dependencies = [ - "astropy", - "numpy<1.22", + "astropy<6.0", + "numpy", "twobody>=0.9", "scipy", "h5py", @@ -27,7 +29,7 @@ dependencies = [ "pymc3", "pymc_ext", "exoplanet>=0.2.2", - "tables" + "tables", ] [project.urls] @@ -57,6 +59,9 @@ docs = [ [tool.setuptools_scm] version_file = "thejoker/_version.py" +[tool.extension-helpers] +use_extension_helpers = true + [tool.pytest.ini_options] testpaths = ["thejoker", "docs"] doctest_plus = "enabled" @@ -70,6 +75,10 @@ filterwarnings = [ "ignore:unclosed file:ResourceWarning", "ignore:unclosed Date: Mon, 15 Jan 2024 21:24:13 -0500 Subject: [PATCH 09/50] ditch python 3.10 and scipy path --- .github/workflows/packaging.yml | 2 +- .github/workflows/tests.yml | 18 +++++++++--------- thejoker/src/setup_package.py | 3 --- 3 files changed, 10 insertions(+), 13 deletions(-) diff --git a/.github/workflows/packaging.yml b/.github/workflows/packaging.yml index af0b6aff..6ad1d482 100644 --- a/.github/workflows/packaging.yml +++ b/.github/workflows/packaging.yml @@ -8,7 +8,7 @@ on: - main env: - CIBW_BUILD: "cp38-* cp39-* cp310-*" + CIBW_BUILD: "cp38-* cp39-*" CIBW_SKIP: "*-win32 *musllinux* *i686*" CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014 diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 6a97bb88..5852e0f8 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -16,15 +16,15 @@ jobs: matrix: include: - - name: Python 3.10 with minimal dependencies and coverage + - name: Python 3.9 with minimal dependencies and coverage os: ubuntu-latest - python: "3.10" - toxenv: py310-test-cov + python: "3.9" + toxenv: py39-test-cov - - name: Python 3.10 + - name: Python 3.9 os: ubuntu-latest - python: '3.10' - toxenv: py310-test + python: '3.9' + toxenv: py39-test - name: Python 3.9 dev dependencies (allowed failure! check logs) os: ubuntu-latest @@ -33,10 +33,10 @@ jobs: toxposargs: --durations=50 || true # override exit code # Mac: - - name: Python 3.10 standard tests (macOS) + - name: Python 3.9 standard tests (macOS) os: macos-latest - python: "3.10" - toxenv: py310-test + python: "3.9" + toxenv: py39-test # Older Python versions: - name: Python 3.9 with oldest supported version of all dependencies diff --git a/thejoker/src/setup_package.py b/thejoker/src/setup_package.py index 21729583..1c9da759 100644 --- a/thejoker/src/setup_package.py +++ b/thejoker/src/setup_package.py @@ -7,16 +7,13 @@ def get_extensions(): exts = [] import numpy as np - import scipy import twobody cfg = defaultdict(list) cfg['include_dirs'].append(np.get_include()) twobody_path = os.path.dirname(twobody.__file__) - scipy_path = os.path.dirname(scipy.__file__) cfg['include_dirs'].append(twobody_path) - cfg['include_dirs'].append(scipy_path) cfg['sources'].append(os.path.join(twobody_path, 'src/twobody.c')) cfg['extra_compile_args'].append('--std=gnu99') From f1f46bb249bda66801ef37bec8d3825e2383e5c2 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Sat, 2 Mar 2024 19:11:20 -0500 Subject: [PATCH 10/50] change random_state to rng and some reformat --- docs/examples/1-Getting-started.ipynb | 4 +- docs/examples/2-Customize-prior.ipynb | 36 +- .../3-Polynomial-velocity-trend.ipynb | 16 +- docs/examples/4-Continue-sampling-mcmc.ipynb | 16 +- docs/examples/5-Calibration-offsets.ipynb | 14 +- docs/examples/Strader-circular-only.ipynb | 34 +- docs/examples/Thompson-black-hole.ipynb | 44 +-- pyproject.toml | 2 +- thejoker/likelihood_helpers.py | 115 +++--- thejoker/multiproc_helpers.py | 329 ++++++++++-------- thejoker/prior.py | 316 ++++++++++------- thejoker/src/fast_likelihood.pyx | 4 +- thejoker/tests/test_sampler.py | 156 +++++---- thejoker/thejoker.py | 235 ++++++------- thejoker/utils.py | 144 ++++---- 15 files changed, 805 insertions(+), 660 deletions(-) diff --git a/docs/examples/1-Getting-started.ipynb b/docs/examples/1-Getting-started.ipynb index 1d67a2df..ee49ee1b 100644 --- a/docs/examples/1-Getting-started.ipynb +++ b/docs/examples/1-Getting-started.ipynb @@ -215,7 +215,7 @@ "outputs": [], "source": [ "prior_samples = prior.sample(size=250_000,\n", - " random_state=rnd)\n", + " rng=rnd)\n", "prior_samples" ] }, @@ -291,7 +291,7 @@ "metadata": {}, "outputs": [], "source": [ - "joker = tj.TheJoker(prior, random_state=rnd)\n", + "joker = tj.TheJoker(prior, rng=rnd)\n", "joker_samples = joker.rejection_sample(data, prior_samples, \n", " max_posterior_samples=256)" ] diff --git a/docs/examples/2-Customize-prior.ipynb b/docs/examples/2-Customize-prior.ipynb index a9ce1500..09014a3b 100644 --- a/docs/examples/2-Customize-prior.ipynb +++ b/docs/examples/2-Customize-prior.ipynb @@ -74,13 +74,13 @@ "with pm.Model() as model:\n", " P = xu.with_unit(pm.Normal('P', 50., 1),\n", " u.day)\n", - " \n", + "\n", " prior = tj.JokerPrior.default(\n", " sigma_K0=30*u.km/u.s,\n", " sigma_v=100*u.km/u.s,\n", " pars={'P': P})\n", - " \n", - "samples1 = prior.sample(size=100_000, random_state=rnd)" + "\n", + "samples1 = prior.sample(size=100_000, rng=rnd)" ] }, { @@ -120,12 +120,12 @@ " u.day)\n", " K = xu.with_unit(pm.Normal('K', 0., 15),\n", " u.km/u.s)\n", - " \n", + "\n", " prior = tj.JokerPrior.default(\n", " sigma_v=100*u.km/u.s,\n", " pars={'P': P, 'K': K})\n", - " \n", - "samples2 = prior.sample(size=100_000, random_state=rnd)\n", + "\n", + "samples2 = prior.sample(size=100_000, rng=rnd)\n", "samples2" ] }, @@ -142,8 +142,8 @@ "metadata": {}, "outputs": [], "source": [ - "samples3 = prior.sample(size=100_000, generate_linear=True, \n", - " random_state=rnd)\n", + "samples3 = prior.sample(size=100_000, generate_linear=True,\n", + " rng=rnd)\n", "samples3" ] }, @@ -175,7 +175,7 @@ " sigma_K0=30*u.km/u.s,\n", " sigma_v=75*u.km/u.s)\n", "default_samples = default_prior.sample(size=20, generate_linear=True,\n", - " random_state=rnd,\n", + " rng=rnd,\n", " t_ref=Time('J2000')) # set arbitrary time zero-point" ] }, @@ -189,14 +189,14 @@ " K = xu.with_unit(pm.Normal('K', 0., 30),\n", " u.km/u.s)\n", " custom_prior = tj.JokerPrior.default(\n", - " P_min=1e1*u.day, \n", - " P_max=1e3*u.day, \n", + " P_min=1e1*u.day,\n", + " P_max=1e3*u.day,\n", " sigma_v=75*u.km/u.s,\n", " pars={'K': K})\n", - " \n", - "custom_samples = custom_prior.sample(size=len(default_samples), \n", - " generate_linear=True, \n", - " random_state=rnd,\n", + "\n", + "custom_samples = custom_prior.sample(size=len(default_samples),\n", + " generate_linear=True,\n", + " rng=rnd,\n", " t_ref=Time('J2000')) # set arbitrary time zero-point" ] }, @@ -207,13 +207,13 @@ "outputs": [], "source": [ "now_mjd = Time.now().mjd\n", - "t_grid = Time(np.linspace(now_mjd - 1000, now_mjd + 1000, 16384), \n", + "t_grid = Time(np.linspace(now_mjd - 1000, now_mjd + 1000, 16384),\n", " format='mjd')\n", "\n", "fig, axes = plt.subplots(2, 1, sharex=True, sharey=True, figsize=(8, 8))\n", - "_ = tj.plot_rv_curves(default_samples, t_grid=t_grid, \n", + "_ = tj.plot_rv_curves(default_samples, t_grid=t_grid,\n", " ax=axes[0], add_labels=False)\n", - "_ = tj.plot_rv_curves(custom_samples, t_grid=t_grid, \n", + "_ = tj.plot_rv_curves(custom_samples, t_grid=t_grid,\n", " ax=axes[1])\n", "axes[0].set_ylim(-200, 200)\n", "fig.tight_layout()" diff --git a/docs/examples/3-Polynomial-velocity-trend.ipynb b/docs/examples/3-Polynomial-velocity-trend.ipynb index 6f66c03a..f86e1cf0 100644 --- a/docs/examples/3-Polynomial-velocity-trend.ipynb +++ b/docs/examples/3-Polynomial-velocity-trend.ipynb @@ -96,7 +96,7 @@ " P_min=2*u.day, P_max=1e3*u.day,\n", " sigma_K0=30*u.km/u.s,\n", " sigma_v=100*u.km/u.s)\n", - "prior_samples = prior.sample(size=250_000, random_state=rnd)" + "prior_samples = prior.sample(size=250_000, rng=rnd)" ] }, { @@ -112,8 +112,8 @@ "metadata": {}, "outputs": [], "source": [ - "joker = tj.TheJoker(prior, random_state=rnd)\n", - "samples = joker.rejection_sample(data, prior_samples, \n", + "joker = tj.TheJoker(prior, rng=rnd)\n", + "samples = joker.rejection_sample(data, prior_samples,\n", " max_posterior_samples=128)\n", "samples" ] @@ -143,8 +143,8 @@ "prior_trend = tj.JokerPrior.default(\n", " P_min=2*u.day, P_max=1e3*u.day,\n", " sigma_K0=30*u.km/u.s,\n", - " sigma_v=[100*u.km/u.s, \n", - " 0.5*u.km/u.s/u.day, \n", + " sigma_v=[100*u.km/u.s,\n", + " 0.5*u.km/u.s/u.day,\n", " 1e-2*u.km/u.s/u.day**2],\n", " poly_trend=3)" ] @@ -179,9 +179,9 @@ "outputs": [], "source": [ "prior_samples_trend = prior_trend.sample(size=250_000,\n", - " random_state=rnd)\n", - "joker_trend = tj.TheJoker(prior_trend, random_state=rnd)\n", - "samples_trend = joker_trend.rejection_sample(data, prior_samples_trend, \n", + " rng=rnd)\n", + "joker_trend = tj.TheJoker(prior_trend, rng=rnd)\n", + "samples_trend = joker_trend.rejection_sample(data, prior_samples_trend,\n", " max_posterior_samples=128)\n", "samples_trend" ] diff --git a/docs/examples/4-Continue-sampling-mcmc.ipynb b/docs/examples/4-Continue-sampling-mcmc.ipynb index 1cf5bd7b..3aae6b32 100644 --- a/docs/examples/4-Continue-sampling-mcmc.ipynb +++ b/docs/examples/4-Continue-sampling-mcmc.ipynb @@ -123,8 +123,8 @@ "metadata": {}, "outputs": [], "source": [ - "prior_samples = prior.sample(size=250_000, \n", - " random_state=rnd)" + "prior_samples = prior.sample(size=250_000,\n", + " rng=rnd)" ] }, { @@ -133,8 +133,8 @@ "metadata": {}, "outputs": [], "source": [ - "joker = tj.TheJoker(prior, random_state=rnd)\n", - "joker_samples = joker.rejection_sample(data, prior_samples, \n", + "joker = tj.TheJoker(prior, rng=rnd)\n", + "joker_samples = joker.rejection_sample(data, prior_samples,\n", " max_posterior_samples=256)\n", "joker_samples" ] @@ -172,8 +172,8 @@ "source": [ "with prior.model:\n", " mcmc_init = joker.setup_mcmc(data, joker_samples)\n", - " \n", - " trace = pmx.sample(tune=500, draws=500, \n", + "\n", + " trace = pmx.sample(tune=500, draws=500,\n", " start=mcmc_init,\n", " cores=1, chains=2)" ] @@ -228,11 +228,11 @@ "import pickle\n", "with open('true-orbit.pkl', 'rb') as f:\n", " truth = pickle.load(f)\n", - " \n", + "\n", "# make sure the angles are wrapped the same way\n", "if np.median(mcmc_samples['omega']) < 0:\n", " truth['omega'] = coord.Angle(truth['omega']).wrap_at(np.pi*u.radian)\n", - " \n", + "\n", "if np.median(mcmc_samples['M0']) < 0:\n", " truth['M0'] = coord.Angle(truth['M0']).wrap_at(np.pi*u.radian)" ] diff --git a/docs/examples/5-Calibration-offsets.ipynb b/docs/examples/5-Calibration-offsets.ipynb index 218df335..245286b2 100644 --- a/docs/examples/5-Calibration-offsets.ipynb +++ b/docs/examples/5-Calibration-offsets.ipynb @@ -112,7 +112,7 @@ "with pm.Model() as model:\n", " dv0_1 = xu.with_unit(pm.Normal('dv0_1', 0, 10),\n", " u.km/u.s)\n", - " \n", + "\n", " prior = tj.JokerPrior.default(\n", " P_min=2*u.day, P_max=256*u.day,\n", " sigma_K0=30*u.km/u.s,\n", @@ -134,7 +134,7 @@ "outputs": [], "source": [ "prior_samples = prior.sample(size=1_000_000,\n", - " random_state=rnd)" + " rng=rnd)" ] }, { @@ -143,8 +143,8 @@ "metadata": {}, "outputs": [], "source": [ - "joker = tj.TheJoker(prior, random_state=rnd)\n", - "joker_samples = joker.rejection_sample(data, prior_samples, \n", + "joker = tj.TheJoker(prior, rng=rnd)\n", + "joker_samples = joker.rejection_sample(data, prior_samples,\n", " max_posterior_samples=128)\n", "joker_samples" ] @@ -180,7 +180,7 @@ "metadata": {}, "outputs": [], "source": [ - "_ = tj.plot_rv_curves(joker_samples, data=data, \n", + "_ = tj.plot_rv_curves(joker_samples, data=data,\n", " apply_mean_v0_offset=False)" ] }, @@ -199,9 +199,9 @@ "source": [ "with prior.model:\n", " mcmc_init = joker.setup_mcmc(data, joker_samples)\n", - " \n", + "\n", " trace = pmx.sample(\n", - " tune=500, draws=500, \n", + " tune=500, draws=500,\n", " start=mcmc_init,\n", " cores=1, chains=2)" ] diff --git a/docs/examples/Strader-circular-only.ipynb b/docs/examples/Strader-circular-only.ipynb index 368c11c0..9747d77c 100644 --- a/docs/examples/Strader-circular-only.ipynb +++ b/docs/examples/Strader-circular-only.ipynb @@ -113,7 +113,7 @@ " 2458247.5062024 131.9 11.5\n", " 2458247.5435496 160.5 14.2\n", " 2458278.5472619 197.1 15.9\n", - " 2458278.5613912 183.7 15.7\"\"\", \n", + " 2458278.5613912 183.7 15.7\"\"\",\n", " names=['BJD', 'rv', 'rv_err'])\n", "tbl['rv'].unit = u.km/u.s\n", "tbl['rv_err'].unit = u.km/u.s" @@ -134,7 +134,7 @@ "source": [ "data = tj.RVData(\n", " t=Time(tbl['BJD'], format='jd', scale='tcb'),\n", - " rv=u.Quantity(tbl['rv']), \n", + " rv=u.Quantity(tbl['rv']),\n", " rv_err=u.Quantity(tbl['rv_err']))" ] }, @@ -169,9 +169,9 @@ "source": [ "with pm.Model() as model:\n", " # Allow extra error to account for under-estimated error bars\n", - " e = xu.with_unit(pm.Constant('e', 0), \n", + " e = xu.with_unit(pm.Constant('e', 0),\n", " u.one)\n", - " \n", + "\n", " prior = tj.JokerPrior.default(\n", " P_min=0.1*u.day, P_max=100*u.day, # Range of periods to consider\n", " sigma_K0=50*u.km/u.s, P0=1*u.year, # scale of the prior on semiamplitude, K\n", @@ -194,8 +194,8 @@ "outputs": [], "source": [ "# Run rejection sampling with The Joker:\n", - "joker = tj.TheJoker(prior, random_state=rnd)\n", - "samples = joker.rejection_sample(data, \n", + "joker = tj.TheJoker(prior, rng=rnd)\n", + "samples = joker.rejection_sample(data,\n", " prior_samples=100_000,\n", " max_posterior_samples=256)\n", "samples" @@ -248,15 +248,15 @@ "source": [ "import aesara_theano_fallback.tensor as tt\n", "with pm.Model():\n", - " \n", - " # To sample with pymc3, we have to set any constant variables \n", - " # as \"Deterministic\" objects. We can ignore eccentricity and \n", + "\n", + " # To sample with pymc3, we have to set any constant variables\n", + " # as \"Deterministic\" objects. We can ignore eccentricity and\n", " # the argument of pericenter by setting them both to 0:\n", - " e = xu.with_unit(pm.Deterministic('e', tt.constant(0)), \n", + " e = xu.with_unit(pm.Deterministic('e', tt.constant(0)),\n", " u.one)\n", - " omega = xu.with_unit(pm.Deterministic('omega', tt.constant(0)), \n", + " omega = xu.with_unit(pm.Deterministic('omega', tt.constant(0)),\n", " u.radian)\n", - " \n", + "\n", " # We use the same prior parameters as before:\n", " prior_mcmc = tj.JokerPrior.default(\n", " P_min=0.1*u.day, P_max=10*u.day,\n", @@ -264,12 +264,12 @@ " sigma_v=50*u.km/u.s,\n", " pars={'e': e, 'omega': omega}\n", " )\n", - " \n", - " # Now we use the sample returned from The Joker to set up \n", + "\n", + " # Now we use the sample returned from The Joker to set up\n", " # our initialization for standard MCMC:\n", - " joker_mcmc = tj.TheJoker(prior_mcmc, random_state=rnd)\n", + " joker_mcmc = tj.TheJoker(prior_mcmc, rng=rnd)\n", " mcmc_init = joker_mcmc.setup_mcmc(data, samples)\n", - " \n", + "\n", " trace = pmx.sample(\n", " tune=500, draws=1000,\n", " start=mcmc_init,\n", @@ -352,7 +352,7 @@ "\n", "for ax in axes:\n", " ax.set_ylabel(f'RV [{data.rv.unit:latex_inline}]')\n", - " \n", + "\n", "axes[1].axhline(0, zorder=-10, color='tab:green', alpha=0.5)\n", "axes[1].set_ylim(-50, 50)" ] diff --git a/docs/examples/Thompson-black-hole.ipynb b/docs/examples/Thompson-black-hole.ipynb index 5e80cd37..8f4f3417 100644 --- a/docs/examples/Thompson-black-hole.ipynb +++ b/docs/examples/Thompson-black-hole.ipynb @@ -81,7 +81,7 @@ " 8112.81800 -44.863 0.088\n", " 8123.79627 -25.810 0.115\n", " 8136.59960 15.691 0.146\n", - " 8143.78352 34.281 0.087\"\"\", \n", + " 8143.78352 34.281 0.087\"\"\",\n", " names=['HJD', 'rv', 'rv_err'])\n", "tres_tbl['rv'].unit = u.km/u.s\n", "tres_tbl['rv_err'].unit = u.km/u.s" @@ -96,7 +96,7 @@ "apogee_tbl = ascii.read(\n", " \"\"\"6204.95544 -37.417 0.011\n", " 6229.92499 34.846 0.010\n", - " 6233.87715 42.567 0.010\"\"\", \n", + " 6233.87715 42.567 0.010\"\"\",\n", " names=['HJD', 'rv', 'rv_err'])\n", "apogee_tbl['rv'].unit = u.km/u.s\n", "apogee_tbl['rv_err'].unit = u.km/u.s" @@ -110,12 +110,12 @@ "source": [ "tres_data = tj.RVData(\n", " t=Time(tres_tbl['HJD'] + 2450000, format='jd', scale='tcb'),\n", - " rv=u.Quantity(tres_tbl['rv']), \n", + " rv=u.Quantity(tres_tbl['rv']),\n", " rv_err=u.Quantity(tres_tbl['rv_err']))\n", "\n", "apogee_data = tj.RVData(\n", " t=Time(apogee_tbl['HJD'] + 2450000, format='jd', scale='tcb'),\n", - " rv=u.Quantity(apogee_tbl['rv']), \n", + " rv=u.Quantity(apogee_tbl['rv']),\n", " rv_err=u.Quantity(apogee_tbl['rv_err']))" ] }, @@ -178,7 +178,7 @@ " # Allow extra error to account for under-estimated error bars\n", " s = xu.with_unit(pm.Lognormal('s', -2, 1),\n", " u.km/u.s)\n", - " \n", + "\n", " prior = tj.JokerPrior.default(\n", " P_min=16*u.day, P_max=128*u.day, # Range of periods to consider\n", " sigma_K0=30*u.km/u.s, P0=1*u.year, # scale of the prior on semiamplitude, K\n", @@ -202,7 +202,7 @@ "source": [ "# Generate a large number of prior samples:\n", "prior_samples = prior.sample(size=1_000_000,\n", - " random_state=rnd)" + " rng=rnd)" ] }, { @@ -212,7 +212,7 @@ "outputs": [], "source": [ "# Run rejection sampling with The Joker:\n", - "joker = tj.TheJoker(prior, random_state=rnd)\n", + "joker = tj.TheJoker(prior, rng=rnd)\n", "samples = joker.rejection_sample(tres_data, prior_samples,\n", " max_posterior_samples=256)\n", "samples" @@ -337,11 +337,11 @@ " # APOGEE and TRES:\n", " dv0_1 = xu.with_unit(pm.Normal('dv0_1', 0, 5.),\n", " u.km/u.s)\n", - " \n", + "\n", " # The same extra uncertainty parameter as previously defined\n", " s = xu.with_unit(pm.Lognormal('s', -2, 1),\n", " u.km/u.s)\n", - " \n", + "\n", " # We can restrict the prior on prior now, using the above\n", " prior_joint = tj.JokerPrior.default(\n", " # P_min=16*u.day, P_max=128*u.day,\n", @@ -351,9 +351,9 @@ " v0_offsets=[dv0_1],\n", " s=s\n", " )\n", - " \n", - "prior_samples_joint = prior_joint.sample(size=10_000_000, \n", - " random_state=rnd)" + "\n", + "prior_samples_joint = prior_joint.sample(size=10_000_000,\n", + " rng=rnd)" ] }, { @@ -363,8 +363,8 @@ "outputs": [], "source": [ "# Run rejection sampling with The Joker:\n", - "joker_joint = tj.TheJoker(prior_joint, random_state=rnd)\n", - "samples_joint = joker_joint.rejection_sample(data, \n", + "joker_joint = tj.TheJoker(prior_joint, rng=rnd)\n", + "samples_joint = joker_joint.rejection_sample(data,\n", " prior_samples_joint,\n", " max_posterior_samples=256)\n", "samples_joint" @@ -402,7 +402,7 @@ "from pymc3_ext.distributions import Angle\n", "\n", "with pm.Model():\n", - " \n", + "\n", " # See note above: when running MCMC, we will sample in the parameters\n", " # (M0 - omega, omega) instead of (M0, omega)\n", " M0_m_omega = xu.with_unit(Angle('M0_m_omega'), u.radian)\n", @@ -410,12 +410,12 @@ " # M0 = xu.with_unit(Angle('M0'), u.radian)\n", " M0 = xu.with_unit(pm.Deterministic('M0', M0_m_omega + omega),\n", " u.radian)\n", - " \n", + "\n", " # The same offset and extra uncertainty parameters as above:\n", " dv0_1 = xu.with_unit(pm.Normal('dv0_1', 0, 5.), u.km/u.s)\n", " s = xu.with_unit(pm.Lognormal('s', -2, 0.5),\n", " u.km/u.s)\n", - " \n", + "\n", " prior_mcmc = tj.JokerPrior.default(\n", " P_min=16*u.day, P_max=128*u.day,\n", " sigma_K0=30*u.km/u.s, P0=1*u.year,\n", @@ -424,10 +424,10 @@ " s=s,\n", " pars={'M0': M0, 'omega': omega}\n", " )\n", - " \n", - " joker_mcmc = tj.TheJoker(prior_mcmc, random_state=rnd)\n", + "\n", + " joker_mcmc = tj.TheJoker(prior_mcmc, rng=rnd)\n", " mcmc_init = joker_mcmc.setup_mcmc(data, samples_joint)\n", - " \n", + "\n", " trace = pmx.sample(\n", " tune=500, draws=1000,\n", " start=mcmc_init,\n", @@ -505,7 +505,7 @@ "\n", "for ax in axes:\n", " ax.set_ylabel(f'RV [{apogee_data.rv.unit:latex_inline}]')\n", - " \n", + "\n", "axes[1].axhline(0, zorder=-10, color='tab:green', alpha=0.5)\n", "axes[1].set_ylim(-1, 1)" ] @@ -532,7 +532,7 @@ "plt.axvline(0.766, zorder=100, color='tab:orange')\n", "plt.axvspan(0.766 - 0.00637,\n", " 0.766 + 0.00637,\n", - " zorder=10, color='tab:orange', \n", + " zorder=10, color='tab:orange',\n", " alpha=0.4, lw=0)" ] }, diff --git a/pyproject.toml b/pyproject.toml index bc8692c9..d8cc7d9a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ dependencies = [ "scipy", "h5py", "schwimmbad>=0.3.1", - "pymc3", + "pymc", "pymc_ext", "exoplanet>=0.2.2", "tables", diff --git a/thejoker/likelihood_helpers.py b/thejoker/likelihood_helpers.py index 750d625e..257a9181 100644 --- a/thejoker/likelihood_helpers.py +++ b/thejoker/likelihood_helpers.py @@ -18,9 +18,9 @@ def get_constant_term_design_matrix(data, ids=None): unq_ids = np.unique(ids) constant_part = np.zeros((len(data), len(unq_ids))) - constant_part[:, 0] = 1. + constant_part[:, 0] = 1.0 for j, id_ in enumerate(unq_ids[1:]): - constant_part[ids == id_, j+1] = 1. + constant_part[ids == id_, j + 1] = 1.0 return constant_part @@ -66,32 +66,37 @@ def marginal_ln_likelihood_inmem(joker_helper, prior_samples_batch): return np.array(ll) -def make_full_samples_inmem(joker_helper, prior_samples_batch, random_state, - n_linear_samples=1): +def make_full_samples_inmem(joker_helper, prior_samples_batch, rng, n_linear_samples=1): from .samples import JokerSamples if prior_samples_batch.dtype != np.float64: prior_samples_batch = prior_samples_batch.astype(np.float64) raw_samples, _ = joker_helper.batch_get_posterior_samples( - prior_samples_batch, n_linear_samples, random_state) + prior_samples_batch, n_linear_samples, rng + ) # unpack the raw samples - samples = JokerSamples.unpack(raw_samples, - joker_helper.internal_units, - t_ref=joker_helper.data.t_ref, - poly_trend=joker_helper.prior.poly_trend, - n_offsets=joker_helper.prior.n_offsets) + samples = JokerSamples.unpack( + raw_samples, + joker_helper.internal_units, + t_ref=joker_helper.data.t_ref, + poly_trend=joker_helper.prior.poly_trend, + n_offsets=joker_helper.prior.n_offsets, + ) return samples -def rejection_sample_inmem(joker_helper, prior_samples_batch, random_state, - ln_prior=None, - max_posterior_samples=None, - n_linear_samples=1, - return_all_logprobs=False): - +def rejection_sample_inmem( + joker_helper, + prior_samples_batch, + rng, + ln_prior=None, + max_posterior_samples=None, + n_linear_samples=1, + return_all_logprobs=False, +): if max_posterior_samples is None: max_posterior_samples = len(prior_samples_batch) @@ -99,19 +104,21 @@ def rejection_sample_inmem(joker_helper, prior_samples_batch, random_state, lls = marginal_ln_likelihood_inmem(joker_helper, prior_samples_batch) # get indices of samples that pass rejection step - uu = random_state.uniform(size=len(lls)) + uu = rng.uniform(size=len(lls)) good_samples_idx = np.where(np.exp(lls - lls.max()) > uu)[0] good_samples_idx = good_samples_idx[:max_posterior_samples] # generate linear parameters - samples = make_full_samples_inmem(joker_helper, - prior_samples_batch[good_samples_idx], - random_state, - n_linear_samples=n_linear_samples) + samples = make_full_samples_inmem( + joker_helper, + prior_samples_batch[good_samples_idx], + rng, + n_linear_samples=n_linear_samples, + ) if ln_prior is not None and ln_prior is not False: - samples['ln_prior'] = ln_prior[good_samples_idx] - samples['ln_likelihood'] = lls[good_samples_idx] + samples["ln_prior"] = ln_prior[good_samples_idx] + samples["ln_likelihood"] = lls[good_samples_idx] if return_all_logprobs: return samples, lls @@ -120,13 +127,16 @@ def rejection_sample_inmem(joker_helper, prior_samples_batch, random_state, return samples -def iterative_rejection_inmem(joker_helper, prior_samples_batch, random_state, - n_requested_samples, - ln_prior=None, - init_batch_size=None, - growth_factor=128, - n_linear_samples=1): - +def iterative_rejection_inmem( + joker_helper, + prior_samples_batch, + rng, + n_requested_samples, + ln_prior=None, + init_batch_size=None, + growth_factor=128, + n_linear_samples=1, +): n_total_samples = len(prior_samples_batch) # The "magic numbers" below control how fast the iterative batches grow @@ -139,12 +149,14 @@ def iterative_rejection_inmem(joker_helper, prior_samples_batch, random_state, n_process = init_batch_size if n_process > n_total_samples: - raise ValueError("Prior sample library not big enough! For " - "iterative sampling, you have to have at least " - "growth_factor * n_requested_samples = " - f"{growth_factor * n_requested_samples} samples in " - "the prior samples cache file. You have, or have " - f"limited to, {n_total_samples} samples.") + raise ValueError( + "Prior sample library not big enough! For " + "iterative sampling, you have to have at least " + "growth_factor * n_requested_samples = " + f"{growth_factor * n_requested_samples} samples in " + "the prior samples cache file. You have, or have " + f"limited to, {n_total_samples} samples." + ) all_idx = np.arange(0, n_total_samples, 1) @@ -154,18 +166,21 @@ def iterative_rejection_inmem(joker_helper, prior_samples_batch, random_state, logger.log(1, f"iteration {i}, computing {n_process} likelihoods") marg_lls = marginal_ln_likelihood_inmem( - joker_helper, prior_samples_batch[start_idx:start_idx + n_process]) + joker_helper, prior_samples_batch[start_idx : start_idx + n_process] + ) all_marg_lls = np.concatenate((all_marg_lls, marg_lls)) if np.any(~np.isfinite(all_marg_lls)): - return RuntimeError("There are NaN or Inf likelihood values in " - f"iteration step {i}!") + return RuntimeError( + "There are NaN or Inf likelihood values in " f"iteration step {i}!" + ) elif len(all_marg_lls) == 0: - return RuntimeError("No likelihood values returned in iteration " - f"step {i}") + return RuntimeError( + "No likelihood values returned in iteration " f"step {i}" + ) # get indices of samples that pass rejection step - uu = random_state.uniform(size=len(all_marg_lls)) + uu = rng.uniform(size=len(all_marg_lls)) aa = np.exp(all_marg_lls - all_marg_lls.max()) good_samples_idx = np.where(aa > uu)[0] @@ -199,18 +214,20 @@ def iterative_rejection_inmem(joker_helper, prior_samples_batch, random_state, full_samples_idx = all_idx[good_samples_idx] # generate linear parameters - samples = make_full_samples_inmem(joker_helper, - prior_samples_batch[full_samples_idx], - random_state, - n_linear_samples=n_linear_samples) + samples = make_full_samples_inmem( + joker_helper, + prior_samples_batch[full_samples_idx], + rng, + n_linear_samples=n_linear_samples, + ) # FIXME: copy-pasted from function above if ln_prior is not None and ln_prior is not False: - samples['ln_prior'] = ln_prior[full_samples_idx] - samples['ln_likelihood'] = all_marg_lls[good_samples_idx] + samples["ln_prior"] = ln_prior[full_samples_idx] + samples["ln_likelihood"] = all_marg_lls[good_samples_idx] return samples def ln_normal(x, mu, var): - return -0.5 * (np.log(2*np.pi * var) + (x - mu)**2 / var) + return -0.5 * (np.log(2 * np.pi * var) + (x - mu) ** 2 / var) diff --git a/thejoker/multiproc_helpers.py b/thejoker/multiproc_helpers.py index 12d146ba..1d242f6f 100644 --- a/thejoker/multiproc_helpers.py +++ b/thejoker/multiproc_helpers.py @@ -6,14 +6,25 @@ # Project from .logging import logger from .samples import JokerSamples -from .utils import (batch_tasks, read_batch, table_contains_column, - tempfile_decorator, NUMPY_LT_1_17) - - -def run_worker(worker, pool, prior_samples_file, task_args=(), n_batches=None, - n_prior_samples=None, samples_idx=None, random_state=None): - - with tb.open_file(prior_samples_file, mode='r') as f: +from .utils import ( + batch_tasks, + read_batch, + table_contains_column, + tempfile_decorator, +) + + +def run_worker( + worker, + pool, + prior_samples_file, + task_args=(), + n_batches=None, + n_prior_samples=None, + samples_idx=None, + rng=None, +): + with tb.open_file(prior_samples_file, mode="r") as f: n_samples = f.root[JokerSamples._hdf5_path].shape[0] if n_prior_samples is not None and samples_idx is not None: @@ -29,20 +40,18 @@ def run_worker(worker, pool, prior_samples_file, task_args=(), n_batches=None, n_batches = max(1, pool.size) if samples_idx is not None: - tasks = batch_tasks(n_samples, n_batches=n_batches, arr=samples_idx, - args=task_args) + tasks = batch_tasks( + n_samples, n_batches=n_batches, arr=samples_idx, args=task_args + ) else: tasks = batch_tasks(n_samples, n_batches=n_batches, args=task_args) - if random_state is not None: - if not NUMPY_LT_1_17: - from numpy.random import Generator, PCG64 - sg = random_state.bit_generator._seed_seq.spawn(len(tasks)) - for i in range(len(tasks)): - tasks[i] = tuple(tasks[i]) + (Generator(PCG64(sg[i])), ) - else: # TODO: remove when we drop numpy 1.16 support - for i in range(len(tasks)): - tasks[i] = tuple(tasks[i]) + (random_state, ) + if rng is not None: + from numpy.random import PCG64, Generator + + sg = rng.bit_generator._seed_seq.spawn(len(tasks)) + for i in range(len(tasks)): + tasks[i] = tuple(tasks[i]) + (Generator(PCG64(sg[i])),) results = [] for res in pool.map(worker, tasks): @@ -73,8 +82,12 @@ def marginal_ln_likelihood_worker(task): slice_or_idx, task_id, prior_samples_file, joker_helper = task # Read the batch of prior samples - batch = read_batch(prior_samples_file, joker_helper.packed_order, - slice_or_idx, units=joker_helper.internal_units) + batch = read_batch( + prior_samples_file, + joker_helper.packed_order, + slice_or_idx, + units=joker_helper.internal_units, + ) if batch.dtype != np.float64: batch = batch.astype(np.float64) @@ -86,87 +99,105 @@ def marginal_ln_likelihood_worker(task): @tempfile_decorator -def marginal_ln_likelihood_helper(joker_helper, prior_samples_file, pool, - n_batches=None, n_prior_samples=None, - samples_idx=None): - - task_args = (prior_samples_file, - joker_helper) - results = run_worker(marginal_ln_likelihood_worker, pool, - prior_samples_file, - task_args=task_args, n_batches=n_batches, - samples_idx=samples_idx, - n_prior_samples=n_prior_samples) +def marginal_ln_likelihood_helper( + joker_helper, + prior_samples_file, + pool, + n_batches=None, + n_prior_samples=None, + samples_idx=None, +): + task_args = (prior_samples_file, joker_helper) + results = run_worker( + marginal_ln_likelihood_worker, + pool, + prior_samples_file, + task_args=task_args, + n_batches=n_batches, + samples_idx=samples_idx, + n_prior_samples=n_prior_samples, + ) return np.concatenate(results) def make_full_samples_worker(task): - (slice_or_idx, - task_id, - prior_samples_file, - joker_helper, - n_linear_samples, - random_state) = task + ( + slice_or_idx, + task_id, + prior_samples_file, + joker_helper, + n_linear_samples, + rng, + ) = task # Read the batch of prior samples - batch = read_batch(prior_samples_file, - columns=joker_helper.packed_order, - slice_or_idx=slice_or_idx, - units=joker_helper.internal_units) - - # TODO: remove this when we drop numpy 1.16 support - if isinstance(random_state, np.random.RandomState): - tmp = np.random.RandomState() - tmp.set_state(random_state.get_state()) - tmp.seed(task_id) # TODO: is this safe? - random_state = tmp + batch = read_batch( + prior_samples_file, + columns=joker_helper.packed_order, + slice_or_idx=slice_or_idx, + units=joker_helper.internal_units, + ) if batch.dtype != np.float64: batch = batch.astype(np.float64) - raw_samples, _ = joker_helper.batch_get_posterior_samples(batch, - n_linear_samples, - random_state) + raw_samples, _ = joker_helper.batch_get_posterior_samples( + batch, n_linear_samples, rng + ) return raw_samples -def make_full_samples(joker_helper, prior_samples_file, pool, random_state, - samples_idx, n_linear_samples=1, n_batches=None): - - task_args = (prior_samples_file, - joker_helper, - n_linear_samples) - results = run_worker(make_full_samples_worker, pool, prior_samples_file, - task_args=task_args, n_batches=n_batches, - samples_idx=samples_idx, - random_state=random_state) +def make_full_samples( + joker_helper, + prior_samples_file, + pool, + rng, + samples_idx, + n_linear_samples=1, + n_batches=None, +): + task_args = (prior_samples_file, joker_helper, n_linear_samples) + results = run_worker( + make_full_samples_worker, + pool, + prior_samples_file, + task_args=task_args, + n_batches=n_batches, + samples_idx=samples_idx, + rng=rng, + ) # Concatenate all of the raw samples arrays raw_samples = np.concatenate(results) # unpack the raw samples - samples = JokerSamples.unpack(raw_samples, - joker_helper.internal_units, - t_ref=joker_helper.data.t_ref, - poly_trend=joker_helper.prior.poly_trend, - n_offsets=joker_helper.prior.n_offsets) + samples = JokerSamples.unpack( + raw_samples, + joker_helper.internal_units, + t_ref=joker_helper.data.t_ref, + poly_trend=joker_helper.prior.poly_trend, + n_offsets=joker_helper.prior.n_offsets, + ) return samples @tempfile_decorator -def rejection_sample_helper(joker_helper, prior_samples_file, pool, - random_state, - n_prior_samples=None, - max_posterior_samples=None, - n_linear_samples=1, - return_logprobs=False, - n_batches=None, - randomize_prior_order=False, - return_all_logprobs=False): - +def rejection_sample_helper( + joker_helper, + prior_samples_file, + pool, + rng, + n_prior_samples=None, + max_posterior_samples=None, + n_linear_samples=1, + return_logprobs=False, + n_batches=None, + randomize_prior_order=False, + return_all_logprobs=False, +): # Total number of samples in the cache: - with tb.open_file(prior_samples_file, mode='r') as f: + with tb.open_file(prior_samples_file, mode="r") as f: n_total_samples = f.root[JokerSamples._hdf5_path].shape[0] # TODO: pytables doesn't support variable length strings @@ -179,46 +210,50 @@ def rejection_sample_helper(joker_helper, prior_samples_file, pool, # "the prior samples.") # TODO: pytables doesn't support variable length strings - with h5py.File(prior_samples_file, mode='r') as f: + with h5py.File(prior_samples_file, mode="r") as f: if return_logprobs: - if not table_contains_column(f, 'ln_prior'): + if not table_contains_column(f, "ln_prior"): raise RuntimeError( "return_logprobs=True but ln_prior values not found in " "prior cache: make sure you generate prior samples with " "prior.sample (..., return_logprobs=True) before saving " - "the prior samples.") + "the prior samples." + ) if n_prior_samples is None: n_prior_samples = n_total_samples elif n_prior_samples > n_total_samples: - raise ValueError("Number of prior samples to use is greater than the " - "number of prior samples passed, or cached to a " - "filename specified. " - f"n_prior_samples={n_prior_samples} vs. " - f"n_total_samples={n_total_samples}") + raise ValueError( + "Number of prior samples to use is greater than the " + "number of prior samples passed, or cached to a " + "filename specified. " + f"n_prior_samples={n_prior_samples} vs. " + f"n_total_samples={n_total_samples}" + ) if max_posterior_samples is None: max_posterior_samples = n_prior_samples # Keyword arguments to be passed to marginal_ln_likelihood_helper: - ll_kw = dict(joker_helper=joker_helper, - prior_samples_file=prior_samples_file, - pool=pool, - n_batches=n_batches) + ll_kw = dict( + joker_helper=joker_helper, + prior_samples_file=prior_samples_file, + pool=pool, + n_batches=n_batches, + ) if randomize_prior_order: # Generate a random ordering for the samples - idx = random_state.choice(n_total_samples, size=n_prior_samples, - replace=False) - ll_kw['samples_idx'] = idx + idx = rng.choice(n_total_samples, size=n_prior_samples, replace=False) + ll_kw["samples_idx"] = idx else: - ll_kw['n_prior_samples'] = n_prior_samples + ll_kw["n_prior_samples"] = n_prior_samples # compute likelihoods lls = marginal_ln_likelihood_helper(**ll_kw) # get indices of samples that pass rejection step - uu = random_state.uniform(size=len(lls)) + uu = rng.uniform(size=len(lls)) good_samples_idx = np.where(np.exp(lls - lls.max()) > uu)[0] good_samples_idx = good_samples_idx[:max_posterior_samples] @@ -228,17 +263,22 @@ def rejection_sample_helper(joker_helper, prior_samples_file, pool, full_samples_idx = good_samples_idx # generate linear parameters - samples = make_full_samples(joker_helper, prior_samples_file, pool, - random_state, full_samples_idx, - n_linear_samples=n_linear_samples, - n_batches=n_batches) + samples = make_full_samples( + joker_helper, + prior_samples_file, + pool, + rng, + full_samples_idx, + n_linear_samples=n_linear_samples, + n_batches=n_batches, + ) if return_logprobs: - samples['ln_likelihood'] = lls[good_samples_idx] + samples["ln_likelihood"] = lls[good_samples_idx] - with tb.open_file(prior_samples_file, mode='r') as f: + with tb.open_file(prior_samples_file, mode="r") as f: data = f.root[JokerSamples._hdf5_path] - samples['ln_prior'] = data.read_coordinates(full_samples_idx) + samples["ln_prior"] = data.read_coordinates(full_samples_idx) if return_all_logprobs: return samples, lls @@ -247,19 +287,22 @@ def rejection_sample_helper(joker_helper, prior_samples_file, pool, @tempfile_decorator -def iterative_rejection_helper(joker_helper, prior_samples_file, pool, - random_state, - n_requested_samples, - init_batch_size=None, - growth_factor=128, - max_prior_samples=None, - n_linear_samples=1, - return_logprobs=False, - n_batches=None, - randomize_prior_order=False): - +def iterative_rejection_helper( + joker_helper, + prior_samples_file, + pool, + rng, + n_requested_samples, + init_batch_size=None, + growth_factor=128, + max_prior_samples=None, + n_linear_samples=1, + return_logprobs=False, + n_batches=None, + randomize_prior_order=False, +): # Total number of samples in the cache: - with tb.open_file(prior_samples_file, mode='r') as f: + with tb.open_file(prior_samples_file, mode="r") as f: n_total_samples = f.root[JokerSamples._hdf5_path].shape[0] # TODO: pytables doesn't support variable length strings @@ -272,14 +315,15 @@ def iterative_rejection_helper(joker_helper, prior_samples_file, pool, # "the prior samples.") # TODO: pytables doesn't support variable length strings - with h5py.File(prior_samples_file, mode='r') as f: + with h5py.File(prior_samples_file, mode="r") as f: if return_logprobs: - if not table_contains_column(f, 'ln_prior'): + if not table_contains_column(f, "ln_prior"): raise RuntimeError( "return_logprobs=True but ln_prior values not found in " "prior cache: make sure you generate prior samples with " "prior.sample (..., return_logprobs=True) before saving " - "the prior samples.") + "the prior samples." + ) if max_prior_samples is None: max_prior_samples = n_total_samples @@ -294,17 +338,18 @@ def iterative_rejection_helper(joker_helper, prior_samples_file, pool, n_process = init_batch_size if n_process > max_prior_samples: - raise ValueError("Prior sample library not big enough! For " - "iterative sampling, you have to have at least " - "growth_factor * n_requested_samples = " - f"{growth_factor * n_requested_samples} samples in " - "the prior samples cache file. You have, or have " - f"limited to, {max_prior_samples} samples.") + raise ValueError( + "Prior sample library not big enough! For " + "iterative sampling, you have to have at least " + "growth_factor * n_requested_samples = " + f"{growth_factor * n_requested_samples} samples in " + "the prior samples cache file. You have, or have " + f"limited to, {max_prior_samples} samples." + ) if randomize_prior_order: # Generate a random ordering for the samples - all_idx = random_state.choice(n_total_samples, size=max_prior_samples, - replace=False) + all_idx = rng.choice(n_total_samples, size=max_prior_samples, replace=False) else: all_idx = np.arange(0, max_prior_samples, 1) @@ -314,14 +359,18 @@ def iterative_rejection_helper(joker_helper, prior_samples_file, pool, logger.log(1, f"iteration {i}, computing {n_process} likelihoods") marg_lls = marginal_ln_likelihood_helper( - joker_helper, prior_samples_file, pool=pool, - n_batches=None, n_prior_samples=None, - samples_idx=all_idx[start_idx:start_idx + n_process]) + joker_helper, + prior_samples_file, + pool=pool, + n_batches=None, + n_prior_samples=None, + samples_idx=all_idx[start_idx : start_idx + n_process], + ) all_marg_lls = np.concatenate((all_marg_lls, marg_lls)) # get indices of samples that pass rejection step - uu = random_state.uniform(size=len(all_marg_lls)) + uu = rng.uniform(size=len(all_marg_lls)) aa = np.exp(all_marg_lls - all_marg_lls.max()) good_samples_idx = np.where(aa > uu)[0] @@ -355,18 +404,24 @@ def iterative_rejection_helper(joker_helper, prior_samples_file, pool, full_samples_idx = all_idx[good_samples_idx] # generate linear parameters - samples = make_full_samples(joker_helper, prior_samples_file, pool, - random_state, full_samples_idx, - n_linear_samples=n_linear_samples, - n_batches=n_batches) + samples = make_full_samples( + joker_helper, + prior_samples_file, + pool, + rng, + full_samples_idx, + n_linear_samples=n_linear_samples, + n_batches=n_batches, + ) # FIXME: copy-pasted from function above if return_logprobs: - samples['ln_likelihood'] = all_marg_lls[good_samples_idx] + samples["ln_likelihood"] = all_marg_lls[good_samples_idx] - with tb.open_file(prior_samples_file, mode='r') as f: + with tb.open_file(prior_samples_file, mode="r") as f: data = f.root[JokerSamples._hdf5_path] - samples['ln_prior'] = data.read_coordinates(full_samples_idx, - field='ln_prior') + samples["ln_prior"] = data.read_coordinates( + full_samples_idx, field="ln_prior" + ) return samples diff --git a/thejoker/prior.py b/thejoker/prior.py index a1d7c401..71f7b024 100644 --- a/thejoker/prior.py +++ b/thejoker/prior.py @@ -2,17 +2,20 @@ import astropy.units as u import numpy as np from aesara_theano_fallback.graph import fg +from astropy.utils.decorators import deprecated_renamed_argument # Project from .logging import logger -from .prior_helpers import (get_nonlinear_equiv_units, - get_linear_equiv_units, - validate_poly_trend, - get_v0_offsets_equiv_units, - validate_sigma_v) -from .utils import random_state_context +from .prior_helpers import ( + get_linear_equiv_units, + get_nonlinear_equiv_units, + get_v0_offsets_equiv_units, + validate_poly_trend, + validate_sigma_v, +) +from .utils import rng_context -__all__ = ['JokerPrior'] +__all__ = ["JokerPrior"] def _validate_model(model): @@ -28,14 +31,14 @@ def _validate_model(model): model = pm.Model() if not isinstance(model, pm.Model): - raise TypeError("Input model must be a pymc3.Model instance, not " - f"a {type(model)}") + raise TypeError( + "Input model must be a pymc3.Model instance, not " f"a {type(model)}" + ) return model class JokerPrior: - def __init__(self, pars=None, poly_trend=1, v0_offsets=None, model=None): """ This class controls the prior probability distributions for the @@ -69,8 +72,8 @@ def __init__(self, pars=None, poly_trend=1, v0_offsets=None, model=None): """ import aesara_theano_fallback.tensor as tt - import pymc3 as pm import exoplanet.units as xu + import pymc3 as pm self.model = _validate_model(model) @@ -92,10 +95,12 @@ def __init__(self, pars=None, poly_trend=1, v0_offsets=None, model=None): try: pars = {p.name: p for p in pars} except Exception: - raise ValueError("Invalid input parameters: The input " - "`pars` must either be a dictionary, " - "list, or a single pymc3 variable, not a " - "'{}'.".format(type(pars))) + raise ValueError( + "Invalid input parameters: The input " + "`pars` must either be a dictionary, " + "list, or a single pymc3 variable, not a " + f"'{type(pars)}'." + ) # Set the number of polynomial trend parameters self.poly_trend, self._v_trend_names = validate_poly_trend(poly_trend) @@ -107,9 +112,11 @@ def __init__(self, pars=None, poly_trend=1, v0_offsets=None, model=None): try: v0_offsets = list(v0_offsets) except Exception: - raise TypeError("Constant velocity offsets must be an iterable " - "of pymc3 variables that define the priors on " - "each offset term.") + raise TypeError( + "Constant velocity offsets must be an iterable " + "of pymc3 variables that define the priors on " + "each offset term." + ) self.v0_offsets = v0_offsets pars.update({p.name: p for p in self.v0_offsets}) @@ -120,49 +127,67 @@ def __init__(self, pars=None, poly_trend=1, v0_offsets=None, model=None): # equivalent to these self._nonlinear_equiv_units = get_nonlinear_equiv_units() self._linear_equiv_units = get_linear_equiv_units(self.poly_trend) - self._v0_offsets_equiv_units = get_v0_offsets_equiv_units( - self.n_offsets) - self._all_par_unit_equiv = {**self._nonlinear_equiv_units, - **self._linear_equiv_units, - **self._v0_offsets_equiv_units} + self._v0_offsets_equiv_units = get_v0_offsets_equiv_units(self.n_offsets) + self._all_par_unit_equiv = { + **self._nonlinear_equiv_units, + **self._linear_equiv_units, + **self._v0_offsets_equiv_units, + } # At this point, pars must be a dictionary: validate that all # parameters are specified and that they all have units for name in self.par_names: if name not in pars: - raise ValueError(f"Missing prior for parameter '{name}': " - "you must specify a prior distribution for " - "all parameters.") + raise ValueError( + f"Missing prior for parameter '{name}': " + "you must specify a prior distribution for " + "all parameters." + ) if not hasattr(pars[name], xu.UNIT_ATTR_NAME): - raise ValueError(f"Parameter '{name}' does not have associated " - "units: Use exoplanet.units to specify units " - "for your pymc3 variables. See the " - "documentation for examples: thejoker.rtfd.io") + raise ValueError( + f"Parameter '{name}' does not have associated " + "units: Use exoplanet.units to specify units " + "for your pymc3 variables. See the " + "documentation for examples: thejoker.rtfd.io" + ) equiv_unit = self._all_par_unit_equiv[name] - if not getattr(pars[name], - xu.UNIT_ATTR_NAME).is_equivalent(equiv_unit): - raise ValueError(f"Parameter '{name}' has an invalid unit: " - f"The units for this parameter must be " - f"transformable to '{equiv_unit}'") + if not getattr(pars[name], xu.UNIT_ATTR_NAME).is_equivalent(equiv_unit): + raise ValueError( + f"Parameter '{name}' has an invalid unit: " + f"The units for this parameter must be " + f"transformable to '{equiv_unit}'" + ) # Enforce that the priors on all linear parameters are Normal (or a # subclass of Normal) - for name in (list(self._linear_equiv_units.keys()) - + list(self._v0_offsets_equiv_units.keys())): + for name in list(self._linear_equiv_units.keys()) + list( + self._v0_offsets_equiv_units.keys() + ): if not isinstance(pars[name].distribution, pm.Normal): - raise ValueError("Priors on the linear parameters (K, v0, " - "etc.) must be independent Normal " - "distributions, not '{}'" - .format(type(pars[name].distribution))) + raise ValueError( + "Priors on the linear parameters (K, v0, " + "etc.) must be independent Normal " + f"distributions, not '{type(pars[name].distribution)}'" + ) self.pars = pars @classmethod - def default(cls, P_min=None, P_max=None, sigma_K0=None, P0=1*u.year, - sigma_v=None, s=None, poly_trend=1, v0_offsets=None, - model=None, pars=None): + def default( + cls, + P_min=None, + P_max=None, + sigma_K0=None, + P0=1 * u.year, + sigma_v=None, + s=None, + poly_trend=1, + v0_offsets=None, + model=None, + pars=None, + ): r""" An alternative initializer to set up the default prior for The Joker. @@ -226,29 +251,36 @@ def default(cls, P_min=None, P_max=None, sigma_K0=None, P0=1*u.year, model = _validate_model(model) - nl_pars = default_nonlinear_prior(P_min, P_max, s=s, - model=model, pars=pars) - l_pars = default_linear_prior(sigma_K0=sigma_K0, P0=P0, sigma_v=sigma_v, - poly_trend=poly_trend, model=model, - pars=pars) + nl_pars = default_nonlinear_prior(P_min, P_max, s=s, model=model, pars=pars) + l_pars = default_linear_prior( + sigma_K0=sigma_K0, + P0=P0, + sigma_v=sigma_v, + poly_trend=poly_trend, + model=model, + pars=pars, + ) pars = {**nl_pars, **l_pars} - obj = cls(pars=pars, model=model, poly_trend=poly_trend, - v0_offsets=v0_offsets) + obj = cls(pars=pars, model=model, poly_trend=poly_trend, v0_offsets=v0_offsets) return obj @property def par_names(self): - return (list(self._nonlinear_equiv_units.keys()) - + list(self._linear_equiv_units.keys()) - + list(self._v0_offsets_equiv_units)) + return ( + list(self._nonlinear_equiv_units.keys()) + + list(self._linear_equiv_units.keys()) + + list(self._v0_offsets_equiv_units) + ) @property def par_units(self): import exoplanet.units as xu - return {p.name: getattr(p, xu.UNIT_ATTR_NAME, u.one) - for _, p in self.pars.items()} + + return { + p.name: getattr(p, xu.UNIT_ATTR_NAME, u.one) for _, p in self.pars.items() + } @property def n_offsets(self): @@ -260,8 +292,18 @@ def __repr__(self): def __str__(self): return ", ".join(self.par_names) - def sample(self, size=1, generate_linear=False, return_logprobs=False, - random_state=None, dtype=None, **kwargs): + @deprecated_renamed_argument( + "random_state", "rng", since="v1.3", warning_type=DeprecationWarning + ) + def sample( + self, + size=1, + generate_linear=False, + return_logprobs=False, + rng=None, + dtype=None, + **kwargs, + ): """ Generate random samples from the prior. @@ -291,18 +333,23 @@ def sample(self, size=1, generate_linear=False, return_logprobs=False, The random samples. """ - from .samples import JokerSamples - from pymc3.distributions import draw_values import exoplanet.units as xu + from pymc3.distributions import draw_values + + from .samples import JokerSamples if dtype is None: dtype = np.float64 - sub_pars = {k: p for k, p in self.pars.items() - if k in self._nonlinear_equiv_units - or ((k in self._linear_equiv_units - or k in self._v0_offsets_equiv_units) - and generate_linear)} + sub_pars = { + k: p + for k, p in self.pars.items() + if k in self._nonlinear_equiv_units + or ( + (k in self._linear_equiv_units or k in self._v0_offsets_equiv_units) + and generate_linear + ) + } if generate_linear: par_names = self.par_names @@ -312,19 +359,19 @@ def sample(self, size=1, generate_linear=False, return_logprobs=False, # MAJOR HACK RELATED TO UPSTREAM ISSUES WITH pymc3: init_shapes = dict() for name, par in sub_pars.items(): - if hasattr(par, 'distribution'): + if hasattr(par, "distribution"): init_shapes[name] = par.distribution.shape - par.distribution.shape = (size, ) + par.distribution.shape = (size,) par_names = list(sub_pars.keys()) par_list = [sub_pars[k] for k in par_names] - with random_state_context(random_state): + with rng_context(rng): samples_values = draw_values(par_list) - raw_samples = {name: samples.astype(dtype) - for name, p, samples in zip(par_names, - par_list, - samples_values)} + raw_samples = { + name: samples.astype(dtype) + for name, p, samples in zip(par_names, par_list, samples_values) + } if return_logprobs: logp = [] @@ -332,15 +379,19 @@ def sample(self, size=1, generate_linear=False, return_logprobs=False, try: _logp = par.distribution.logp(raw_samples[par.name]).eval() except AttributeError: - logger.warning("Cannot auto-compute log-prior value for " - f"parameter {par} because it is defined " - "as a transformation from another " - "variable.") + logger.warning( + "Cannot auto-compute log-prior value for " + f"parameter {par} because it is defined " + "as a transformation from another " + "variable." + ) continue except fg.MissingInputError: - logger.warning("Cannot auto-compute log-prior value for " - f"parameter {par} because it depends on " - "other variables.") + logger.warning( + "Cannot auto-compute log-prior value for " + f"parameter {par} because it depends on " + "other variables." + ) continue logp.append(_logp) @@ -348,13 +399,13 @@ def sample(self, size=1, generate_linear=False, return_logprobs=False, # CONTINUED MAJOR HACK RELATED TO UPSTREAM ISSUES WITH pymc3: for name, par in sub_pars.items(): - if hasattr(par, 'distribution'): + if hasattr(par, "distribution"): par.distribution.shape = init_shapes[name] # Apply units if they are specified: - prior_samples = JokerSamples(poly_trend=self.poly_trend, - n_offsets=self.n_offsets, - **kwargs) + prior_samples = JokerSamples( + poly_trend=self.poly_trend, n_offsets=self.n_offsets, **kwargs + ) for name in par_names: p = sub_pars[name] unit = getattr(p, xu.UNIT_ATTR_NAME, u.one) @@ -365,7 +416,7 @@ def sample(self, size=1, generate_linear=False, return_logprobs=False, prior_samples[name] = np.atleast_1d(raw_samples[name]) * unit if return_logprobs: - prior_samples['ln_prior'] = log_prior + prior_samples["ln_prior"] = log_prior # TODO: right now, elsewhere, we assume the log_prior is a single value # for each sample (i.e. the total prior value). In principle, we could @@ -379,8 +430,7 @@ def sample(self, size=1, generate_linear=False, return_logprobs=False, @u.quantity_input(P_min=u.day, P_max=u.day) -def default_nonlinear_prior(P_min=None, P_max=None, s=None, - model=None, pars=None): +def default_nonlinear_prior(P_min=None, P_max=None, s=None, model=None, pars=None): r""" Retrieve pymc3 variables that specify the default prior on the nonlinear parameters of The Joker. See docstring of `JokerPrior.default()` for more @@ -408,12 +458,14 @@ def default_nonlinear_prior(P_min=None, P_max=None, s=None, """ import aesara_theano_fallback.tensor as tt import pymc3 as pm + try: from pymc3_ext.distributions import Angle except ImportError: from exoplanet.distributions import Angle import exoplanet.units as xu - from .distributions import UniformLog, Kipping13Global + + from .distributions import Kipping13Global, UniformLog model = pm.modelcontext(model) @@ -421,12 +473,12 @@ def default_nonlinear_prior(P_min=None, P_max=None, s=None, pars = dict() if s is None: - s = 0 * u.m/u.s + s = 0 * u.m / u.s if isinstance(s, pm.model.TensorVariable): - pars['s'] = pars.get('s', s) + pars["s"] = pars.get("s", s) else: - if not hasattr(s, 'unit') or not s.unit.is_equivalent(u.km/u.s): + if not hasattr(s, "unit") or not s.unit.is_equivalent(u.km / u.s): raise u.UnitsError("Invalid unit for s: must be equivalent to km/s") # dictionary of parameters to return @@ -438,31 +490,31 @@ def default_nonlinear_prior(P_min=None, P_max=None, s=None, # Note: we have to do it this way (as opposed to with .get(..., default) # because this can only get executed if the param is not already # defined, otherwise variables are defined twice in the model - if 'e' not in pars: - out_pars['e'] = xu.with_unit(Kipping13Global('e'), - u.one) + if "e" not in pars: + out_pars["e"] = xu.with_unit(Kipping13Global("e"), u.one) # If either omega or M0 is specified by user, default to U(0,2Ï€) - if 'omega' not in pars: - out_pars['omega'] = xu.with_unit(Angle('omega'), u.rad) + if "omega" not in pars: + out_pars["omega"] = xu.with_unit(Angle("omega"), u.rad) - if 'M0' not in pars: - out_pars['M0'] = xu.with_unit(Angle('M0'), u.rad) + if "M0" not in pars: + out_pars["M0"] = xu.with_unit(Angle("M0"), u.rad) - if 's' not in pars: - out_pars['s'] = xu.with_unit(pm.Deterministic('s', - tt.constant(s.value)), - s.unit) + if "s" not in pars: + out_pars["s"] = xu.with_unit( + pm.Deterministic("s", tt.constant(s.value)), s.unit + ) - if 'P' not in pars: + if "P" not in pars: if P_min is None or P_max is None: - raise ValueError("If you are using the default period prior, " - "you must pass in both P_min and P_max to set " - "the period prior domain.") - out_pars['P'] = xu.with_unit(UniformLog('P', - P_min.value, - P_max.to_value(P_min.unit)), - P_min.unit) + raise ValueError( + "If you are using the default period prior, " + "you must pass in both P_min and P_max to set " + "the period prior domain." + ) + out_pars["P"] = xu.with_unit( + UniformLog("P", P_min.value, P_max.to_value(P_min.unit)), P_min.unit + ) for k in pars.keys(): out_pars[k] = pars[k] @@ -470,9 +522,10 @@ def default_nonlinear_prior(P_min=None, P_max=None, s=None, return out_pars -@u.quantity_input(sigma_K0=u.km/u.s, P0=u.day) -def default_linear_prior(sigma_K0=None, P0=None, sigma_v=None, - poly_trend=1, model=None, pars=None): +@u.quantity_input(sigma_K0=u.km / u.s, P0=u.day) +def default_linear_prior( + sigma_K0=None, P0=None, sigma_v=None, poly_trend=1, model=None, pars=None +): r""" Retrieve pymc3 variables that specify the default prior on the linear parameters of The Joker. See docstring of `JokerPrior.default()` for more @@ -494,8 +547,9 @@ def default_linear_prior(sigma_K0=None, P0=None, sigma_v=None, This is either required, or this function must be called within a pymc3 model context. """ - import pymc3 as pm import exoplanet.units as xu + import pymc3 as pm + from .distributions import FixedCompanionMass model = pm.modelcontext(model) @@ -510,40 +564,42 @@ def default_linear_prior(sigma_K0=None, P0=None, sigma_v=None, poly_trend, v_names = validate_poly_trend(poly_trend) # get period/ecc from dict of nonlinear parameters - P = model.named_vars.get('P', None) - e = model.named_vars.get('e', None) + P = model.named_vars.get("P", None) + e = model.named_vars.get("e", None) if P is None or e is None: - raise ValueError("Period P and eccentricity e must both be defined as " - "nonlinear parameters on the model.") + raise ValueError( + "Period P and eccentricity e must both be defined as " + "nonlinear parameters on the model." + ) - if v_names and 'v0' not in pars: + if v_names and "v0" not in pars: sigma_v = validate_sigma_v(sigma_v, poly_trend, v_names) with model: - if 'K' not in pars: + if "K" not in pars: if sigma_K0 is None or P0 is None: - raise ValueError("If using the default prior form on K, you " - "must pass in a variance scale (sigma_K0) " - "and a reference period (P0)") + raise ValueError( + "If using the default prior form on K, you " + "must pass in a variance scale (sigma_K0) " + "and a reference period (P0)" + ) # Default prior on semi-amplitude: scales with period and # eccentricity such that it is flat with companion mass v_unit = sigma_K0.unit - out_pars['K'] = xu.with_unit(FixedCompanionMass('K', P=P, e=e, - sigma_K0=sigma_K0, - P0=P0), - v_unit) + out_pars["K"] = xu.with_unit( + FixedCompanionMass("K", P=P, e=e, sigma_K0=sigma_K0, P0=P0), v_unit + ) else: - v_unit = getattr(pars['K'], xu.UNIT_ATTR_NAME, u.one) + v_unit = getattr(pars["K"], xu.UNIT_ATTR_NAME, u.one) for i, name in enumerate(v_names): if name not in pars: # Default priors are independent gaussians # FIXME: make mean, mu_v, customizable out_pars[name] = xu.with_unit( - pm.Normal(name, 0., - sigma_v[name].value), - sigma_v[name].unit) + pm.Normal(name, 0.0, sigma_v[name].value), sigma_v[name].unit + ) for k in pars.keys(): out_pars[k] = pars[k] diff --git a/thejoker/src/fast_likelihood.pyx b/thejoker/src/fast_likelihood.pyx index c8e65848..2fc367c0 100644 --- a/thejoker/src/fast_likelihood.pyx +++ b/thejoker/src/fast_likelihood.pyx @@ -463,7 +463,7 @@ cdef class CJokerHelper: cpdef batch_get_posterior_samples(self, double[:, ::1] chunk, int n_linear_samples_per, - object random_state): + object rng): """TODO: Parameters @@ -519,7 +519,7 @@ cdef class CJokerHelper: # TODO: FIXME: this calls back to numpy at the Python layer # - use https://github.com/bashtage/randomgen instead? # a and Ainv are populated by the likelihood_worker() - linear_pars = random_state.multivariate_normal( + linear_pars = rng.multivariate_normal( self.a, np.linalg.inv(self.Ainv), size=n_linear_samples_per) for j in range(n_linear_samples_per): diff --git a/thejoker/tests/test_sampler.py b/thejoker/tests/test_sampler.py index f1b24521..48446742 100644 --- a/thejoker/tests/test_sampler.py +++ b/thejoker/tests/test_sampler.py @@ -3,48 +3,54 @@ # Third-party import astropy.units as u -from astropy.time import Time import numpy as np import pytest -from schwimmbad import SerialPool, MultiPool +from astropy.time import Time +from schwimmbad import MultiPool, SerialPool from twobody import KeplerOrbit +from ..data import RVData + # Package from ..prior import JokerPrior -from ..data import RVData from ..thejoker import TheJoker from .test_prior import get_prior -from ..utils import DEFAULT_RNG, NUMPY_LT_1_17 -def make_data(n_times=8, random_state=None, v1=None, K=None): - if random_state is None: - random_state = DEFAULT_RNG() - rnd = random_state +def make_data(n_times=8, rng=None, v1=None, K=None): + if rng is None: + rng = np.random.default_rng() P = 51.8239 * u.day if K is None: - K = 54.2473 * u.km/u.s - v0 = 31.48502 * u.km/u.s - EPOCH = Time('J2000') - t = Time('J2000') + P * np.sort(rnd.uniform(0, 3., n_times)) + K = 54.2473 * u.km / u.s + v0 = 31.48502 * u.km / u.s + EPOCH = Time("J2000") + t = Time("J2000") + P * np.sort(rng.uniform(0, 3.0, n_times)) # binary system - random parameters - orbit = KeplerOrbit(P=P, K=K, e=0.3, omega=0.283*u.radian, - M0=2.592*u.radian, t0=EPOCH, - i=90*u.deg, Omega=0*u.deg) # these don't matter + orbit = KeplerOrbit( + P=P, + K=K, + e=0.3, + omega=0.283 * u.radian, + M0=2.592 * u.radian, + t0=EPOCH, + i=90 * u.deg, + Omega=0 * u.deg, + ) # these don't matter rv = orbit.radial_velocity(t) + v0 if v1 is not None: rv = rv + v1 * (t - EPOCH).jd * u.day - err = np.full_like(rv.value, 0.5) * u.km/u.s + err = np.full_like(rv.value, 0.5) * u.km / u.s data = RVData(t, rv, rv_err=err) return data, orbit -@pytest.mark.parametrize('case', range(get_prior())) +@pytest.mark.parametrize("case", range(get_prior())) def test_init(case): prior, _ = get_prior(case) @@ -52,7 +58,7 @@ def test_init(case): TheJoker(prior) with pytest.raises(TypeError): - TheJoker('jsdfkj') + TheJoker("jsdfkj") # Pools: with SerialPool() as pool: @@ -60,27 +66,22 @@ def test_init(case): # fail when pool is invalid: with pytest.raises(TypeError): - TheJoker(prior, pool='sdfks') + TheJoker(prior, pool="sdfks") # Random state: - rnd = DEFAULT_RNG(42) - TheJoker(prior, random_state=rnd) + rnd = np.random.default_rng(42) + TheJoker(prior, rng=rng) # fail when random state is invalid: with pytest.raises(TypeError): - TheJoker(prior, random_state='sdfks') - - if not NUMPY_LT_1_17: - with pytest.warns(DeprecationWarning): - rnd = np.random.RandomState(42) - TheJoker(prior, random_state=rnd) + TheJoker(prior, rng="sdfks") # tempfile location: - joker = TheJoker(prior, tempfile_path='/tmp/joker') + joker = TheJoker(prior, tempfile_path="/tmp/joker") assert os.path.exists(joker.tempfile_path) -@pytest.mark.parametrize('case', range(get_prior())) +@pytest.mark.parametrize("case", range(get_prior())) def test_marginal_ln_likelihood(tmpdir, case): prior, _ = get_prior(case) @@ -93,15 +94,14 @@ def test_marginal_ln_likelihood(tmpdir, case): assert len(ll) == len(prior_samples) # save prior samples to a file and pass that instead - filename = str(tmpdir / 'samples.hdf5') + filename = str(tmpdir / "samples.hdf5") prior_samples.write(filename, overwrite=True) ll = joker.marginal_ln_likelihood(data, filename) assert len(ll) == len(prior_samples) # make sure batches work: - ll = joker.marginal_ln_likelihood(data, filename, - n_batches=10) + ll = joker.marginal_ln_likelihood(data, filename, n_batches=10) assert len(ll) == len(prior_samples) # NOTE: this makes it so I can't parallelize tests, I think @@ -111,19 +111,30 @@ def test_marginal_ln_likelihood(tmpdir, case): assert len(ll) == len(prior_samples) -@pytest.mark.parametrize('prior', [ - JokerPrior.default(P_min=5*u.day, P_max=500*u.day, - sigma_K0=25*u.km/u.s, sigma_v=100*u.km/u.s), - JokerPrior.default(P_min=5*u.day, P_max=500*u.day, - sigma_K0=25*u.km/u.s, poly_trend=2, - sigma_v=[100*u.km/u.s, 0.5*u.km/u.s/u.day]) -]) +@pytest.mark.parametrize( + "prior", + [ + JokerPrior.default( + P_min=5 * u.day, + P_max=500 * u.day, + sigma_K0=25 * u.km / u.s, + sigma_v=100 * u.km / u.s, + ), + JokerPrior.default( + P_min=5 * u.day, + P_max=500 * u.day, + sigma_K0=25 * u.km / u.s, + poly_trend=2, + sigma_v=[100 * u.km / u.s, 0.5 * u.km / u.s / u.day], + ), + ], +) def test_rejection_sample(tmpdir, prior): data, orbit = make_data() - flat_data, orbit = make_data(K=0.1*u.m/u.s) + flat_data, orbit = make_data(K=0.1 * u.m / u.s) prior_samples = prior.sample(size=16384, return_logprobs=True) - filename = str(tmpdir / 'samples.hdf5') + filename = str(tmpdir / "samples.hdf5") prior_samples.write(filename, overwrite=True) joker = TheJoker(prior) @@ -146,61 +157,74 @@ def test_rejection_sample(tmpdir, prior): lls = samples.ln_unmarginalized_likelihood(flat_data) assert np.isfinite(lls).all() - samples, lls = joker.rejection_sample(flat_data, prior_samples, - return_all_logprobs=True) + samples, lls = joker.rejection_sample( + flat_data, prior_samples, return_all_logprobs=True + ) assert len(lls) == len(prior_samples) # Check that setting the random state makes it deterministic all_Ps = [] all_Ks = [] for i in range(10): - joker = TheJoker(prior, random_state=DEFAULT_RNG(42)) + joker = TheJoker(prior, rng=np.random.default_rng(42)) samples = joker.rejection_sample(flat_data, prior_samples) - all_Ps.append(samples['P']) - all_Ks.append(samples['K']) + all_Ps.append(samples["P"]) + all_Ks.append(samples["K"]) for i in range(1, len(all_Ps)): assert u.allclose(all_Ps[0], all_Ps[i]) assert u.allclose(all_Ks[0], all_Ks[i]) -@pytest.mark.parametrize('prior', [ - JokerPrior.default(P_min=5*u.day, P_max=500*u.day, - sigma_K0=25*u.km/u.s, sigma_v=100*u.km/u.s), - JokerPrior.default(P_min=5*u.day, P_max=500*u.day, - sigma_K0=25*u.km/u.s, poly_trend=2, - sigma_v=[100*u.km/u.s, 0.5*u.km/u.s/u.day]) -]) +@pytest.mark.parametrize( + "prior", + [ + JokerPrior.default( + P_min=5 * u.day, + P_max=500 * u.day, + sigma_K0=25 * u.km / u.s, + sigma_v=100 * u.km / u.s, + ), + JokerPrior.default( + P_min=5 * u.day, + P_max=500 * u.day, + sigma_K0=25 * u.km / u.s, + poly_trend=2, + sigma_v=[100 * u.km / u.s, 0.5 * u.km / u.s / u.day], + ), + ], +) def test_iterative_rejection_sample(tmpdir, prior): data, orbit = make_data(n_times=3) prior_samples = prior.sample(size=10_000, return_logprobs=True) - filename = str(tmpdir / 'samples.hdf5') + filename = str(tmpdir / "samples.hdf5") prior_samples.write(filename, overwrite=True) - joker = TheJoker(prior, random_state=DEFAULT_RNG(42)) + joker = TheJoker(prior, rng=np.random.default_rng(42)) for _samples in [prior_samples, filename]: # pass JokerSamples instance, process all samples: - samples = joker.iterative_rejection_sample(data, _samples, - n_requested_samples=4) + samples = joker.iterative_rejection_sample( + data, _samples, n_requested_samples=4 + ) assert len(samples) > 1 - samples = joker.iterative_rejection_sample(data, _samples, - n_requested_samples=4, - return_logprobs=True) + samples = joker.iterative_rejection_sample( + data, _samples, n_requested_samples=4, return_logprobs=True + ) assert len(samples) > 1 # Check that setting the random state makes it deterministic all_Ps = [] all_Ks = [] for i in range(10): - joker = TheJoker(prior, random_state=DEFAULT_RNG(42)) - samples = joker.iterative_rejection_sample(data, prior_samples, - n_requested_samples=4, - randomize_prior_order=True) - all_Ps.append(samples['P']) - all_Ks.append(samples['K']) + joker = TheJoker(prior, rng=np.random.default_rng(42)) + samples = joker.iterative_rejection_sample( + data, prior_samples, n_requested_samples=4, randomize_prior_order=True + ) + all_Ps.append(samples["P"]) + all_Ks.append(samples["K"]) for i in range(1, len(all_Ps)): assert u.allclose(all_Ps[0], all_Ps[i]) diff --git a/thejoker/thejoker.py b/thejoker/thejoker.py index be45b3d3..72771f2b 100644 --- a/thejoker/thejoker.py +++ b/thejoker/thejoker.py @@ -2,21 +2,23 @@ import os import warnings -# Third-party import numpy as np -# Project -from .logging import logger +# Third-party +from astropy.utils.decorators import deprecated_renamed_argument + from .data_helpers import validate_prepare_data -from .samples_analysis import MAP_sample, is_P_unimodal from .likelihood_helpers import get_trend_design_matrix -from .prior_helpers import validate_n_offsets, validate_poly_trend + +# Project +from .logging import logger from .prior import JokerPrior, _validate_model -from .src.fast_likelihood import CJokerHelper +from .prior_helpers import validate_n_offsets, validate_poly_trend from .samples import JokerSamples -from .utils import DEFAULT_RNG, NUMPY_LT_1_17 +from .samples_analysis import is_P_unimodal +from .src.fast_likelihood import CJokerHelper -__all__ = ['TheJoker'] +__all__ = ["TheJoker"] class TheJoker: @@ -30,7 +32,7 @@ class TheJoker: parameters used in The Joker. pool : `schwimmbad.BasePool` (optional) A processing pool (default is a `schwimmbad.SerialPool` instance). - random_state : `numpy.random.Generator` (optional) + rng : `numpy.random.Generator` (optional) A `numpy.random.Generator` instance for controlling random number generation. tempfile_path : str (optional) @@ -40,41 +42,27 @@ class TheJoker: Default: ``~/.thejoker`` """ - def __init__(self, prior, pool=None, random_state=None, - tempfile_path=None): - + @deprecated_renamed_argument( + "random_state", "rng", since="v1.3", warning_type=DeprecationWarning + ) + def __init__(self, prior, pool=None, rng=None, tempfile_path=None): # set the processing pool if pool is None: import schwimmbad + pool = schwimmbad.SerialPool() - elif not hasattr(pool, 'map') or not hasattr(pool, 'close'): - raise TypeError("Input pool object must have .map() and .close() " - "methods. We recommend using `schwimmbad` pools.") + elif not hasattr(pool, "map") or not hasattr(pool, "close"): + raise TypeError( + "Input pool object must have .map() and .close() " + "methods. We recommend using `schwimmbad` pools." + ) self.pool = pool # Set the parent random state - child processes get different states # based on the parent - if random_state is None: - random_state = DEFAULT_RNG() - elif (isinstance(random_state, np.random.RandomState) and - not NUMPY_LT_1_17): - warnings.warn("With thejoker>=v1.2, use numpy.random.Generator " - "objects instead of RandomState objects to control " - "random numbers.", DeprecationWarning) - tmp = np.random.Generator(np.random.MT19937()) - tmp.bit_generator.state = random_state.get_state() - random_state = tmp - elif (not NUMPY_LT_1_17 and - not isinstance(random_state, np.random.Generator)): - raise TypeError("Random state object must be a " - "numpy.random.Generator instance, not " - f"'{type(random_state)}'") - elif (NUMPY_LT_1_17 and - not isinstance(random_state, np.random.RandomState)): - raise TypeError("Random state object must be a " - "numpy.random.RandomState instance, not " - f"'{type(random_state)}'") - self.random_state = random_state + if rng is None: + rng = np.random.default_rng() + self.rng = rng # check if a JokerParams instance was passed in to specify the state if not isinstance(prior, JokerPrior): @@ -82,10 +70,9 @@ def __init__(self, prior, pool=None, random_state=None, self.prior = prior if tempfile_path is None: - self._tempfile_path = os.path.expanduser('~/.thejoker/') + self._tempfile_path = os.path.expanduser("~/.thejoker/") else: - self._tempfile_path = os.path.abspath( - os.path.expanduser(tempfile_path)) + self._tempfile_path = os.path.abspath(os.path.expanduser(tempfile_path)) @property def tempfile_path(self): @@ -93,14 +80,15 @@ def tempfile_path(self): return self._tempfile_path def _make_joker_helper(self, data): - all_data, ids, trend_M = validate_prepare_data(data, - self.prior.poly_trend, - self.prior.n_offsets) + all_data, ids, trend_M = validate_prepare_data( + data, self.prior.poly_trend, self.prior.n_offsets + ) joker_helper = CJokerHelper(all_data, self.prior, trend_M) return joker_helper - def marginal_ln_likelihood(self, data, prior_samples, n_batches=None, - in_memory=False): + def marginal_ln_likelihood( + self, data, prior_samples, n_batches=None, in_memory=False + ): """ Compute the marginal log-likelihood at each of the input prior samples. @@ -129,32 +117,36 @@ def marginal_ln_likelihood(self, data, prior_samples, n_batches=None, The marginal log-likelihood computed at the location of each prior sample. """ - from .multiproc_helpers import marginal_ln_likelihood_helper from .likelihood_helpers import marginal_ln_likelihood_inmem + from .multiproc_helpers import marginal_ln_likelihood_helper joker_helper = self._make_joker_helper(data) # also validates data if in_memory: if isinstance(prior_samples, JokerSamples): prior_samples, _ = prior_samples.pack( - units=joker_helper.internal_units, - names=joker_helper.packed_order) + units=joker_helper.internal_units, names=joker_helper.packed_order + ) return marginal_ln_likelihood_inmem(joker_helper, prior_samples) else: - return marginal_ln_likelihood_helper(joker_helper, prior_samples, - pool=self.pool, - n_batches=n_batches) - - def rejection_sample(self, data, prior_samples, - n_prior_samples=None, - max_posterior_samples=None, - n_linear_samples=1, - return_logprobs=False, - return_all_logprobs=False, - n_batches=None, - randomize_prior_order=False, - in_memory=False): + return marginal_ln_likelihood_helper( + joker_helper, prior_samples, pool=self.pool, n_batches=n_batches + ) + + def rejection_sample( + self, + data, + prior_samples, + n_prior_samples=None, + max_posterior_samples=None, + n_linear_samples=1, + return_logprobs=False, + return_all_logprobs=False, + n_batches=None, + randomize_prior_order=False, + in_memory=False, + ): """ Run The Joker's rejection sampling on prior samples to get posterior samples for the input data. @@ -212,65 +204,69 @@ def rejection_sample(self, data, prior_samples, The posterior samples produced from The Joker. """ - from .multiproc_helpers import rejection_sample_helper from .likelihood_helpers import rejection_sample_inmem + from .multiproc_helpers import rejection_sample_helper joker_helper = self._make_joker_helper(data) # also validates data if isinstance(prior_samples, int): # If an integer, generate that many prior samples first N = prior_samples - prior_samples = self.prior.sample(size=N, - return_logprobs=return_logprobs) + prior_samples = self.prior.sample(size=N, return_logprobs=return_logprobs) if in_memory: if isinstance(prior_samples, JokerSamples): ln_prior = None if return_logprobs: - ln_prior = prior_samples['ln_prior'] + ln_prior = prior_samples["ln_prior"] prior_samples, _ = prior_samples.pack( - units=joker_helper.internal_units, - names=joker_helper.packed_order) + units=joker_helper.internal_units, names=joker_helper.packed_order + ) else: ln_prior = return_logprobs samples = rejection_sample_inmem( joker_helper, prior_samples, - random_state=self.random_state, + rng=self.rng, ln_prior=ln_prior, max_posterior_samples=max_posterior_samples, n_linear_samples=n_linear_samples, - return_all_logprobs=return_all_logprobs) + return_all_logprobs=return_all_logprobs, + ) else: samples = rejection_sample_helper( joker_helper, prior_samples, pool=self.pool, - random_state=self.random_state, + rng=self.rng, n_prior_samples=n_prior_samples, max_posterior_samples=max_posterior_samples, n_linear_samples=n_linear_samples, return_logprobs=return_logprobs, n_batches=n_batches, randomize_prior_order=randomize_prior_order, - return_all_logprobs=return_all_logprobs) + return_all_logprobs=return_all_logprobs, + ) return samples - def iterative_rejection_sample(self, data, prior_samples, - n_requested_samples, - max_prior_samples=None, - n_linear_samples=1, - return_logprobs=False, - n_batches=None, - randomize_prior_order=False, - init_batch_size=None, - growth_factor=128, - in_memory=False): - + def iterative_rejection_sample( + self, + data, + prior_samples, + n_requested_samples, + max_prior_samples=None, + n_linear_samples=1, + return_logprobs=False, + n_batches=None, + randomize_prior_order=False, + init_batch_size=None, + growth_factor=128, + in_memory=False, + ): """This is an experimental sampling method that adaptively generates posterior samples given a large library of prior samples. The advantage of this function over the standard ``rejection_sample`` method is that @@ -324,8 +320,8 @@ def iterative_rejection_sample(self, data, prior_samples, samples : `~thejoker.JokerSamples` The posterior samples produced from The Joker. """ - from .multiproc_helpers import iterative_rejection_helper from .likelihood_helpers import iterative_rejection_inmem + from .multiproc_helpers import iterative_rejection_helper joker_helper = self._make_joker_helper(data) # also validates data @@ -333,23 +329,24 @@ def iterative_rejection_sample(self, data, prior_samples, if isinstance(prior_samples, JokerSamples): ln_prior = None if return_logprobs: - ln_prior = prior_samples['ln_prior'] + ln_prior = prior_samples["ln_prior"] prior_samples, _ = prior_samples.pack( - units=joker_helper.internal_units, - names=joker_helper.packed_order) + units=joker_helper.internal_units, names=joker_helper.packed_order + ) else: ln_prior = return_logprobs samples = iterative_rejection_inmem( joker_helper, prior_samples, - random_state=self.random_state, + rng=self.rng, n_requested_samples=n_requested_samples, ln_prior=ln_prior, init_batch_size=init_batch_size, growth_factor=growth_factor, - n_linear_samples=n_linear_samples) + n_linear_samples=n_linear_samples, + ) else: samples = iterative_rejection_helper( @@ -358,13 +355,14 @@ def iterative_rejection_sample(self, data, prior_samples, init_batch_size=init_batch_size, growth_factor=growth_factor, pool=self.pool, - random_state=self.random_state, + rng=self.rng, n_requested_samples=n_requested_samples, max_prior_samples=max_prior_samples, n_linear_samples=n_linear_samples, return_logprobs=return_logprobs, n_batches=n_batches, - randomize_prior_order=randomize_prior_order) + randomize_prior_order=randomize_prior_order, + ) return samples @@ -392,25 +390,27 @@ def setup_mcmc(self, data, joker_samples, model=None, custom_func=None): mcmc_init : dict """ - import pymc3 as pm + import aesara_theano_fallback.tensor as tt import exoplanet as xo import exoplanet.units as xu - import aesara_theano_fallback.tensor as tt + import pymc3 as pm model = _validate_model(model) # Reduce data, strip units: - data, ids, _ = validate_prepare_data(data, - self.prior.poly_trend, - self.prior.n_offsets) + data, ids, _ = validate_prepare_data( + data, self.prior.poly_trend, self.prior.n_offsets + ) x = data._t_bmjd - data._t_ref_bmjd y = data.rv.value err = data.rv_err.to_value(data.rv.unit) # First, prepare the joker_samples: if not isinstance(joker_samples, JokerSamples): - raise TypeError("You must pass in a JokerSamples instance to the " - "joker_samples argument.") + raise TypeError( + "You must pass in a JokerSamples instance to the " + "joker_samples argument." + ) if len(joker_samples) > 1: # check if unimodal in P, if not, warn @@ -432,20 +432,20 @@ def setup_mcmc(self, data, joker_samples, model=None, custom_func=None): p = self.prior.pars - if 't_peri' not in model.named_vars: + if "t_peri" not in model.named_vars: with model: - pm.Deterministic('t_peri', p['P'] * p['M0'] / (2*np.pi)) + pm.Deterministic("t_peri", p["P"] * p["M0"] / (2 * np.pi)) - if 'obs' in model.named_vars: + if "obs" in model.named_vars: return mcmc_init with model: # Set up the orbit model orbit = xo.orbits.KeplerianOrbit( - period=p['P'], - ecc=p['e'], - omega=p['omega'], - t_periastron=model.named_vars['t_peri'] + period=p["P"], + ecc=p["e"], + omega=p["omega"], + t_periastron=model.named_vars["t_peri"], ) # design matrix @@ -456,28 +456,28 @@ def setup_mcmc(self, data, joker_samples, model=None, custom_func=None): _, vtrend_names = validate_poly_trend(self.prior.poly_trend) with model: - v_pars = ([p['v0']] - + [p[name] for name in offset_names] - + [p[name] for name in vtrend_names[1:]]) # skip v0 + v_pars = ( + [p["v0"]] + + [p[name] for name in offset_names] + + [p[name] for name in vtrend_names[1:]] + ) # skip v0 v_trend_vec = tt.stack(v_pars, axis=0) trend = tt.dot(M, v_trend_vec) - rv_model = orbit.get_radial_velocity(x, K=p['K']) + trend + rv_model = orbit.get_radial_velocity(x, K=p["K"]) + trend pm.Deterministic("model_rv", rv_model) - err = tt.sqrt(err**2 + p['s']**2) + err = tt.sqrt(err**2 + p["s"] ** 2) pm.Normal("obs", mu=rv_model, sd=err, observed=y) - pm.Deterministic('logp', model.logpt) + pm.Deterministic("logp", model.logpt) dist = pm.Normal.dist(model.model_rv, data.rv_err.value) lnlike = pm.Deterministic( - 'ln_likelihood', - dist.logp(data.rv.value).sum(axis=-1)) + "ln_likelihood", dist.logp(data.rv.value).sum(axis=-1) + ) - pm.Deterministic( - 'ln_prior', - model.logpt - lnlike) + pm.Deterministic("ln_prior", model.logpt - lnlike) return mcmc_init @@ -490,10 +490,11 @@ def trace_to_samples(self, trace, data, names=None): trace : `~pymc3.backends.base.MultiTrace` """ warnings.warn( - 'This method is deprecated: Use ' - 'thejoker.samples_helpers.trace_to_samples() instead', - UserWarning + "This method is deprecated: Use " + "thejoker.samples_helpers.trace_to_samples() instead", + UserWarning, ) from thejoker.samples_helpers import trace_to_samples + return trace_to_samples(self, trace, data, names) diff --git a/thejoker/utils.py b/thejoker/utils.py index 67b22f20..d1aaf542 100644 --- a/thejoker/utils.py +++ b/thejoker/utils.py @@ -1,38 +1,25 @@ """Miscellaneous utilities""" # Standard library -import inspect import contextlib -from functools import wraps +import inspect import os -from distutils.version import LooseVersion +from functools import wraps from tempfile import NamedTemporaryFile # Third-party import astropy.units as u -from astropy.table.meta import get_header_from_yaml -from astropy.io.misc.hdf5 import meta_path import h5py import numpy as np import tables as tb +from astropy.io.misc.hdf5 import meta_path +from astropy.table.meta import get_header_from_yaml +from astropy.utils.decorators import deprecated_renamed_argument # Package from .samples import JokerSamples -# TODO: remove this when we drop support for numpy 1.16 -# Numpy version: -NUMPY_LT_1_17 = LooseVersion(np.__version__) < '1.17' - -if NUMPY_LT_1_17: - DEFAULT_RNG = np.random.RandomState - integers = lambda obj, *args, **kwargs: obj.randint(*args, **kwargs) -else: - DEFAULT_RNG = np.random.default_rng - integers = lambda obj, *args, **kwargs: obj.integers(*args, **kwargs) - - -__all__ = ['batch_tasks', 'table_header_to_units', 'read_batch', - 'tempfile_decorator'] +__all__ = ["batch_tasks", "table_header_to_units", "read_batch", "tempfile_decorator"] def batch_tasks(n_tasks, n_batches, arr=None, args=None, start_idx=0): @@ -80,10 +67,10 @@ def batch_tasks(n_tasks, n_batches, arr=None, args=None, start_idx=0): else: if arr is None: # store indices - tasks.append([(start_idx, n_tasks+start_idx), start_idx] + args) + tasks.append([(start_idx, n_tasks + start_idx), start_idx] + args) else: # store sliced array - tasks.append([arr[start_idx:n_tasks+start_idx], start_idx] + args) + tasks.append([arr[start_idx : n_tasks + start_idx], start_idx] + args) return tasks @@ -94,12 +81,11 @@ def table_header_to_units(header_dataset): column name to unit. """ - header = get_header_from_yaml(h.decode('utf-8') - for h in header_dataset) + header = get_header_from_yaml(h.decode("utf-8") for h in header_dataset) units = dict() - for row in header['datatype']: - units[row['name']] = u.Unit(row.get('unit', u.one)) + for row in header["datatype"]: + units[row["name"]] = u.Unit(row.get("unit", u.one)) return units @@ -108,17 +94,19 @@ def table_contains_column(root, column): from .samples import JokerSamples path = meta_path(JokerSamples._hdf5_path) - header = get_header_from_yaml(h.decode('utf-8') for h in root[path]) + header = get_header_from_yaml(h.decode("utf-8") for h in root[path]) columns = [] - for row in header['datatype']: - columns.append(row['name']) + for row in header["datatype"]: + columns.append(row["name"]) return column in columns -def read_batch(prior_samples_file, columns, slice_or_idx, units=None, - random_state=None): +@deprecated_renamed_argument( + "random_state", "rng", since="v1.3", warning_type=DeprecationWarning +) +def read_batch(prior_samples_file, columns, slice_or_idx, units=None, rng=None): """ Single-point interface to all read_batch functions below that infers the type of read from the ``slice_or_idx`` argument. @@ -137,7 +125,7 @@ def read_batch(prior_samples_file, columns, slice_or_idx, units=None, interpreted as the row indices to read. units : dict (optional) The desired output units to convert the prior samples to. - random_state : `numpy.random.RandomState` + rng : `numpy.random.Generator` Used to generate random row indices if passing an integer (size) in to ``slice_or_idx``. This argument is ignored for other read methods. @@ -150,28 +138,32 @@ def read_batch(prior_samples_file, columns, slice_or_idx, units=None, """ if isinstance(slice_or_idx, tuple): # read a contiguous batch of prior samples - batch = read_batch(prior_samples_file, columns, - slice(*slice_or_idx), units=units) + batch = read_batch( + prior_samples_file, columns, slice(*slice_or_idx), units=units + ) elif isinstance(slice_or_idx, slice): # read a contiguous batch of prior samples - batch = read_batch_slice(prior_samples_file, columns, - slice_or_idx, units=units) + batch = read_batch_slice(prior_samples_file, columns, slice_or_idx, units=units) elif isinstance(slice_or_idx, int): # read a random batch of samples of size "slice_or_idx" - batch = read_random_batch(prior_samples_file, columns, - slice_or_idx, units=units, - random_state=random_state) + batch = read_random_batch( + prior_samples_file, + columns, + slice_or_idx, + units=units, + rng=rng, + ) elif isinstance(slice_or_idx, np.ndarray): # read a random batch of samples of size "slice_or_idx" - batch = read_batch_idx(prior_samples_file, columns, - slice_or_idx, units=units) + batch = read_batch_idx(prior_samples_file, columns, slice_or_idx, units=units) else: - raise ValueError("Invalid input for slice_or_idx: must be a slice, " - "int, or numpy array.") + raise ValueError( + "Invalid input for slice_or_idx: must be a slice, " "int, or numpy array." + ) return batch @@ -187,15 +179,13 @@ def read_batch_slice(prior_samples_file, columns, slice, units=None): # We have to do this with h5py because current (2021-02-05) versions of # pytables don't support variable length strings, which h5py is using to # serialize units in the astropy table metadata - with h5py.File(prior_samples_file, mode='r') as f: + with h5py.File(prior_samples_file, mode="r") as f: table_units = table_header_to_units(f[meta_path(path)]) batch = None - with tb.open_file(prior_samples_file, mode='r') as f: - + with tb.open_file(prior_samples_file, mode="r") as f: for i, name in enumerate(columns): - arr = f.root[path].read(slice.start, slice.stop, slice.step, - field=name) + arr = f.root[path].read(slice.start, slice.stop, slice.step, field=name) if batch is None: batch = np.zeros((len(arr), len(columns)), dtype=arr.dtype) batch[:, i] = arr @@ -220,11 +210,11 @@ def read_batch_idx(prior_samples_file, columns, idx, units=None): # We have to do this with h5py because current (2021-02-05) versions of # pytables don't support variable length strings, which h5py is using to # serialize units in the astropy table metadata - with h5py.File(prior_samples_file, mode='r') as f: + with h5py.File(prior_samples_file, mode="r") as f: table_units = table_header_to_units(f[meta_path(path)]) batch = np.zeros((len(idx), len(columns))) - with tb.open_file(prior_samples_file, mode='r') as f: + with tb.open_file(prior_samples_file, mode="r") as f: for i, name in enumerate(columns): batch[:, i] = f.root[path].read_coordinates(idx, field=name) @@ -238,19 +228,18 @@ def read_batch_idx(prior_samples_file, columns, idx, units=None): return batch -def read_random_batch(prior_samples_file, columns, size, units=None, - random_state=None): +def read_random_batch(prior_samples_file, columns, size, units=None, rng=None): """ Read a random batch (row block) of prior samples into a plain numpy array, converting units where necessary. """ - if random_state is None: - random_state = DEFAULT_RNG() + if rng is None: + rng = np.random.default_rng() path = JokerSamples._hdf5_path - with tb.open_file(prior_samples_file, mode='r') as f: - idx = integers(random_state, 0, f.root[path].shape[0], size=size) + with tb.open_file(prior_samples_file, mode="r") as f: + idx = rng.integers(rng, 0, f.root[path].shape[0], size=size) return read_batch_idx(prior_samples_file, columns, idx=idx, units=units) @@ -258,48 +247,51 @@ def read_random_batch(prior_samples_file, columns, size, units=None, def tempfile_decorator(func): wrapped_signature = inspect.signature(func) func_args = list(wrapped_signature.parameters.keys()) - if 'prior_samples_file' not in func_args: - raise ValueError("Cant decorate function because it doesn't contain an " - "argument called 'prior_samples_file'") + if "prior_samples_file" not in func_args: + raise ValueError( + "Cant decorate function because it doesn't contain an " + "argument called 'prior_samples_file'" + ) @wraps(func) def wrapper(*args, **kwargs): args = list(args) - if 'prior_samples_file' in kwargs: - prior_samples = kwargs['prior_samples_file'] + if "prior_samples_file" in kwargs: + prior_samples = kwargs["prior_samples_file"] else: - prior_samples = args.pop(func_args.index('prior_samples_file')) + prior_samples = args.pop(func_args.index("prior_samples_file")) - in_memory = kwargs.get('in_memory', False) + in_memory = kwargs.get("in_memory", False) if not isinstance(prior_samples, str) and not in_memory: if not isinstance(prior_samples, JokerSamples): - raise TypeError("prior_samples_file must either be a string " - "filename specifying a cache file contining " - "prior samples, or must be a JokerSamples " - f"instance, not: {type(prior_samples)}") + raise TypeError( + "prior_samples_file must either be a string " + "filename specifying a cache file contining " + "prior samples, or must be a JokerSamples " + f"instance, not: {type(prior_samples)}" + ) # This is required (instead of a context) because the named file # can't be opened a second time on Windows...see, e.g., # https://github.com/Kotaimen/awscfncli/issues/93 - f = NamedTemporaryFile(mode='r+', suffix='.hdf5', delete=False) + f = NamedTemporaryFile(mode="r+", suffix=".hdf5", delete=False) f.close() try: # write samples to tempfile and recursively call this method prior_samples.write(f.name, overwrite=True) - kwargs['prior_samples_file'] = f.name + kwargs["prior_samples_file"] = f.name func_return = func(*args, **kwargs) except Exception as e: raise e finally: os.unlink(f.name) - else: # FIXME: it's a string, so it's probably a filename, but we should # validate the contents of the file! - kwargs['prior_samples_file'] = prior_samples + kwargs["prior_samples_file"] = prior_samples func_return = func(*args, **kwargs) return func_return @@ -308,12 +300,12 @@ def wrapper(*args, **kwargs): @contextlib.contextmanager -def random_state_context(random_state): - state = np.random.get_state() - if random_state is not None: - np.random.seed(integers(random_state, 2**32-1)) # HACK +def rng_context(rng): + pre_state = np.random.get_bit_generator() + if rng is not None: + np.random.set_bit_generator(rng.bit_generator) try: yield finally: - if random_state is not None: - np.random.set_state(state) + if rng is not None: + np.random.set_bit_generator(pre_state) From 102ae34cd3bbba263c7417e4c217aa75112353a7 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Sun, 3 Mar 2024 12:24:10 -0500 Subject: [PATCH 11/50] wowowow got it working with pymc v5 --- pyproject.toml | 20 ++- setup.py | 29 ++++ thejoker/_keplerian_orbit.py | 276 +++++++++++++++++++++++++++++++ thejoker/distributions.py | 158 ++++++++++-------- thejoker/prior.py | 142 +++++++--------- thejoker/samples_helpers.py | 210 +++++++++++++---------- thejoker/src/fast_likelihood.pyx | 19 ++- thejoker/src/setup_package.py | 17 +- thejoker/thejoker.py | 44 ++--- thejoker/units.py | 52 ++++++ thejoker/utils.py | 9 +- 11 files changed, 684 insertions(+), 292 deletions(-) create mode 100644 setup.py create mode 100644 thejoker/_keplerian_orbit.py create mode 100644 thejoker/units.py diff --git a/pyproject.toml b/pyproject.toml index d8cc7d9a..3b390c6b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,6 @@ requires = [ "setuptools>=64", "setuptools_scm>=8", "wheel", - "extension-helpers==1.*", "numpy", "scipy", "cython", @@ -16,19 +15,19 @@ name = "thejoker" authors = [{name = "Adrian Price-Whelan", email = "adrianmpw@gmail.com"}] description = "A custom Monte Carlo sampler for the two-body problem." readme = "README.rst" -requires-python = ">=3.7,<3.10" +requires-python = ">=3.9" license.file = "LICENSE" dynamic = ["version"] dependencies = [ - "astropy<6.0", + "astropy", "numpy", "twobody>=0.9", "scipy", "h5py", "schwimmbad>=0.3.1", "pymc", - "pymc_ext", - "exoplanet>=0.2.2", + "pymc_ext @ git+https://github.com/exoplanet-dev/pymc-ext", + "exoplanet-core[pymc]", "tables", ] @@ -56,12 +55,17 @@ docs = [ "ipykernel" ] +[tool.setuptools.packages.find] +where = ["."] +include = ["thejoker", "thejoker.*"] + +[tool.setuptools.package-data] +"*" = ["*.c"] +"thejoker.src" = ["*.pyx", "*.pxd", "*.h"] + [tool.setuptools_scm] version_file = "thejoker/_version.py" -[tool.extension-helpers] -use_extension_helpers = true - [tool.pytest.ini_options] testpaths = ["thejoker", "docs"] doctest_plus = "enabled" diff --git a/setup.py b/setup.py new file mode 100644 index 00000000..1190b3ae --- /dev/null +++ b/setup.py @@ -0,0 +1,29 @@ +import os +from collections import defaultdict + +from setuptools import Extension, setup + +exts = [] + +import numpy as np + +try: + import twobody +except ImportError: + msg = "The twobody package is required to install TheJoker. " + raise ImportError(msg) + +cfg = defaultdict(list) +cfg["include_dirs"].append(np.get_include()) + +twobody_path = os.path.dirname(twobody.__file__) +cfg["include_dirs"].append(twobody_path) +cfg["sources"].append(os.path.join(twobody_path, "src/twobody.c")) + +cfg["extra_compile_args"].append("--std=gnu99") +cfg["sources"].append("thejoker/src/fast_likelihood.pyx") +exts.append(Extension("thejoker.src.fast_likelihood", **cfg)) + +setup( + ext_modules=exts, +) diff --git a/thejoker/_keplerian_orbit.py b/thejoker/_keplerian_orbit.py new file mode 100644 index 00000000..be0c2bab --- /dev/null +++ b/thejoker/_keplerian_orbit.py @@ -0,0 +1,276 @@ +"""A port from exoplanet, to support pymc recent versions""" + +__all__ = ["KeplerianOrbit"] + + +import numpy as np +import pytensor.tensor as pt +from astropy import units as u +from exoplanet_core.pymc import ops + +tt = pt + + +class KeplerianOrbit: + """A system of bodies on Keplerian orbits around a common central + + Given the input parameters, the values of all other parameters will be + computed so a ``KeplerianOrbit`` instance will always have attributes for + each argument. Note that the units of the computed attributes will all be + in the standard units of this class (``R_sun``, ``M_sun``, and ``days``) + except for ``rho_star`` which will be in ``g / cm^3``. + + There are only specific combinations of input parameters that can be used: + + 1. First, either ``period`` or ``a`` must be given. If values are given + for both parameters, then neither ``m_star`` or ``rho_star`` can be + defined because the stellar density implied by each planet will be + computed in ``rho_star``. + 2. Only one of ``incl`` and ``b`` can be given. + 3. If a value is given for ``ecc`` then ``omega`` must also be given. + 4. If no stellar parameters are given, the central body is assumed to be + the sun. If only ``rho_star`` is defined, the radius of the central is + assumed to be ``1 * R_sun``. Otherwise, at most two of ``m_star``, + ``r_star``, and ``rho_star`` can be defined. + 5. Either ``t0`` (reference transit) or ``t_periastron`` must be given, + but not both. + + + Args: + period: The orbital periods of the bodies in days. + a: The semimajor axes of the orbits in ``R_sun``. + t0: The time of a reference transit for each orbits in days. + t_periastron: The epoch of a reference periastron passage in days. + incl: The inclinations of the orbits in radians. + b: The impact parameters of the orbits. + ecc: The eccentricities of the orbits. Must be ``0 <= ecc < 1``. + omega: The arguments of periastron for the orbits in radians. + Omega: The position angles of the ascending nodes in radians. + m_planet: The masses of the planets in units of ``m_planet_units``. + m_star: The mass of the star in ``M_sun``. + r_star: The radius of the star in ``R_sun``. + rho_star: The density of the star in units of ``rho_star_units``. + + """ + + def __init__( + self, + period=None, + a=None, + t0=None, + t_periastron=None, + ecc=None, + omega=None, + sin_omega=None, + cos_omega=None, + Omega=None, + model=None, + **kwargs, + ): + self.a = a + self.period = period + self.m_planet = tt.zeros_like(period) + # self.m_total = self.m_star + self.m_planet + + self.n = 2 * np.pi / self.period + self.K0 = self.n * self.a / self.m_total + + if Omega is None: + self.Omega = None + else: + self.Omega = pt.as_tensor_variable(Omega) + self.cos_Omega = tt.cos(self.Omega) + self.sin_Omega = tt.sin(self.Omega) + + # Eccentricity + if ecc is None: + self.ecc = None + self.M0 = 0.5 * np.pi + tt.zeros_like(self.n) + else: + self.ecc = pt.as_tensor_variable(ecc) + if omega is not None: + if sin_omega is not None and cos_omega is not None: + raise ValueError( + "either 'omega' or 'sin_omega' and 'cos_omega' can be " + "provided" + ) + self.omega = pt.as_tensor_variable(omega) + self.cos_omega = tt.cos(self.omega) + self.sin_omega = tt.sin(self.omega) + elif sin_omega is not None and cos_omega is not None: + self.cos_omega = pt.as_tensor_variable(cos_omega) + self.sin_omega = pt.as_tensor_variable(sin_omega) + self.omega = tt.arctan2(self.sin_omega, self.cos_omega) + + else: + raise ValueError("both e and omega must be provided") + + opsw = 1 + self.sin_omega + E0 = 2 * tt.arctan2( + tt.sqrt(1 - self.ecc) * self.cos_omega, + tt.sqrt(1 + self.ecc) * opsw, + ) + self.M0 = E0 - self.ecc * tt.sin(E0) + + ome2 = 1 - self.ecc**2 + self.K0 /= tt.sqrt(ome2) + + zla = tt.zeros_like(self.P) + self.incl = 0.5 * np.pi + zla + self.cos_incl = zla + self.b = zla + + if t0 is not None and t_periastron is not None: + raise ValueError("you can't define both t0 and t_periastron") + if t0 is None and t_periastron is None: + t0 = tt.zeros_like(self.period) + + if t0 is None: + self.t_periastron = pt.as_tensor_variable(t_periastron) + self.t0 = self.t_periastron + self.M0 / self.n + else: + self.t0 = pt.as_tensor_variable(t0) + self.t_periastron = self.t0 - self.M0 / self.n + + self.tref = self.t_periastron - self.t0 + + self.sin_incl = tt.sin(self.incl) + + def _rotate_vector(self, x, y): + """Apply the rotation matrices to go from orbit to observer frame + + In order, + 1. rotate about the z axis by an amount omega -> x1, y1, z1 + 2. rotate about the x1 axis by an amount -incl -> x2, y2, z2 + 3. rotate about the z2 axis by an amount Omega -> x3, y3, z3 + + Args: + x: A tensor representing the x coodinate in the plane of the + orbit. + y: A tensor representing the y coodinate in the plane of the + orbit. + + Returns: + Three tensors representing ``(X, Y, Z)`` in the observer frame. + + """ + + # 1) rotate about z0 axis by omega + if self.ecc is None: + x1 = x + y1 = y + else: + x1 = self.cos_omega * x - self.sin_omega * y + y1 = self.sin_omega * x + self.cos_omega * y + + # 2) rotate about x1 axis by -incl + x2 = x1 + y2 = self.cos_incl * y1 + # z3 = z2, subsequent rotation by Omega doesn't affect it + Z = -self.sin_incl * y1 + + # 3) rotate about z2 axis by Omega + if self.Omega is None: + return (x2, y2, Z) + + X = self.cos_Omega * x2 - self.sin_Omega * y2 + Y = self.sin_Omega * x2 + self.cos_Omega * y2 + return X, Y, Z + + def _warp_times(self, t, _pad=True): + if _pad: + return tt.shape_padright(t) - self.t0 + return t - self.t0 + + def _get_true_anomaly(self, t, _pad=True): + M = (self._warp_times(t, _pad=_pad) - self.tref) * self.n + if self.ecc is None: + return tt.sin(M), tt.cos(M) + sinf, cosf = ops.kepler(M, self.ecc + tt.zeros_like(M)) + return sinf, cosf + + def _get_velocity(self, m, t): + """Get the velocity vector of a body in the observer frame""" + sinf, cosf = self._get_true_anomaly(t) + K = self.K0 * m + if self.ecc is None: + return self._rotate_vector(-K * sinf, K * cosf) + return self._rotate_vector(-K * sinf, K * (cosf + self.ecc)) + + def get_star_velocity(self, t): + """Get the star's velocity vector + + .. note:: For a system with multiple planets, this will return one + column per planet with the contributions from each planet. The + total velocity can be found by summing along the last axis. + + Args: + t: The times where the velocity should be evaluated. + + Returns: + The components of the velocity vector at ``t`` in units of + ``M_sun/day``. + + """ + return tuple(tt.squeeze(x) for x in self._get_velocity(self.m_planet, t)) + + def get_radial_velocity(self, t, K=None, output_units=None): + """Get the radial velocity of the star + + .. note:: The convention in exoplanet is that positive `z` points + *towards* the observer. However, for consistency with radial + velocity literature this method returns values where positive + radial velocity corresponds to a redshift as expected. + + Args: + t: The times where the radial velocity should be evaluated. + K (Optional): The semi-amplitudes of the orbits. If provided, the + ``m_planet`` and ``incl`` parameters will be ignored and this + amplitude will be used instead. + output_units (Optional): An AstroPy velocity unit. If not given, + the output will be evaluated in ``m/s``. This is ignored if a + value is given for ``K``. + + Returns: + The reflex radial velocity evaluated at ``t`` in units of + ``output_units``. For multiple planets, this will have one row for + each planet. + + """ + + # Special case for K given: m_planet, incl, etc. is ignored + if K is not None: + sinf, cosf = self._get_true_anomaly(t) + if self.ecc is None: + return tt.squeeze(K * cosf) + # cos(w + f) + e * cos(w) from Lovis & Fischer + return tt.squeeze( + K + * ( + self.cos_omega * cosf + - self.sin_omega * sinf + + self.ecc * self.cos_omega + ) + ) + + # Compute the velocity using the full orbit solution + if output_units is None: + output_units = u.m / u.s + conv = (1 * u.R_sun / u.day).to(output_units).value + v = self.get_star_velocity(t) + return -conv * v[2] + + +def get_true_anomaly(M, e, **kwargs): + """Get the true anomaly for a tensor of mean anomalies and eccentricities + + Args: + M: The mean anomaly. + e: The eccentricity. This should have the same shape as ``M``. + + Returns: + The true anomaly of the orbit. + + """ + sinf, cosf = ops.kepler(M, e) + return tt.arctan2(sinf, cosf) diff --git a/thejoker/distributions.py b/thejoker/distributions.py index eb70ba67..07f22cf6 100644 --- a/thejoker/distributions.py +++ b/thejoker/distributions.py @@ -1,48 +1,60 @@ # Third-party import astropy.units as u import numpy as np -import pymc3 as pm -from pymc3.distributions import generate_samples -import aesara_theano_fallback.tensor as tt -import exoplanet.units as xu +import pymc as pm +import pytensor.tensor as pt +from pymc.distributions.dist_math import check_parameters +from pytensor.tensor.random.basic import NormalRV, RandomVariable -__all__ = ['UniformLog', 'FixedCompanionMass'] +from thejoker.units import UNIT_ATTR_NAME +__all__ = ["UniformLog", "FixedCompanionMass"] + + +class UniformLogRV(RandomVariable): + name = "uniformlog" + ndim_supp = 0 + ndims_params = [0, 0] + dtype = "floatX" + + @classmethod + def rng_fn(cls, rng, a, b, size): + _fac = np.log(b) - np.log(a) + uu = rng.uniform(size=size) + return np.exp(uu * _fac + np.log(a)) -class UniformLog(pm.Continuous): - def __init__(self, a, b, **kwargs): - """A distribution over a value, x, that is uniform in log(x) over the - domain :math:`(a, b)`. - """ - - self.a = float(a) - self.b = float(b) - assert (self.a > 0) and (self.b > 0) - self._fac = np.log(self.b) - np.log(self.a) - - shape = kwargs.get("shape", None) - if shape is None: - testval = 0.5 * (self.a + self.b) - else: - testval = 0.5 * (self.a + self.b) + np.zeros(shape) - kwargs["testval"] = kwargs.pop("testval", testval) - super(UniformLog, self).__init__(**kwargs) - - def _random(self, size=None): - uu = np.random.uniform(size=size) - return np.exp(uu * self._fac + np.log(self.a)) - - def random(self, point=None, size=None): - return generate_samples( - self._random, - dist_shape=self.shape, - broadcast_shape=self.shape, - size=size, +uniformlog = UniformLogRV() + + +class UniformLog(pm.Continuous): + rv_op = uniformlog + + @classmethod + def dist(cls, a, b, **kwargs): + a = pt.as_tensor_variable(a) + b = pt.as_tensor_variable(b) + return super().dist([a, b], **kwargs) + + def support_point(rv, size, a, b): + a, b = pt.broadcast_arrays(a, b) + return 0.5 * (a + b) + + def logp(value, a, b): + _fac = pt.log(b) - pt.log(a) + res = -pt.as_tensor_variable(value) - pt.log(_fac) + return check_parameters( + res, + (a > 0) & (a < b), + msg="a > 0 and a < b", ) - def logp(self, value): - return -tt.as_tensor_variable(value) - np.log(self._fac) + +class FixedCompanionMassRV(NormalRV): + _print_name = ("FixedCompanionMass", "\\mathcal{N}") + + +fixedcompanionmass = FixedCompanionMassRV() class FixedCompanionMass(pm.Normal): @@ -62,50 +74,60 @@ class FixedCompanionMass(pm.Normal): be specified. """ - @u.quantity_input(sigma_K0=u.km/u.s, P0=u.day, max_K=u.km/u.s) - def __init__(self, P, e, sigma_K0, P0, mu=0., max_K=500*u.km/u.s, - K_unit=None, **kwargs): - self._sigma_K0 = sigma_K0 - self._P0 = P0 - self._max_K = max_K - + rv_op = fixedcompanionmass + + @classmethod + @u.quantity_input(sigma_K0=u.km / u.s, P0=u.day, max_K=u.km / u.s) + def dist( + cls, + P, + e, + sigma_K0, + P0, + mu=0.0, + max_K=500 * u.km / u.s, + K_unit=None, + *args, + **kwargs, + ): if K_unit is not None: - self._sigma_K0 = self.sigma_K0.to(K_unit) - self._max_K = self._max_K.to(self._sigma_K0.unit) - - if hasattr(P, xu.UNIT_ATTR_NAME): - self._P0 = self._P0.to(getattr(P, xu.UNIT_ATTR_NAME)) + sigma_K0 = sigma_K0.to_value(K_unit) + max_K = max_K.to(sigma_K0.unit) - sigma_K0 = self._sigma_K0.value - P0 = self._P0.value + if hasattr(P, UNIT_ATTR_NAME): + P0 = P0.to(getattr(P, UNIT_ATTR_NAME)) - sigma = tt.min([self._max_K.value, - sigma_K0 * (P/P0)**(-1/3) / np.sqrt(1-e**2)]) - super().__init__(mu=mu, sigma=sigma) + sigma = pt.clip( + sigma_K0.value * (P / P0.value) ** (-1 / 3) / np.sqrt(1 - e**2), + 0.0, + max_K.value, + ) + dist = super().dist(mu=mu, sigma=sigma, *args, **kwargs) + dist._sigma_K0 = sigma_K0 + dist._max_K = max_K + dist._P0 = P0 + return dist class Kipping13Long(pm.Beta): + rv_op = pt.random.beta - def __init__(self): - r""" - The inferred long-period eccentricity distribution from Kipping (2013). - """ - super().__init__(1.12, 3.09) + @classmethod + def dist(cls, *args, **kwargs): + return super().dist(alpha=1.12, beta=3.09, *args, **kwargs) class Kipping13Short(pm.Beta): + rv_op = pt.random.beta - def __init__(self): - r""" - The inferred short-period eccentricity distribution from Kipping (2013). - """ - super().__init__(0.697, 3.27) + @classmethod + def dist(cls, *args, **kwargs): + return super().dist(alpha=0.697, beta=3.27, *args, **kwargs) class Kipping13Global(pm.Beta): + rv_op = pt.random.beta - def __init__(self): - r""" - The inferred global eccentricity distribution from Kipping (2013). - """ - super().__init__(0.867, 3.03) + @classmethod + def dist(cls, *args, **kwargs): + return super().dist(alpha=0.867, beta=3.03, *args, **kwargs) diff --git a/thejoker/prior.py b/thejoker/prior.py index 71f7b024..a10e8807 100644 --- a/thejoker/prior.py +++ b/thejoker/prior.py @@ -1,9 +1,12 @@ # Third-party import astropy.units as u import numpy as np -from aesara_theano_fallback.graph import fg +import pymc as pm +import pytensor.tensor as pt from astropy.utils.decorators import deprecated_renamed_argument +import thejoker.units as xu + # Project from .logging import logger from .prior_helpers import ( @@ -19,8 +22,6 @@ def _validate_model(model): - import pymc3 as pm - # validate input model if model is None: try: @@ -32,7 +33,7 @@ def _validate_model(model): if not isinstance(model, pm.Model): raise TypeError( - "Input model must be a pymc3.Model instance, not " f"a {type(model)}" + "Input model must be a pymc.Model instance, not " f"a {type(model)}" ) return model @@ -53,7 +54,7 @@ def __init__(self, pars=None, poly_trend=1, v0_offsets=None, model=None): Parameters ---------- pars : dict, list (optional) - Either a list of pymc3 variables, or a dictionary of variables with + Either a list of pymc variables, or a dictionary of variables with keys set to the variable names. If any of these variables are defined as deterministic transforms from other variables, see the next parameter below. @@ -66,15 +67,11 @@ def __init__(self, pars=None, poly_trend=1, v0_offsets=None, model=None): v0_offsets : list (optional) A list of additional Gaussian parameters that set systematic offsets of subsets of the data. TODO: link to tutorial here - model : `pymc3.Model` + model : `pymc.Model` This is either required, or this function must be called within a - pymc3 model context. + pymc model context. """ - import aesara_theano_fallback.tensor as tt - import exoplanet.units as xu - import pymc3 as pm - self.model = _validate_model(model) # Parse and clean up the input pars @@ -82,7 +79,7 @@ def __init__(self, pars=None, poly_trend=1, v0_offsets=None, model=None): pars = dict() pars.update(model.named_vars) - elif isinstance(pars, tt.TensorVariable): # a single variable + elif isinstance(pars, pt.TensorVariable): # a single variable # Note: this has to go before the next clause because # TensorVariable instances are iterable... pars = {pars.name: pars} @@ -98,7 +95,7 @@ def __init__(self, pars=None, poly_trend=1, v0_offsets=None, model=None): raise ValueError( "Invalid input parameters: The input " "`pars` must either be a dictionary, " - "list, or a single pymc3 variable, not a " + "list, or a single pymc variable, not a " f"'{type(pars)}'." ) @@ -114,7 +111,7 @@ def __init__(self, pars=None, poly_trend=1, v0_offsets=None, model=None): except Exception: raise TypeError( "Constant velocity offsets must be an iterable " - "of pymc3 variables that define the priors on " + "of pymc variables that define the priors on " "each offset term." ) @@ -138,39 +135,47 @@ def __init__(self, pars=None, poly_trend=1, v0_offsets=None, model=None): # parameters are specified and that they all have units for name in self.par_names: if name not in pars: - raise ValueError( - f"Missing prior for parameter '{name}': " - "you must specify a prior distribution for " - "all parameters." + msg = ( + f"Missing prior for parameter '{name}': you must specify a prior " + "distribution for all parameters." ) + raise ValueError(msg) if not hasattr(pars[name], xu.UNIT_ATTR_NAME): - raise ValueError( - f"Parameter '{name}' does not have associated " - "units: Use exoplanet.units to specify units " - "for your pymc3 variables. See the " + msg = ( + f"Parameter '{name}' does not have associated units: Use " + "thejoker.units to specify units for your pymc variables. See the " "documentation for examples: thejoker.rtfd.io" ) + raise ValueError(msg) equiv_unit = self._all_par_unit_equiv[name] if not getattr(pars[name], xu.UNIT_ATTR_NAME).is_equivalent(equiv_unit): - raise ValueError( - f"Parameter '{name}' has an invalid unit: " - f"The units for this parameter must be " - f"transformable to '{equiv_unit}'" + msg = ( + f"Parameter '{name}' has an invalid unit: The units for this " + f"parameter must be transformable to '{equiv_unit}'" ) + raise ValueError(msg) # Enforce that the priors on all linear parameters are Normal (or a # subclass of Normal) for name in list(self._linear_equiv_units.keys()) + list( self._v0_offsets_equiv_units.keys() ): - if not isinstance(pars[name].distribution, pm.Normal): - raise ValueError( - "Priors on the linear parameters (K, v0, " - "etc.) must be independent Normal " - f"distributions, not '{type(pars[name].distribution)}'" + p = pars[name] + if not hasattr(p, "owner"): + msg = f"Invalid type for prior on linear parameter {name}: {type(p)}" + raise TypeError(msg) + + if not isinstance( + p.owner.op, pt.random.op.RandomVariable + ) or p.owner.op._print_name[0] not in ["Normal", "FixedCompanionMass"]: + msg = ( + "Priors on the linear parameters (K, v0, etc.) must be independent " + f"Normal distributions, not '{p.owner.op._print_name[0]}' (for " + f"{name})" ) + raise ValueError(msg) self.pars = pars @@ -239,11 +244,11 @@ def default( v0_offsets : list (optional) A list of additional Gaussian parameters that set systematic offsets of subsets of the data. TODO: link to tutorial here - model : `pymc3.Model` (optional) + model : `pymc.Model` (optional) If not specified, this will create a model instance and store it on the prior object. pars : dict, list (optional) - Either a list of pymc3 variables, or a dictionary of variables with + Either a list of pymc variables, or a dictionary of variables with keys set to the variable names. If any of these variables are defined as deterministic transforms from other variables, see the next parameter below. @@ -276,8 +281,6 @@ def par_names(self): @property def par_units(self): - import exoplanet.units as xu - return { p.name: getattr(p, xu.UNIT_ATTR_NAME, u.one) for _, p in self.pars.items() } @@ -307,14 +310,6 @@ def sample( """ Generate random samples from the prior. - .. note:: - - Right now, generating samples with the prior values is slow (i.e. - with ``return_logprobs=True``) because of pymc3 issues (see - discussion here: - https://discourse.pymc.io/t/draw-values-speed-scaling-with-transformed-variables/4076). - This will hopefully be resolved in the future... - Parameters ---------- size : int (optional) @@ -333,9 +328,6 @@ def sample( The random samples. """ - import exoplanet.units as xu - from pymc3.distributions import draw_values - from .samples import JokerSamples if dtype is None: @@ -357,16 +349,16 @@ def sample( par_names = list(self._nonlinear_equiv_units.keys()) # MAJOR HACK RELATED TO UPSTREAM ISSUES WITH pymc3: - init_shapes = dict() - for name, par in sub_pars.items(): - if hasattr(par, "distribution"): - init_shapes[name] = par.distribution.shape - par.distribution.shape = (size,) + # init_shapes = dict() + # for name, par in sub_pars.items(): + # if hasattr(par, "distribution"): + # init_shapes[name] = par.distribution.shape + # par.distribution.shape = (size,) par_names = list(sub_pars.keys()) par_list = [sub_pars[k] for k in par_names] with rng_context(rng): - samples_values = draw_values(par_list) + samples_values = pm.draw(par_list, draws=size) raw_samples = { name: samples.astype(dtype) @@ -386,7 +378,7 @@ def sample( "variable." ) continue - except fg.MissingInputError: + except Exception: logger.warning( "Cannot auto-compute log-prior value for " f"parameter {par} because it depends on " @@ -398,9 +390,9 @@ def sample( log_prior = np.sum(logp, axis=0) # CONTINUED MAJOR HACK RELATED TO UPSTREAM ISSUES WITH pymc3: - for name, par in sub_pars.items(): - if hasattr(par, "distribution"): - par.distribution.shape = init_shapes[name] + # for name, par in sub_pars.items(): + # if hasattr(par, "distribution"): + # par.distribution.shape = init_shapes[name] # Apply units if they are specified: prior_samples = JokerSamples( @@ -432,7 +424,7 @@ def sample( @u.quantity_input(P_min=u.day, P_max=u.day) def default_nonlinear_prior(P_min=None, P_max=None, s=None, model=None, pars=None): r""" - Retrieve pymc3 variables that specify the default prior on the nonlinear + Retrieve pymc variables that specify the default prior on the nonlinear parameters of The Joker. See docstring of `JokerPrior.default()` for more information. @@ -449,21 +441,14 @@ def default_nonlinear_prior(P_min=None, P_max=None, s=None, model=None, pars=Non Parameters ---------- - P_min : `~astropy.units.Quantity` [time] - P_max : `~astropy.units.Quantity` [time] - s : `~pm.model.TensorVariable`, ~astropy.units.Quantity` [speed] - model : `pymc3.Model` - This is either required, or this function must be called within a pymc3 + P_min : `astropy.units.Quantity` [time] + P_max : `astropy.units.Quantity` [time] + s : `TensorVariable`, `astropy.units.Quantity` [speed] + model : `pymc.Model` + This is either required, or this function must be called within a pymc model context. """ - import aesara_theano_fallback.tensor as tt - import pymc3 as pm - - try: - from pymc3_ext.distributions import Angle - except ImportError: - from exoplanet.distributions import Angle - import exoplanet.units as xu + from pymc_ext.distributions import angle from .distributions import Kipping13Global, UniformLog @@ -475,7 +460,7 @@ def default_nonlinear_prior(P_min=None, P_max=None, s=None, model=None, pars=Non if s is None: s = 0 * u.m / u.s - if isinstance(s, pm.model.TensorVariable): + if isinstance(s, pt.TensorVariable): pars["s"] = pars.get("s", s) else: if not hasattr(s, "unit") or not s.unit.is_equivalent(u.km / u.s): @@ -495,14 +480,14 @@ def default_nonlinear_prior(P_min=None, P_max=None, s=None, model=None, pars=Non # If either omega or M0 is specified by user, default to U(0,2Ï€) if "omega" not in pars: - out_pars["omega"] = xu.with_unit(Angle("omega"), u.rad) + out_pars["omega"] = xu.with_unit(angle("omega"), u.rad) if "M0" not in pars: - out_pars["M0"] = xu.with_unit(Angle("M0"), u.rad) + out_pars["M0"] = xu.with_unit(angle("M0"), u.rad) if "s" not in pars: out_pars["s"] = xu.with_unit( - pm.Deterministic("s", tt.constant(s.value)), s.unit + pm.Deterministic("s", pt.constant(s.value)), s.unit ) if "P" not in pars: @@ -527,7 +512,7 @@ def default_linear_prior( sigma_K0=None, P0=None, sigma_v=None, poly_trend=1, model=None, pars=None ): r""" - Retrieve pymc3 variables that specify the default prior on the linear + Retrieve pymc variables that specify the default prior on the linear parameters of The Joker. See docstring of `JokerPrior.default()` for more information. @@ -543,13 +528,10 @@ def default_linear_prior( sigma_K0 : `~astropy.units.Quantity` [speed] P0 : `~astropy.units.Quantity` [time] sigma_v : iterable of `~astropy.units.Quantity` - model : `pymc3.Model` - This is either required, or this function must be called within a pymc3 + model : `pymc.Model` + This is either required, or this function must be called within a pymc model context. """ - import exoplanet.units as xu - import pymc3 as pm - from .distributions import FixedCompanionMass model = pm.modelcontext(model) diff --git a/thejoker/samples_helpers.py b/thejoker/samples_helpers.py index 27280fb3..94e32551 100644 --- a/thejoker/samples_helpers.py +++ b/thejoker/samples_helpers.py @@ -2,13 +2,12 @@ import os import warnings +from astropy.io.misc.hdf5 import _encode_mixins, meta_path + # Third-party from astropy.table.meta import get_header_from_yaml, get_yaml_from_table -from astropy.io.misc.hdf5 import _encode_mixins, meta_path -from astropy.utils.exceptions import AstropyUserWarning from astropy.utils import metadata - -from thejoker.thejoker import validate_prepare_data +from astropy.utils.exceptions import AstropyUserWarning def _custom_tbl_dtype_compare(dtype1, dtype2): @@ -19,23 +18,31 @@ def _custom_tbl_dtype_compare(dtype1, dtype2): for d1, d2 in zip(dtype1, dtype2): for k in set(list(d1.keys()) + list(d2.keys())): - if k == 'unit': - if d1.get(k, '') != '' and k not in d2: + if k == "unit": + if d1.get(k, "") != "" and k not in d2: return False - if d2.get(k, '') != '' and k not in d1: + if d2.get(k, "") != "" and k not in d1: return False - if d1.get(k, '') != d2.get(k, ''): + if d1.get(k, "") != d2.get(k, ""): return False else: - if d1.get(k, '1') != d2.get(k, '2'): + if d1.get(k, "1") != d2.get(k, "2"): return False return True -def write_table_hdf5(table, output, path=None, compression=False, - append=False, overwrite=False, serialize_meta=False, - metadata_conflicts='error', **create_dataset_kwargs): +def write_table_hdf5( + table, + output, + path=None, + compression=False, + append=False, + overwrite=False, + serialize_meta=False, + metadata_conflicts="error", + **create_dataset_kwargs, +): """ Write a Table object to an HDF5 file @@ -75,6 +82,7 @@ def write_table_hdf5(table, output, path=None, compression=False, """ from astropy.table import meta + try: import h5py except ImportError: @@ -82,23 +90,27 @@ def write_table_hdf5(table, output, path=None, compression=False, if path is None: # table is just an arbitrary, hardcoded string here. - path = '__astropy_table__' - elif path.endswith('/'): + path = "__astropy_table__" + elif path.endswith("/"): raise ValueError("table path should end with table name, not /") - if '/' in path: - group, name = path.rsplit('/', 1) + if "/" in path: + group, name = path.rsplit("/", 1) else: group, name = None, path if isinstance(output, (h5py.File, h5py.Group)): - if len(list(output.keys())) > 0 and name == '__astropy_table__': - raise ValueError("table path should always be set via the " - "path= argument when writing to existing " - "files") - elif name == '__astropy_table__': - warnings.warn("table path was not set via the path= argument; " - "using default path {}".format(path)) + if len(list(output.keys())) > 0 and name == "__astropy_table__": + raise ValueError( + "table path should always be set via the " + "path= argument when writing to existing " + "files" + ) + elif name == "__astropy_table__": + warnings.warn( + "table path was not set via the path= argument; " + f"using default path {path}" + ) if group: try: @@ -109,7 +121,6 @@ def write_table_hdf5(table, output, path=None, compression=False, output_group = output elif isinstance(output, str): - if os.path.exists(output) and not append: if overwrite and not append: os.remove(output) @@ -117,22 +128,25 @@ def write_table_hdf5(table, output, path=None, compression=False, raise OSError(f"File exists: {output}") # Open the file for appending or writing - f = h5py.File(output, 'a' if append else 'w') + f = h5py.File(output, "a" if append else "w") # Recursively call the write function try: - return write_table_hdf5(table, f, path=path, - compression=compression, append=append, - overwrite=overwrite, - serialize_meta=serialize_meta, - **create_dataset_kwargs) + return write_table_hdf5( + table, + f, + path=path, + compression=compression, + append=append, + overwrite=overwrite, + serialize_meta=serialize_meta, + **create_dataset_kwargs, + ) finally: f.close() else: - - raise TypeError('output should be a string or an h5py File or ' - 'Group object') + raise TypeError("output should be a string or an h5py File or " "Group object") # Check whether table already exists existing_header = None @@ -146,13 +160,16 @@ def write_table_hdf5(table, output, path=None, compression=False, # the table to have been written by this function in the past, so it # should have a metadata header if meta_path(name) not in output_group: - raise ValueError("No metadata exists for existing table. We " - "can only append tables if metadata " - "is consistent for all tables") + raise ValueError( + "No metadata exists for existing table. We " + "can only append tables if metadata " + "is consistent for all tables" + ) # Load existing table header: existing_header = get_header_from_yaml( - h.decode('utf-8') for h in output_group[meta_path(name)]) + h.decode("utf-8") for h in output_group[meta_path(name)] + ) else: raise OSError(f"Table {path} already exists") @@ -162,7 +179,7 @@ def write_table_hdf5(table, output, path=None, compression=False, # Table with numpy unicode strings can't be written in HDF5 so # to write such a table a copy of table is made containing columns as # bytestrings. Now this copy of the table can be written in HDF5. - if any(col.info.dtype.kind == 'U' for col in table.itercols()): + if any(col.info.dtype.kind == "U" for col in table.itercols()): table = table.copy(copy_data=False) table.convert_unicode_to_bytestring() @@ -171,31 +188,36 @@ def write_table_hdf5(table, output, path=None, compression=False, # HDF5 can store natively (name, dtype) with no meta. if serialize_meta is False: for col in table.itercols(): - for attr in ('unit', 'format', 'description', 'meta'): + for attr in ("unit", "format", "description", "meta"): if getattr(col.info, attr, None) not in (None, {}): - warnings.warn("table contains column(s) with defined 'unit', 'format'," - " 'description', or 'meta' info attributes. These will" - " be dropped since serialize_meta=False.", - AstropyUserWarning) + warnings.warn( + "table contains column(s) with defined 'unit', 'format'," + " 'description', or 'meta' info attributes. These will" + " be dropped since serialize_meta=False.", + AstropyUserWarning, + ) if existing_header is None: # Just write the table and metadata # Write the table to the file if compression: if compression is True: - compression = 'gzip' - dset = output_group.create_dataset(name, data=table.as_array(), - compression=compression, - **create_dataset_kwargs) + compression = "gzip" + dset = output_group.create_dataset( + name, + data=table.as_array(), + compression=compression, + **create_dataset_kwargs, + ) else: - dset = output_group.create_dataset(name, data=table.as_array(), - **create_dataset_kwargs) + dset = output_group.create_dataset( + name, data=table.as_array(), **create_dataset_kwargs + ) if serialize_meta: header_yaml = meta.get_yaml_from_table(table) - header_encoded = [h.encode('utf-8') for h in header_yaml] - output_group.create_dataset(meta_path(name), - data=header_encoded) + header_encoded = [h.encode("utf-8") for h in header_yaml] + output_group.create_dataset(meta_path(name), data=header_encoded) else: # Write the Table meta dict key:value pairs to the file as HDF5 @@ -207,44 +229,50 @@ def write_table_hdf5(table, output, path=None, compression=False, try: dset.attrs[key] = val except TypeError: - warnings.warn("Attribute `{}` of type {} cannot be written to " - "HDF5 files - skipping. (Consider specifying " - "serialize_meta=True to write all meta data)" - .format(key, type(val)), AstropyUserWarning) + warnings.warn( + f"Attribute `{key}` of type {type(val)} cannot be written to " + "HDF5 files - skipping. (Consider specifying " + "serialize_meta=True to write all meta data)", + AstropyUserWarning, + ) else: # We need to append the tables! try: # FIXME: do something with the merged metadata! - metadata.merge(existing_header['meta'], - table.meta, - metadata_conflicts=metadata_conflicts) + metadata.merge( + existing_header["meta"], + table.meta, + metadata_conflicts=metadata_conflicts, + ) except metadata.MergeConflictError: raise metadata.MergeConflictError( "Cannot append table to existing file because " "the existing file table metadata and this " "table object's metadata do not match. If you " "want to ignore this issue, or change to a " - "warning, set metadata_conflicts='silent' or 'warn'.") + "warning, set metadata_conflicts='silent' or 'warn'." + ) # Now compare datatype of this object and on disk this_header = get_header_from_yaml(get_yaml_from_table(table)) - if not _custom_tbl_dtype_compare(existing_header['datatype'], - this_header['datatype']): + if not _custom_tbl_dtype_compare( + existing_header["datatype"], this_header["datatype"] + ): raise ValueError( "Cannot append table to existing file because " "the existing file table datatype and this " "object's table datatype do not match. " - f"{existing_header['datatype']} vs. {this_header['datatype']}") + f"{existing_header['datatype']} vs. {this_header['datatype']}" + ) # If we got here, we can now try to append: current_size = len(output_group[name]) - output_group[name].resize((current_size + len(table), )) + output_group[name].resize((current_size + len(table),)) output_group[name][current_size:] = table.as_array() -def inferencedata_to_samples(joker_prior, inferencedata, data, - prune_divergences=True): +def inferencedata_to_samples(joker_prior, inferencedata, data, prune_divergences=True): """ Create a ``JokerSamples`` instance from an arviz object. @@ -256,23 +284,26 @@ def inferencedata_to_samples(joker_prior, inferencedata, data, prune_divergences : bool (optional) """ + import thejoker.units as xu from thejoker.samples import JokerSamples - import exoplanet.units as xu + from thejoker.thejoker import validate_prepare_data - if hasattr(inferencedata, 'posterior'): + if hasattr(inferencedata, "posterior"): posterior = inferencedata.posterior else: posterior = inferencedata inferencedata = None - data, *_ = validate_prepare_data(data, - joker_prior.poly_trend, - joker_prior.n_offsets) + data, *_ = validate_prepare_data( + data, joker_prior.poly_trend, joker_prior.n_offsets + ) - samples = JokerSamples(poly_trend=joker_prior.poly_trend, - n_offsets=joker_prior.n_offsets, - t_ref=data.t_ref) + samples = JokerSamples( + poly_trend=joker_prior.poly_trend, + n_offsets=joker_prior.n_offsets, + t_ref=data.t_ref, + ) names = joker_prior.par_names @@ -284,10 +315,10 @@ def inferencedata_to_samples(joker_prior, inferencedata, data, else: samples[name] = posterior[name].values.ravel() - if hasattr(posterior, 'logp'): - samples['ln_posterior'] = posterior.logp.values.ravel() + if hasattr(posterior, "logp"): + samples["ln_posterior"] = posterior.logp.values.ravel() - for name in ['ln_likelihood', 'ln_prior']: + for name in ["ln_likelihood", "ln_prior"]: if hasattr(posterior, name): samples[name] = getattr(posterior, name).values.ravel() @@ -296,7 +327,8 @@ def inferencedata_to_samples(joker_prior, inferencedata, data, raise ValueError( "If you want to remove divergences, you must pass in the root " "level inferencedata object (instead of, e.g., inferencedata. " - "posterior") + "posterior" + ) divergences = inferencedata.sample_stats.diverging.values.ravel() samples = samples[~divergences] @@ -306,25 +338,27 @@ def inferencedata_to_samples(joker_prior, inferencedata, data, def trace_to_samples(self, trace, data, names=None): """ - Create a ``JokerSamples`` instance from a pymc3 trace object. + Create a ``JokerSamples`` instance from a pymc trace object. Parameters ---------- - trace : `~pymc3.backends.base.MultiTrace` + trace : `~pymc.backends.base.MultiTrace` """ - import pymc3 as pm - import exoplanet.units as xu + import pymc as pm + + import thejoker.units as xu from thejoker.samples import JokerSamples + from thejoker.thejoker import validate_prepare_data df = pm.trace_to_dataframe(trace) - data, *_ = validate_prepare_data(data, - self.prior.poly_trend, - self.prior.n_offsets) + data, *_ = validate_prepare_data(data, self.prior.poly_trend, self.prior.n_offsets) - samples = JokerSamples(poly_trend=self.prior.poly_trend, - n_offsets=self.prior.n_offsets, - t_ref=data.t_ref) + samples = JokerSamples( + poly_trend=self.prior.poly_trend, + n_offsets=self.prior.n_offsets, + t_ref=data.t_ref, + ) if names is None: names = self.prior.par_names diff --git a/thejoker/src/fast_likelihood.pyx b/thejoker/src/fast_likelihood.pyx index 2fc367c0..f3021e3a 100644 --- a/thejoker/src/fast_likelihood.pyx +++ b/thejoker/src/fast_likelihood.pyx @@ -15,6 +15,7 @@ np.import_array() import cython cimport cython cimport scipy.linalg.cython_lapack as lapack +import thejoker.units as xu # from libc.stdio cimport printf from libc.math cimport pow, log, fabs, pi @@ -202,7 +203,6 @@ cdef class CJokerHelper: # put v0_offsets variances into Lambda # - validated to be Normal() in JokerPrior - import exoplanet.units as xu for i in range(self.n_offsets): name = prior.v0_offsets[i].name dist = prior.v0_offsets[i].distribution @@ -216,16 +216,23 @@ cdef class CJokerHelper: # --------------------------------------------------------------------- # TODO: This is a bit of a hack: from ..distributions import FixedCompanionMass - if isinstance(prior.pars['K'].distribution, FixedCompanionMass): + if prior.pars['K'].owner.op._print_name[0] == "FixedCompanionMass": self.fixed_K_prior = 0 else: self.fixed_K_prior = 1 for i, name in enumerate(prior._linear_equiv_units.keys()): - dist = prior.model[name].distribution _unit = getattr(prior.model[name], xu.UNIT_ATTR_NAME) to_unit = self.internal_units[name] - mu = (dist.mean.eval() * _unit).to_value(to_unit) + + dist = prior.model[name] + # first three are (rng, size, dtype) as per + # https://github.com/pymc-devs/pymc/blob/main/pymc/printing.py#L43 + pars = dist.owner.inputs[3:] + + # mean is par 0 + mu = (pars[0].eval() * _unit).to_value(to_unit) + std = (pars[1].eval() * _unit).to_value(to_unit) if name == 'K' and self.fixed_K_prior == 0: # TODO: here's the major hack @@ -236,12 +243,12 @@ cdef class CJokerHelper: self.mu[i] = mu elif name == 'v0': - self.Lambda[i] = (dist.sd.eval() * _unit).to_value(to_unit) ** 2 + self.Lambda[i] = std ** 2 self.mu[i] = mu else: # v1, v2, etc. j = i + self.n_offsets - self.Lambda[j] = (dist.sd.eval() * _unit).to_value(to_unit) ** 2 + self.Lambda[j] = std ** 2 self.mu[j] = mu # --------------------------------------------------------------------- diff --git a/thejoker/src/setup_package.py b/thejoker/src/setup_package.py index 1c9da759..524d3cd0 100644 --- a/thejoker/src/setup_package.py +++ b/thejoker/src/setup_package.py @@ -1,6 +1,7 @@ -from distutils.core import Extension -from collections import defaultdict import os +from collections import defaultdict + +from setuptools import Extension def get_extensions(): @@ -10,14 +11,14 @@ def get_extensions(): import twobody cfg = defaultdict(list) - cfg['include_dirs'].append(np.get_include()) + cfg["include_dirs"].append(np.get_include()) twobody_path = os.path.dirname(twobody.__file__) - cfg['include_dirs'].append(twobody_path) - cfg['sources'].append(os.path.join(twobody_path, 'src/twobody.c')) + cfg["include_dirs"].append(twobody_path) + cfg["sources"].append(os.path.join(twobody_path, "src/twobody.c")) - cfg['extra_compile_args'].append('--std=gnu99') - cfg['sources'].append('thejoker/src/fast_likelihood.pyx') - exts.append(Extension('thejoker.src.fast_likelihood', **cfg)) + cfg["extra_compile_args"].append("--std=gnu99") + cfg["sources"].append("thejoker/src/fast_likelihood.pyx") + exts.append(Extension("thejoker.src.fast_likelihood", **cfg)) return exts diff --git a/thejoker/thejoker.py b/thejoker/thejoker.py index 72771f2b..9dab5f82 100644 --- a/thejoker/thejoker.py +++ b/thejoker/thejoker.py @@ -1,6 +1,5 @@ # Standard library import os -import warnings import numpy as np @@ -368,7 +367,7 @@ def iterative_rejection_sample( def setup_mcmc(self, data, joker_samples, model=None, custom_func=None): """ - Setup the model to run MCMC using pymc3. + Setup the model to run MCMC using pymc. Parameters ---------- @@ -376,13 +375,13 @@ def setup_mcmc(self, data, joker_samples, model=None, custom_func=None): The radial velocity data, or an iterable containing ``RVData`` objects for each data source. joker_samples : `~thejoker.JokerSamples` - If a single sample is passed in, this is packed into a pymc3 + If a single sample is passed in, this is packed into a pymc initialization dictionary and returned after setting up. If multiple samples are passed in, the median (along period) sample is taken and returned after setting up for MCMC. - model : `pymc3.Model` + model : `pymc.Model` This is either required, or this function must be called within a - pymc3 model context. + pymc model context. custom_func : callable (optional) Returns @@ -390,10 +389,11 @@ def setup_mcmc(self, data, joker_samples, model=None, custom_func=None): mcmc_init : dict """ - import aesara_theano_fallback.tensor as tt - import exoplanet as xo - import exoplanet.units as xu - import pymc3 as pm + import pymc as pm + import pytensor.tensor as pt + + import thejoker.units as xu + from thejoker._keplerian_orbit import KeplerianOrbit model = _validate_model(model) @@ -441,7 +441,7 @@ def setup_mcmc(self, data, joker_samples, model=None, custom_func=None): with model: # Set up the orbit model - orbit = xo.orbits.KeplerianOrbit( + orbit = KeplerianOrbit( period=p["P"], ecc=p["e"], omega=p["omega"], @@ -461,13 +461,13 @@ def setup_mcmc(self, data, joker_samples, model=None, custom_func=None): + [p[name] for name in offset_names] + [p[name] for name in vtrend_names[1:]] ) # skip v0 - v_trend_vec = tt.stack(v_pars, axis=0) - trend = tt.dot(M, v_trend_vec) + v_trend_vec = pt.stack(v_pars, axis=0) + trend = pt.dot(M, v_trend_vec) rv_model = orbit.get_radial_velocity(x, K=p["K"]) + trend pm.Deterministic("model_rv", rv_model) - err = tt.sqrt(err**2 + p["s"] ** 2) + err = pt.sqrt(err**2 + p["s"] ** 2) pm.Normal("obs", mu=rv_model, sd=err, observed=y) pm.Deterministic("logp", model.logpt) @@ -480,21 +480,3 @@ def setup_mcmc(self, data, joker_samples, model=None, custom_func=None): pm.Deterministic("ln_prior", model.logpt - lnlike) return mcmc_init - - def trace_to_samples(self, trace, data, names=None): - """ - Create a ``JokerSamples`` instance from a pymc3 trace object. - - Parameters - ---------- - trace : `~pymc3.backends.base.MultiTrace` - """ - warnings.warn( - "This method is deprecated: Use " - "thejoker.samples_helpers.trace_to_samples() instead", - UserWarning, - ) - - from thejoker.samples_helpers import trace_to_samples - - return trace_to_samples(self, trace, data, names) diff --git a/thejoker/units.py b/thejoker/units.py new file mode 100644 index 00000000..83d9c31e --- /dev/null +++ b/thejoker/units.py @@ -0,0 +1,52 @@ +"""Originally from the exoplanet project""" + +__all__ = ["with_unit", "has_unit", "to_unit"] + +from pytensor.tensor import as_tensor_variable + +UNIT_ATTR_NAME = "__tensor_unit__" + + +def with_unit(obj, unit): + """Decorate a tensor with Astropy units + + Parameters + ---------- + obj + The tensor object + unit : astropy.units.Unit + The units for this object + + """ + if hasattr(obj, UNIT_ATTR_NAME): + msg = f"{obj!r} already has units" + raise TypeError(msg) + obj = as_tensor_variable(obj) + setattr(obj, UNIT_ATTR_NAME, unit) + return obj + + +def has_unit(obj): + """Does an object have units as defined by exoplanet?""" + return hasattr(obj, UNIT_ATTR_NAME) + + +def to_unit(obj, target): + """Convert a Theano tensor with units to a target set of units + + Parameters + ---------- + obj + The Theano tensor + target : astropy.units.Unit + The target units + + Returns + ------- + A tensor in the right units + + """ + if not has_unit(obj): + return obj + base = getattr(obj, UNIT_ATTR_NAME) + return obj * base.to(target) diff --git a/thejoker/utils.py b/thejoker/utils.py index d1aaf542..efd587f9 100644 --- a/thejoker/utils.py +++ b/thejoker/utils.py @@ -16,9 +16,6 @@ from astropy.table.meta import get_header_from_yaml from astropy.utils.decorators import deprecated_renamed_argument -# Package -from .samples import JokerSamples - __all__ = ["batch_tasks", "table_header_to_units", "read_batch", "tempfile_decorator"] @@ -173,6 +170,7 @@ def read_batch_slice(prior_samples_file, columns, slice, units=None): Read a batch (row block) of prior samples into a plain numpy array, converting units where necessary. """ + from .samples import JokerSamples path = JokerSamples._hdf5_path @@ -205,6 +203,8 @@ def read_batch_idx(prior_samples_file, columns, idx, units=None): Read a batch (row block) of prior samples specified by the input index array, ``idx``, into a plain numpy array, converting units where necessary. """ + from .samples import JokerSamples + path = JokerSamples._hdf5_path # We have to do this with h5py because current (2021-02-05) versions of @@ -233,6 +233,7 @@ def read_random_batch(prior_samples_file, columns, size, units=None, rng=None): Read a random batch (row block) of prior samples into a plain numpy array, converting units where necessary. """ + from .samples import JokerSamples if rng is None: rng = np.random.default_rng() @@ -245,6 +246,8 @@ def read_random_batch(prior_samples_file, columns, size, units=None, rng=None): def tempfile_decorator(func): + from .samples import JokerSamples + wrapped_signature = inspect.signature(func) func_args = list(wrapped_signature.parameters.keys()) if "prior_samples_file" not in func_args: From 5228c99c05952afe7aaf9e9d2bcd163a655b9378 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Sun, 3 Mar 2024 12:30:04 -0500 Subject: [PATCH 12/50] Remove old setup_package files --- thejoker/src/setup_package.py | 24 ------------------------ thejoker/tests/setup_package.py | 3 --- 2 files changed, 27 deletions(-) delete mode 100644 thejoker/src/setup_package.py delete mode 100644 thejoker/tests/setup_package.py diff --git a/thejoker/src/setup_package.py b/thejoker/src/setup_package.py deleted file mode 100644 index 524d3cd0..00000000 --- a/thejoker/src/setup_package.py +++ /dev/null @@ -1,24 +0,0 @@ -import os -from collections import defaultdict - -from setuptools import Extension - - -def get_extensions(): - exts = [] - - import numpy as np - import twobody - - cfg = defaultdict(list) - cfg["include_dirs"].append(np.get_include()) - - twobody_path = os.path.dirname(twobody.__file__) - cfg["include_dirs"].append(twobody_path) - cfg["sources"].append(os.path.join(twobody_path, "src/twobody.c")) - - cfg["extra_compile_args"].append("--std=gnu99") - cfg["sources"].append("thejoker/src/fast_likelihood.pyx") - exts.append(Extension("thejoker.src.fast_likelihood", **cfg)) - - return exts diff --git a/thejoker/tests/setup_package.py b/thejoker/tests/setup_package.py deleted file mode 100644 index f2fd9ed4..00000000 --- a/thejoker/tests/setup_package.py +++ /dev/null @@ -1,3 +0,0 @@ -def get_package_data(): - return { - _ASTROPY_PACKAGE_NAME_ + '.tests': ['coveragerc']} From 97970a2c0f7c1598153665b0e203f08885fdc3e8 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Sun, 3 Mar 2024 12:30:13 -0500 Subject: [PATCH 13/50] pymc version req --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 3b390c6b..a06cb12f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ dependencies = [ "scipy", "h5py", "schwimmbad>=0.3.1", - "pymc", + "pymc>=5", "pymc_ext @ git+https://github.com/exoplanet-dev/pymc-ext", "exoplanet-core[pymc]", "tables", From 2b5f9b793431395fae23d12f2a5e96a364ac029a Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Sun, 3 Mar 2024 12:40:32 -0500 Subject: [PATCH 14/50] formatting and getting tests to at least run (but fail) --- docs/conf.py | 116 ++++++----- docs/install.rst | 41 +--- pyproject.toml | 2 +- thejoker/src/tests/py_likelihood.py | 69 +++--- thejoker/src/tests/test_fast_likelihood.py | 57 ++--- thejoker/tests/test_data.py | 231 ++++++++++----------- thejoker/tests/test_likelihood_helpers.py | 23 +- thejoker/tests/test_prior.py | 158 +++++++------- thejoker/tests/test_utils.py | 116 ++++++----- 9 files changed, 400 insertions(+), 413 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index 27d33120..2851ae96 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- # Licensed under a 3-clause BSD style license - see LICENSE.rst # # Astropy documentation build configuration file. @@ -25,31 +24,34 @@ # Thus, any C-extensions that are needed to build the documentation will *not* # be accessible, and the documentation will not build correctly. +import datetime import os import sys -import datetime from importlib import import_module try: from sphinx_astropy.conf.v1 import * # noqa except ImportError: - print('ERROR: the documentation requires the sphinx-astropy package to be installed') + print( + "ERROR: the documentation requires the sphinx-astropy package to be installed" + ) sys.exit(1) # Get configuration information from setup.cfg from configparser import ConfigParser + conf = ConfigParser() -conf.read([os.path.join(os.path.dirname(__file__), '..', 'setup.cfg')]) -setup_cfg = dict(conf.items('metadata')) +conf.read([os.path.join(os.path.dirname(__file__), "..", "setup.cfg")]) +setup_cfg = dict(conf.items("metadata")) # -- General configuration ---------------------------------------------------- # By default, highlight as Python 3. -highlight_language = 'python3' +highlight_language = "python3" # If your documentation needs a minimal Sphinx version, state it here. -#needs_sphinx = '1.2' +# needs_sphinx = '1.2' # To perform a Sphinx version check that needs to be more specific than # major.minor, call `check_sphinx_version("x.y.z")` here. @@ -57,8 +59,8 @@ # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files. -exclude_patterns.append('_templates') -exclude_patterns.append('**.ipynb_checkpoints') +exclude_patterns.append("_templates") +exclude_patterns.append("**.ipynb_checkpoints") # This is added to the end of RST files - a good place to put substitutions to # be used globally. @@ -69,20 +71,19 @@ # -- Project information ------------------------------------------------------ # This does not *have* to match the package name, but typically does -project = setup_cfg['name'] -author = setup_cfg['author'] -copyright = '{0}, {1}'.format( - datetime.datetime.now().year, setup_cfg['author']) +project = setup_cfg["name"] +author = setup_cfg["author"] +copyright = "{0}, {1}".format(datetime.datetime.now().year, setup_cfg["author"]) # The version info for the project you're documenting, acts as replacement for # |version| and |release|, also used in various other places throughout the # built documents. -import_module(setup_cfg['name']) -package = sys.modules[setup_cfg['name']] +import_module(setup_cfg["name"]) +package = sys.modules[setup_cfg["name"]] # The short X.Y version. -version = package.__version__.split('-', 1)[0] +version = package.__version__.split("-", 1)[0] # The full version, including alpha/beta/rc tags. release = package.__version__ @@ -99,71 +100,71 @@ # Add any paths that contain custom themes here, relative to this directory. # To use a different custom theme, add the directory containing the theme. -#html_theme_path = [] +# html_theme_path = [] # The theme to use for HTML and HTML Help pages. See the documentation for # a list of builtin themes. To override the custom theme, set this to the # name of a builtin theme or the name of a custom theme in html_theme_path. -#html_theme = None +# html_theme = None # Please update these texts to match the name of your package. html_theme_options = { - 'logotext1': 'The', # white, semi-bold - 'logotext2': 'Joker', # orange, light - 'logotext3': ':docs' # white, light + "logotext1": "The", # white, semi-bold + "logotext2": "Joker", # orange, light + "logotext3": ":docs", # white, light } # Custom sidebar templates, maps document names to template names. -#html_sidebars = {} +# html_sidebars = {} # The name of an image file (relative to this directory) to place at the top # of the sidebar. -#html_logo = '' +# html_logo = '' # The name of an image file (within the static path) to use as favicon of the # docs. This file should be a Windows icon file (.ico) being 16x16 or 32x32 # pixels large. -path = os.path.abspath(os.path.join(os.path.dirname(__file__), '_static')) -html_favicon = os.path.join(path, 'icon.ico') +path = os.path.abspath(os.path.join(os.path.dirname(__file__), "_static")) +html_favicon = os.path.join(path, "icon.ico") # If not '', a 'Last updated on:' timestamp is inserted at every page bottom, # using the given strftime format. -#html_last_updated_fmt = '' +# html_last_updated_fmt = '' # The name for this set of Sphinx documents. If None, it defaults to # " v documentation". -html_title = '{0} v{1}'.format(project, release) +html_title = f"{project} v{release}" # Output file base name for HTML help builder. -htmlhelp_basename = project + 'doc' +htmlhelp_basename = project + "doc" -html_static_path = ['_static'] -html_style = 'thejoker.css' +html_static_path = ["_static"] +html_style = "thejoker.css" # -- Options for LaTeX output ------------------------------------------------- # Grouping the document tree into LaTeX files. List of tuples # (source start file, target name, title, author, documentclass [howto/manual]). -latex_documents = [('index', project + '.tex', project + u' Documentation', - author, 'manual')] +latex_documents = [ + ("index", project + ".tex", project + " Documentation", author, "manual") +] # -- Options for manual page output ------------------------------------------- # One entry per manual page. List of tuples # (source start file, name, description, authors, manual section). -man_pages = [('index', project.lower(), project + u' Documentation', - [author], 1)] +man_pages = [("index", project.lower(), project + " Documentation", [author], 1)] # -- Options for the edit_on_github extension --------------------------------- -if eval(setup_cfg.get('edit_on_github')): - extensions += ['sphinx_astropy.ext.edit_on_github'] +if eval(setup_cfg.get("edit_on_github")): + extensions += ["sphinx_astropy.ext.edit_on_github"] - versionmod = __import__(setup_cfg['package_name'] + '.version') - edit_on_github_project = setup_cfg['github_project'] + versionmod = __import__(setup_cfg["package_name"] + ".version") + edit_on_github_project = setup_cfg["github_project"] if versionmod.version.release: edit_on_github_branch = "v" + versionmod.version.version else: @@ -173,14 +174,15 @@ edit_on_github_doc_root = "docs" # -- Resolving issue number to links in changelog ----------------------------- -github_issues_url = 'https://github.com/{0}/issues/'.format(setup_cfg['github_project']) +github_issues_url = "https://github.com/{0}/issues/".format(setup_cfg["github_project"]) -intersphinx_mapping['h5py'] = ('http://docs.h5py.org/en/latest/', None) -intersphinx_mapping['pymc3'] = ('https://docs.pymc.io/', None) -intersphinx_mapping['twobody'] = ('https://twobody.readthedocs.io/en/latest/', - None) -intersphinx_mapping['scwhimmbad'] = ( - 'https://schwimmbad.readthedocs.io/en/latest/', None) +intersphinx_mapping["h5py"] = ("http://docs.h5py.org/en/latest/", None) +intersphinx_mapping["pymc"] = ("https://docs.pymc.io/", None) +intersphinx_mapping["twobody"] = ("https://twobody.readthedocs.io/en/latest/", None) +intersphinx_mapping["scwhimmbad"] = ( + "https://schwimmbad.readthedocs.io/en/latest/", + None, +) # see if we're running on CI: ON_CI = os.environ.get("CI", False) @@ -189,21 +191,21 @@ # Use astropy plot style plot_rcparams = dict() if not ON_CI: - plot_rcparams['text.usetex'] = True -plot_rcparams['savefig.facecolor'] = 'none' -plot_rcparams['savefig.bbox'] = 'tight' + plot_rcparams["text.usetex"] = True +plot_rcparams["savefig.facecolor"] = "none" +plot_rcparams["savefig.bbox"] = "tight" plot_apply_rcparams = True -plot_formats = [('png', 512)] +plot_formats = [("png", 512)] # nbsphinx config: -exclude_patterns.append('make-data.*') -exclude_patterns.append('*/make-data.*') +exclude_patterns.append("make-data.*") +exclude_patterns.append("*/make-data.*") -extensions += ['nbsphinx'] -extensions += ['IPython.sphinxext.ipython_console_highlighting'] +extensions += ["nbsphinx"] +extensions += ["IPython.sphinxext.ipython_console_highlighting"] -extensions += ['sphinx.ext.mathjax'] +extensions += ["sphinx.ext.mathjax"] mathjax_path = "https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_HTMLorMML" # nbsphinx_execute_arguments = [ @@ -213,10 +215,10 @@ nbsphinx_timeout = 600 if PR: - nbsphinx_execute = 'never' + nbsphinx_execute = "never" if ON_CI: - nbsphinx_kernel_name = 'python3' + nbsphinx_kernel_name = "python3" else: - nbsphinx_kernel_name = os.environ.get('NBSPHINX_KERNEL', 'python3') + nbsphinx_kernel_name = os.environ.get("NBSPHINX_KERNEL", "python3") diff --git a/docs/install.rst b/docs/install.rst index 26b2800a..6c389711 100644 --- a/docs/install.rst +++ b/docs/install.rst @@ -21,47 +21,14 @@ GitHub project link: pip install git+https://github.com/adrn/thejoker.git -From source -=========== - -Alternatively, you can clone the repo (or download a zip of the latest code from -the `GitHub page `_) and use the provided -`environment.yml `_ file to create a new -environment for |thejoker| that is set up with all of the dependencies -installed. To install in a new ``conda`` environment, change to the top-level -directory of the cloned repository, and run: - -.. code-block:: bash - - conda env create - -When this finishes, activate the environment +To install from source (i.e. from the cloned Git repository), also use pip: .. code-block:: bash - source activate thejoker - -The project is installable using the standard - -.. code-block:: bash - - python setup.py install + cd /path/to/thejoker + pip install . Dependencies ============ -- numpy -- scipy -- astropy -- h5py -- emcee -- pytables -- exoplanet -- pymc3 -- `schwimmbad `_ -- `twobody `_ - -Optional Dependencies ---------------------- - -- matplotlib +See the `pyproject.toml `_ file for the most up-to-date list of dependencies. diff --git a/pyproject.toml b/pyproject.toml index a06cb12f..cf9b548e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ dependencies = [ "schwimmbad>=0.3.1", "pymc>=5", "pymc_ext @ git+https://github.com/exoplanet-dev/pymc-ext", - "exoplanet-core[pymc]", + "exoplanet-core[pymc] @ git+https://github.com/exoplanet-dev/exoplanet-core", "tables", ] diff --git a/thejoker/src/tests/py_likelihood.py b/thejoker/src/tests/py_likelihood.py index de6a9665..832751f9 100644 --- a/thejoker/src/tests/py_likelihood.py +++ b/thejoker/src/tests/py_likelihood.py @@ -5,16 +5,13 @@ # Third-party import astropy.units as u import numpy as np -from twobody.wrap import cy_rv_from_elements from astroML.utils import log_multivariate_gaussian -# from scipy.stats import multivariate_normal +from twobody.wrap import cy_rv_from_elements -# Project -from ...samples import JokerSamples from ...distributions import FixedCompanionMass -from ...utils import DEFAULT_RNG +from ...samples import JokerSamples -__all__ = ['get_ivar', 'likelihood_worker', 'marginal_ln_likelihood'] +__all__ = ["get_ivar", "likelihood_worker", "marginal_ln_likelihood"] def get_ivar(data, s): @@ -126,7 +123,7 @@ def design_matrix(nonlinear_p, data, prior): t = data._t_bmjd t0 = data._t_ref_bmjd - zdot = cy_rv_from_elements(t, P, 1., ecc, omega, M0, t0, 1e-8, 128) + zdot = cy_rv_from_elements(t, P, 1.0, ecc, omega, M0, t0, 1e-8, 128) M1 = np.vander(t - t0, N=prior.poly_trend, increasing=True) M = np.hstack((zdot[:, None], M1)) @@ -136,7 +133,7 @@ def design_matrix(nonlinear_p, data, prior): def get_M_Lambda_ivar(samples, prior, data): v_unit = data.rv.unit - units = {'K': v_unit, 's': v_unit} + units = {"K": v_unit, "s": v_unit} for i, k in enumerate(list(prior._linear_equiv_units.keys())[1:]): # skip K units[k] = v_unit / u.day**i packed_samples, _ = samples.pack(units=units) @@ -146,14 +143,14 @@ def get_M_Lambda_ivar(samples, prior, data): Lambda = np.zeros(n_linear) for i, k in enumerate(prior._linear_equiv_units.keys()): - if k == 'K': + if k == "K": continue # set below Lambda[i] = prior.pars[k].distribution.sd.eval() ** 2 - K_dist = prior.pars['K'].distribution + K_dist = prior.pars["K"].distribution if isinstance(K_dist, FixedCompanionMass): sigma_K0 = K_dist._sigma_K0.to_value(v_unit) - P0 = K_dist._P0.to_value(samples['P'].unit) + P0 = K_dist._P0.to_value(samples["P"].unit) max_K2 = K_dist._max_K.to_value(v_unit) ** 2 else: Lambda[0] = K_dist.sd.eval() ** 2 @@ -161,9 +158,9 @@ def get_M_Lambda_ivar(samples, prior, data): for n in range(n_samples): M = design_matrix(packed_samples[n], data, prior) if isinstance(K_dist, FixedCompanionMass): - P = samples['P'][n].value - e = samples['e'][n] - Lambda[0] = sigma_K0**2 / (1 - e**2) * (P / P0)**(-2/3) + P = samples["P"][n].value + e = samples["e"][n] + Lambda[0] = sigma_K0**2 / (1 - e**2) * (P / P0) ** (-2 / 3) Lambda[0] = min(max_K2, Lambda[0]) # jitter must be in same units as the data RV's / ivar! @@ -197,9 +194,9 @@ def marginal_ln_likelihood(samples, prior, data): marg_ll = np.zeros(n_samples) for n, M, Lambda, ivar, *_ in get_M_Lambda_ivar(samples, prior, data): try: - marg_ll[n], *_ = likelihood_worker(data.rv.value, ivar, M, - mu, np.diag(Lambda), - make_aA=False) + marg_ll[n], *_ = likelihood_worker( + data.rv.value, ivar, M, mu, np.diag(Lambda), make_aA=False + ) except np.linalg.LinAlgError as e: raise e @@ -220,7 +217,7 @@ def rejection_sample(samples, prior, data, rnd=None): mu = np.zeros(n_linear) if rnd is None: - rnd = DEFAULT_RNG() + rnd = np.random.default_rng() ll = marginal_ln_likelihood(samples, prior, data) uu = rnd.uniform(size=len(ll)) @@ -231,11 +228,12 @@ def rejection_sample(samples, prior, data, rnd=None): all_packed = np.zeros((n_good_samples, len(prior.par_names))) for n, M, Lambda, ivar, packed_nonlinear, units in get_M_Lambda_ivar( - good_samples, prior, data): + good_samples, prior, data + ): try: - _, b, B, a, A = likelihood_worker(data.rv.value, ivar, M, - mu, np.diag(Lambda), - make_aA=True) + _, b, B, a, A = likelihood_worker( + data.rv.value, ivar, M, mu, np.diag(Lambda), make_aA=True + ) except np.linalg.LinAlgError as e: raise e @@ -249,8 +247,7 @@ def rejection_sample(samples, prior, data, rnd=None): else: unpack_units[k] = samples[k].unit - return JokerSamples.unpack(all_packed, unpack_units, prior.poly_trend, - data.t_ref) + return JokerSamples.unpack(all_packed, unpack_units, prior.poly_trend, data.t_ref) def get_aAbB(samples, prior, data): @@ -269,22 +266,24 @@ def get_aAbB(samples, prior, data): n_times = len(data) mu = np.zeros(n_linear) - out = {'a': np.zeros((n_samples, n_linear)), - 'A': np.zeros((n_samples, n_linear, n_linear)), - 'b': np.zeros((n_samples, n_times)), - 'B': np.zeros((n_samples, n_times, n_times))} + out = { + "a": np.zeros((n_samples, n_linear)), + "A": np.zeros((n_samples, n_linear, n_linear)), + "b": np.zeros((n_samples, n_times)), + "B": np.zeros((n_samples, n_times, n_times)), + } for n, M, Lambda, ivar, *_ in get_M_Lambda_ivar(samples, prior, data): try: - _, b, B, a, A = likelihood_worker(data.rv.value, ivar, M, - mu, np.diag(Lambda), - make_aA=True) + _, b, B, a, A = likelihood_worker( + data.rv.value, ivar, M, mu, np.diag(Lambda), make_aA=True + ) except np.linalg.LinAlgError as e: raise e - out['a'][n] = a - out['A'][n] = A - out['b'][n] = b - out['B'][n] = B + out["a"][n] = a + out["A"][n] = A + out["b"][n] = b + out["B"][n] = B return out diff --git a/thejoker/src/tests/test_fast_likelihood.py b/thejoker/src/tests/test_fast_likelihood.py index 1f7761a5..618a9691 100644 --- a/thejoker/src/tests/test_fast_likelihood.py +++ b/thejoker/src/tests/test_fast_likelihood.py @@ -1,35 +1,39 @@ # Third-party +import time + import astropy.units as u import numpy as np -import time -import pymc3 as pm -import exoplanet.units as xu +import pymc as pm + +import thejoker.units as xu # Package from ...data import RVData +from ...likelihood_helpers import get_constant_term_design_matrix from ...prior import JokerPrior from ..fast_likelihood import CJokerHelper -from ...likelihood_helpers import get_constant_term_design_matrix -from .py_likelihood import marginal_ln_likelihood, get_aAbB +from .py_likelihood import get_aAbB, marginal_ln_likelihood # TODO: horrible copy-pasta test code below def test_against_py(): with pm.Model(): - K = xu.with_unit(pm.Normal('K', 0, 10.), - u.km/u.s) + K = xu.with_unit(pm.Normal("K", 0, 10.0), u.km / u.s) - prior = JokerPrior.default(P_min=8*u.day, P_max=32768*u.day, - s=0*u.m/u.s, - sigma_v=100*u.km/u.s, - pars={'K': K}) + prior = JokerPrior.default( + P_min=8 * u.day, + P_max=32768 * u.day, + s=0 * u.m / u.s, + sigma_v=100 * u.km / u.s, + pars={"K": K}, + ) # t = np.random.uniform(0, 250, 16) + 56831.324 t = np.sort(np.random.uniform(0, 250, 3)) + 56831.324 rv = np.cos(t) rv_err = np.random.uniform(0.1, 0.2, t.size) - data = RVData(t=t, rv=rv*u.km/u.s, rv_err=rv_err*u.km/u.s) + data = RVData(t=t, rv=rv * u.km / u.s, rv_err=rv_err * u.km / u.s) trend_M = get_constant_term_design_matrix(data) samples = prior.sample(size=8192) @@ -50,16 +54,19 @@ def test_against_py(): def test_scale_varK_against_py(): - prior = JokerPrior.default(P_min=8*u.day, P_max=32768*u.day, - s=0*u.m/u.s, - sigma_K0=25*u.km/u.s, - sigma_v=100*u.km/u.s) + prior = JokerPrior.default( + P_min=8 * u.day, + P_max=32768 * u.day, + s=0 * u.m / u.s, + sigma_K0=25 * u.km / u.s, + sigma_v=100 * u.km / u.s, + ) # t = np.random.uniform(0, 250, 16) + 56831.324 t = np.sort(np.random.uniform(0, 250, 3)) + 56831.324 rv = np.cos(t) rv_err = np.random.uniform(0.1, 0.2, t.size) - data = RVData(t=t, rv=rv*u.km/u.s, rv_err=rv_err*u.km/u.s) + data = RVData(t=t, rv=rv * u.km / u.s, rv_err=rv_err * u.km / u.s) trend_M = get_constant_term_design_matrix(data) samples = prior.sample(size=8192) @@ -81,19 +88,21 @@ def test_scale_varK_against_py(): def test_likelihood_helpers(): with pm.Model(): - K = xu.with_unit(pm.Normal('K', 0, 1.), - u.km/u.s) + K = xu.with_unit(pm.Normal("K", 0, 1.0), u.km / u.s) - prior = JokerPrior.default(P_min=8*u.day, P_max=32768*u.day, - s=0*u.m/u.s, - sigma_v=1*u.km/u.s, - pars={'K': K}) + prior = JokerPrior.default( + P_min=8 * u.day, + P_max=32768 * u.day, + s=0 * u.m / u.s, + sigma_v=1 * u.km / u.s, + pars={"K": K}, + ) # t = np.random.uniform(0, 250, 16) + 56831.324 t = np.sort(np.random.uniform(0, 250, 3)) + 56831.324 rv = np.cos(t) rv_err = np.random.uniform(0.1, 0.2, t.size) - data = RVData(t=t, rv=rv*u.km/u.s, rv_err=rv_err*u.km/u.s) + data = RVData(t=t, rv=rv * u.km / u.s, rv_err=rv_err * u.km / u.s) trend_M = get_constant_term_design_matrix(data) samples = prior.sample(size=16) # HACK: MAGIC NUMBER 16! diff --git a/thejoker/tests/test_data.py b/thejoker/tests/test_data.py index b3f7ad9f..c271de78 100644 --- a/thejoker/tests/test_data.py +++ b/thejoker/tests/test_data.py @@ -2,11 +2,13 @@ import warnings +import astropy.units as u + # Third-party from astropy.table import Table from astropy.time import Time from astropy.timeseries import TimeSeries -import astropy.units as u + try: from erfa import ErfaWarning except ImportError: # lts version of Astropy @@ -16,12 +18,14 @@ try: import matplotlib.pyplot as plt + HAS_MPL = True except ImportError: HAS_MPL = False try: import fuzzywuzzy # noqa + HAS_FUZZY = True except ImportError: HAS_FUZZY = False @@ -30,57 +34,54 @@ from ..data import RVData from ..data_helpers import guess_time_format, validate_prepare_data from ..prior import JokerPrior -from ..utils import DEFAULT_RNG def test_guess_time_format(): with warnings.catch_warnings(): - warnings.simplefilter('ignore', category=ErfaWarning) + warnings.simplefilter("ignore", category=ErfaWarning) for yr in np.arange(1975, 2040, 5): - assert guess_time_format(Time(f'{yr}-05-23').jd) == 'jd' - assert guess_time_format(Time(f'{yr}-05-23').mjd) == 'mjd' + assert guess_time_format(Time(f"{yr}-05-23").jd) == "jd" + assert guess_time_format(Time(f"{yr}-05-23").mjd) == "mjd" with pytest.raises(NotImplementedError): - guess_time_format('asdfasdf') + guess_time_format("asdfasdf") - for bad_val in np.array([0., 1450., 2500., 5000.]): + for bad_val in np.array([0.0, 1450.0, 2500.0, 5000.0]): with pytest.raises(ValueError): guess_time_format(bad_val) def get_valid_input(rnd=None, size=32): if rnd is None: - rnd = DEFAULT_RNG(42) + rnd = np.random.default_rng(42) - t_arr = rnd.uniform(55555., 56012., size=size) - t_obj = Time(t_arr, format='mjd') + t_arr = rnd.uniform(55555.0, 56012.0, size=size) + t_obj = Time(t_arr, format="mjd") - rv = 100 * np.sin(2*np.pi * t_arr / 15.) * u.km / u.s - err = rnd.uniform(0.1, 0.5, size=len(t_arr)) * u.km/u.s + rv = 100 * np.sin(2 * np.pi * t_arr / 15.0) * u.km / u.s + err = rnd.uniform(0.1, 0.5, size=len(t_arr)) * u.km / u.s cov = (np.diag(err.value) * err.unit) ** 2 _tbl = Table() - _tbl['rv'] = rnd.uniform(size=len(rv)) - _tbl['rv'].unit = u.km/u.s - _tbl['rv_err'] = rnd.uniform(size=len(rv)) - _tbl['rv_err'].unit = u.km/u.s + _tbl["rv"] = rnd.uniform(size=len(rv)) + _tbl["rv"].unit = u.km / u.s + _tbl["rv_err"] = rnd.uniform(size=len(rv)) + _tbl["rv_err"].unit = u.km / u.s - raw = {'t_arr': t_arr, - 't_obj': t_obj, - 'rv': rv, - 'err': err, - 'cov': cov} + raw = {"t_arr": t_arr, "t_obj": t_obj, "rv": rv, "err": err, "cov": cov} - return [dict(t=t_arr, rv=rv, rv_err=err), - (t_arr, rv, err), - (t_obj, rv, err), - (t_obj, _tbl['rv'], _tbl['rv_err']), - (t_arr, rv, cov), - (t_obj, rv, cov)], raw + return [ + dict(t=t_arr, rv=rv, rv_err=err), + (t_arr, rv, err), + (t_obj, rv, err), + (t_obj, _tbl["rv"], _tbl["rv_err"]), + (t_arr, rv, cov), + (t_obj, rv, cov), + ], raw def test_rvdata_init(): - rnd = DEFAULT_RNG(42) + rnd = np.random.default_rng(42) # Test valid initialization combos # These should succeed: @@ -91,11 +92,11 @@ def test_rvdata_init(): else: RVData(**x) - t_arr = raw['t_arr'] - t_obj = raw['t_obj'] - rv = raw['rv'] - err = raw['err'] - cov = raw['cov'] + t_arr = raw["t_arr"] + t_obj = raw["t_obj"] + rv = raw["rv"] + err = raw["err"] + cov = raw["cov"] # With/without clean: for i in range(1, 3): # skip time, because Time() catches nan values @@ -105,10 +106,10 @@ def test_rvdata_init(): inputs[i] = arr data = RVData(*inputs) - assert len(data) == (len(arr)-1) + assert len(data) == (len(arr) - 1) data = RVData(*inputs, clean=True) - assert len(data) == (len(arr)-1) + assert len(data) == (len(arr) - 1) data = RVData(*inputs, clean=False) assert len(data) == len(arr) @@ -148,7 +149,7 @@ def test_rvdata_init(): with pytest.raises(ValueError): RVData(t_obj, rv, cov[:-1]) - bad_cov = np.arange(8).reshape((2, 2, 2)) * (u.km/u.s)**2 + bad_cov = np.arange(8).reshape((2, 2, 2)) * (u.km / u.s) ** 2 with pytest.raises(ValueError): RVData(t_obj, rv, bad_cov) @@ -157,10 +158,8 @@ def test_rvdata_init(): RVData(t_arr, rv, err, t_ref=t_arr[3]) -@pytest.mark.parametrize("inputs", - get_valid_input()[0]) +@pytest.mark.parametrize("inputs", get_valid_input()[0]) def test_data_methods(tmpdir, inputs): - # check that copy works if isinstance(inputs, tuple): data1 = RVData(*inputs) @@ -187,9 +186,9 @@ def test_data_methods(tmpdir, inputs): ts = data1.to_timeseries() assert isinstance(ts, TimeSeries) - filename = str(tmpdir / 'test.hdf5') + filename = str(tmpdir / "test.hdf5") with warnings.catch_warnings(): - warnings.simplefilter('ignore', category=UserWarning) + warnings.simplefilter("ignore", category=UserWarning) ts.write(filename, serialize_meta=True) data2 = RVData.from_timeseries(filename) assert u.allclose(data1.t.mjd, data2.t.mjd) @@ -203,11 +202,11 @@ def test_data_methods(tmpdir, inputs): assert len(warns) != 0 # get phase from data object - phase1 = data1.phase(P=15.*u.day) + phase1 = data1.phase(P=15.0 * u.day) assert phase1.min() >= 0 assert phase1.max() <= 1 - phase2 = data1.phase(P=15.*u.day, t0=Time(58585.24, format='mjd')) + phase2 = data1.phase(P=15.0 * u.day, t0=Time(58585.24, format="mjd")) assert not np.allclose(phase1, phase2) # compute inverse variance @@ -222,43 +221,39 @@ def test_guess_from_table(): """NOTE: this is not an exhaustive set of tests, but at least checks a few common cases""" - for rv_name in ['rv', 'vr', 'radial_velocity']: + for rv_name in ["rv", "vr", "radial_velocity"]: tbl = Table() - tbl['t'] = np.linspace(56423.234, 59324.342, 16) * u.day - tbl[rv_name] = np.random.normal(0, 1, len(tbl['t'])) - tbl[f'{rv_name}_err'] = np.random.uniform(0.1, 0.2, len(tbl['t'])) - data = RVData.guess_from_table(tbl, rv_unit=u.km/u.s) - assert np.allclose(data.t.utc.mjd, tbl['t']) + tbl["t"] = np.linspace(56423.234, 59324.342, 16) * u.day + tbl[rv_name] = np.random.normal(0, 1, len(tbl["t"])) + tbl[f"{rv_name}_err"] = np.random.uniform(0.1, 0.2, len(tbl["t"])) + data = RVData.guess_from_table(tbl, rv_unit=u.km / u.s) + assert np.allclose(data.t.utc.mjd, tbl["t"]) if HAS_FUZZY: with warnings.catch_warnings(): - warnings.simplefilter('ignore', category=UserWarning) - for rv_name in ['VHELIO', 'VHELIO_AVG', 'vr', 'vlos']: + warnings.simplefilter("ignore", category=UserWarning) + for rv_name in ["VHELIO", "VHELIO_AVG", "vr", "vlos"]: tbl = Table() - tbl['t'] = np.linspace(56423.234, 59324.342, 16) * u.day - tbl[rv_name] = np.random.normal(0, 1, len(tbl['t'])) - tbl[f'{rv_name}_err'] = np.random.uniform(0.1, 0.2, - len(tbl['t'])) - data = RVData.guess_from_table(tbl, rv_unit=u.km/u.s, - fuzzy=True) - assert np.allclose(data.t.utc.mjd, tbl['t']) + tbl["t"] = np.linspace(56423.234, 59324.342, 16) * u.day + tbl[rv_name] = np.random.normal(0, 1, len(tbl["t"])) + tbl[f"{rv_name}_err"] = np.random.uniform(0.1, 0.2, len(tbl["t"])) + data = RVData.guess_from_table(tbl, rv_unit=u.km / u.s, fuzzy=True) + assert np.allclose(data.t.utc.mjd, tbl["t"]) tbl = Table() - tbl['t'] = np.linspace(2456423.234, 2459324.342, 16) * u.day - tbl['rv'] = np.random.normal(0, 1, len(tbl['t'])) * u.km/u.s - tbl['rv_err'] = np.random.uniform(0.1, 0.2, len(tbl['t'])) * u.km/u.s + tbl["t"] = np.linspace(2456423.234, 2459324.342, 16) * u.day + tbl["rv"] = np.random.normal(0, 1, len(tbl["t"])) * u.km / u.s + tbl["rv_err"] = np.random.uniform(0.1, 0.2, len(tbl["t"])) * u.km / u.s data = RVData.guess_from_table(tbl) - assert np.allclose(data.t.utc.jd, tbl['t']) + assert np.allclose(data.t.utc.jd, tbl["t"]) - data = RVData.guess_from_table(tbl, time_kwargs=dict(scale='tcb')) - assert np.allclose(data.t.tcb.jd, tbl['t']) + data = RVData.guess_from_table(tbl, time_kwargs=dict(scale="tcb")) + assert np.allclose(data.t.tcb.jd, tbl["t"]) -@pytest.mark.skipif(not HAS_MPL, reason='matplotlib not installed') -@pytest.mark.parametrize("inputs", - get_valid_input()[0]) +@pytest.mark.skipif(not HAS_MPL, reason="matplotlib not installed") +@pytest.mark.parametrize("inputs", get_valid_input()[0]) def test_plotting(inputs): - # check that copy works if isinstance(inputs, tuple): data = RVData(*inputs) @@ -268,98 +263,98 @@ def test_plotting(inputs): data.plot() # style - data.plot(color='r') + data.plot(color="r") # custom axis fig, ax = plt.subplots(1, 1) data.plot(ax=plt.gca()) # formatting - data.plot(rv_unit=u.m/u.s) - data.plot(rv_unit=u.m/u.s, time_format='jd') - data.plot(rv_unit=u.m/u.s, time_format=lambda x: x.utc.mjd) - data.plot(ecolor='r') + data.plot(rv_unit=u.m / u.s) + data.plot(rv_unit=u.m / u.s, time_format="jd") + data.plot(rv_unit=u.m / u.s, time_format=lambda x: x.utc.mjd) + data.plot(ecolor="r") - plt.close('all') + plt.close("all") def test_multi_data(): - import exoplanet.units as xu - import pymc3 as pm + import pymc as pm + + import thejoker.units as xu - rnd = DEFAULT_RNG(42) + rnd = np.random.default_rng(42) # Set up mulitple valid data objects: _, raw1 = get_valid_input(rnd=rnd) - data1 = RVData(raw1['t_obj'], raw1['rv'], raw1['err']) + data1 = RVData(raw1["t_obj"], raw1["rv"], raw1["err"]) _, raw2 = get_valid_input(rnd=rnd, size=8) - data2 = RVData(raw2['t_obj'], raw2['rv'], raw2['err']) + data2 = RVData(raw2["t_obj"], raw2["rv"], raw2["err"]) _, raw3 = get_valid_input(rnd=rnd, size=4) - data3 = RVData(raw3['t_obj'], raw3['rv'], raw3['err']) + data3 = RVData(raw3["t_obj"], raw3["rv"], raw3["err"]) - prior1 = JokerPrior.default(1*u.day, 1*u.year, - 25*u.km/u.s, - sigma_v=100*u.km/u.s) + prior1 = JokerPrior.default( + 1 * u.day, 1 * u.year, 25 * u.km / u.s, sigma_v=100 * u.km / u.s + ) # Object should return input: - multi_data, ids, trend_M = validate_prepare_data(data1, - prior1.poly_trend, - prior1.n_offsets) + multi_data, ids, trend_M = validate_prepare_data( + data1, prior1.poly_trend, prior1.n_offsets + ) assert np.allclose(multi_data.rv.value, data1.rv.value) assert np.all(ids == 0) - assert np.allclose(trend_M[:, 0], 1.) + assert np.allclose(trend_M[:, 0], 1.0) # Three valid objects as a list: with pm.Model(): - dv1 = xu.with_unit(pm.Normal('dv0_1', 0, 1.), - u.km/u.s) - dv2 = xu.with_unit(pm.Normal('dv0_2', 4, 5.), - u.km/u.s) - prior2 = JokerPrior.default(1*u.day, 1*u.year, - 25*u.km/u.s, - sigma_v=100*u.km/u.s, - v0_offsets=[dv1, dv2]) + dv1 = xu.with_unit(pm.Normal("dv0_1", 0, 1.0), u.km / u.s) + dv2 = xu.with_unit(pm.Normal("dv0_2", 4, 5.0), u.km / u.s) + prior2 = JokerPrior.default( + 1 * u.day, + 1 * u.year, + 25 * u.km / u.s, + sigma_v=100 * u.km / u.s, + v0_offsets=[dv1, dv2], + ) datas = [data1, data2, data3] - multi_data, ids, trend_M = validate_prepare_data(datas, - prior2.poly_trend, - prior2.n_offsets) + multi_data, ids, trend_M = validate_prepare_data( + datas, prior2.poly_trend, prior2.n_offsets + ) assert len(np.unique(ids)) == 3 assert len(multi_data) == sum([len(d) for d in datas]) assert 0 in ids and 1 in ids and 2 in ids - assert np.allclose(trend_M[:, 0], 1.) + assert np.allclose(trend_M[:, 0], 1.0) # Three valid objects with names: - datas = {'apogee': data1, 'lamost': data2, 'weave': data3} - multi_data, ids, trend_M = validate_prepare_data(datas, - prior2.poly_trend, - prior2.n_offsets) + datas = {"apogee": data1, "lamost": data2, "weave": data3} + multi_data, ids, trend_M = validate_prepare_data( + datas, prior2.poly_trend, prior2.n_offsets + ) assert len(np.unique(ids)) == 3 assert len(multi_data) == sum([len(d) for d in datas.values()]) - assert 'apogee' in ids and 'lamost' in ids and 'weave' in ids - assert np.allclose(trend_M[:, 0], 1.) + assert "apogee" in ids and "lamost" in ids and "weave" in ids + assert np.allclose(trend_M[:, 0], 1.0) # Check it fails if n_offsets != number of data sources with pytest.raises(ValueError): - validate_prepare_data(datas, - prior1.poly_trend, - prior1.n_offsets) + validate_prepare_data(datas, prior1.poly_trend, prior1.n_offsets) with pytest.raises(ValueError): - validate_prepare_data(data1, - prior2.poly_trend, - prior2.n_offsets) + validate_prepare_data(data1, prior2.poly_trend, prior2.n_offsets) # Check that this fails if one has a covariance matrix - data_cov = RVData(raw3['t_obj'], raw3['rv'], raw3['cov']) + data_cov = RVData(raw3["t_obj"], raw3["rv"], raw3["cov"]) with pytest.raises(NotImplementedError): - validate_prepare_data({'apogee': data1, 'test': data2, - 'weave': data_cov}, - prior2.poly_trend, prior2.n_offsets) + validate_prepare_data( + {"apogee": data1, "test": data2, "weave": data_cov}, + prior2.poly_trend, + prior2.n_offsets, + ) with pytest.raises(NotImplementedError): - validate_prepare_data([data1, data2, data_cov], - prior2.poly_trend, - prior2.n_offsets) + validate_prepare_data( + [data1, data2, data_cov], prior2.poly_trend, prior2.n_offsets + ) diff --git a/thejoker/tests/test_likelihood_helpers.py b/thejoker/tests/test_likelihood_helpers.py index 2d153519..e8c24b68 100644 --- a/thejoker/tests/test_likelihood_helpers.py +++ b/thejoker/tests/test_likelihood_helpers.py @@ -3,12 +3,11 @@ from ..data import RVData from ..data_helpers import validate_prepare_data from .test_data import get_valid_input -from ..utils import DEFAULT_RNG def test_design_matrix(): # implicitly tests get_constant_term_design_matrix - rnd = DEFAULT_RNG(42) + rnd = np.random.default_rng(42) # Set up mulitple valid data objects: ndata1 = 8 @@ -16,22 +15,22 @@ def test_design_matrix(): ndata3 = 3 _, raw1 = get_valid_input(rnd=rnd, size=ndata1) - data1 = RVData(raw1['t_obj'], raw1['rv'], raw1['err']) + data1 = RVData(raw1["t_obj"], raw1["rv"], raw1["err"]) _, raw2 = get_valid_input(rnd=rnd, size=ndata2) - data2 = RVData(raw2['t_obj'], raw2['rv'], raw2['err']) + data2 = RVData(raw2["t_obj"], raw2["rv"], raw2["err"]) _, raw3 = get_valid_input(rnd=rnd, size=ndata3) - data3 = RVData(raw3['t_obj'], raw3['rv'], raw3['err']) + data3 = RVData(raw3["t_obj"], raw3["rv"], raw3["err"]) data, ids, M = validate_prepare_data([data1, data2, data3], 1, 2) - assert np.allclose(M[:, 0], 1.) + assert np.allclose(M[:, 0], 1.0) idx = np.arange(len(data), dtype=int) - mask = (idx >= ndata1) & (idx < (ndata1+ndata2)) - assert np.allclose(M[mask, 1], 1.) - assert np.allclose(M[~mask, 1], 0.) + mask = (idx >= ndata1) & (idx < (ndata1 + ndata2)) + assert np.allclose(M[mask, 1], 1.0) + assert np.allclose(M[~mask, 1], 0.0) - mask = (idx >= (ndata1+ndata2)) & (idx < (ndata1+ndata2+ndata3)) - assert np.allclose(M[mask, 2], 1.) - assert np.allclose(M[~mask, 2], 0.) + mask = (idx >= (ndata1 + ndata2)) & (idx < (ndata1 + ndata2 + ndata3)) + assert np.allclose(M[mask, 2], 1.0) + assert np.allclose(M[~mask, 2], 0.0) diff --git a/thejoker/tests/test_prior.py b/thejoker/tests/test_prior.py index dbd06eb6..dd20b231 100644 --- a/thejoker/tests/test_prior.py +++ b/thejoker/tests/test_prior.py @@ -1,28 +1,34 @@ # Third-party import astropy.units as u import numpy as np +import pymc as pm import pytest -import pymc3 as pm -import exoplanet.units as xu + +import thejoker.units as xu # Project -from ..prior import JokerPrior, default_nonlinear_prior, default_linear_prior +from ..prior import JokerPrior, default_linear_prior, default_nonlinear_prior def get_prior(case=None): - default_expected_units = {'P': u.day, 'e': u.one, - 'omega': u.radian, 'M0': u.radian, - 's': u.m/u.s, - 'K': u.km/u.s, 'v0': u.km/u.s, - 'ln_prior': u.one} + default_expected_units = { + "P": u.day, + "e": u.one, + "omega": u.radian, + "M0": u.radian, + "s": u.m / u.s, + "K": u.km / u.s, + "v0": u.km / u.s, + "ln_prior": u.one, + } if case == 0: # Default prior with standard parameters: with pm.Model() as model: - default_nonlinear_prior(P_min=1*u.day, P_max=1*u.year) - default_linear_prior(sigma_K0=25*u.km/u.s, - P0=1*u.year, - sigma_v=10*u.km/u.s) + default_nonlinear_prior(P_min=1 * u.day, P_max=1 * u.year) + default_linear_prior( + sigma_K0=25 * u.km / u.s, P0=1 * u.year, sigma_v=10 * u.km / u.s + ) prior = JokerPrior(model=model) return prior, default_expected_units @@ -31,14 +37,13 @@ def get_prior(case=None): # Replace a nonlinear parameter units = default_expected_units.copy() with pm.Model() as model: - P = xu.with_unit(pm.Normal('P', 10, 0.5), - u.year) - units['P'] = u.year + P = xu.with_unit(pm.Normal("P", 10, 0.5), u.year) + units["P"] = u.year - default_nonlinear_prior(pars={'P': P}) - default_linear_prior(sigma_K0=25*u.km/u.s, - P0=1*u.year, - sigma_v=10*u.km/u.s) + default_nonlinear_prior(pars={"P": P}) + default_linear_prior( + sigma_K0=25 * u.km / u.s, P0=1 * u.year, sigma_v=10 * u.km / u.s + ) prior = JokerPrior(model=model) @@ -48,13 +53,11 @@ def get_prior(case=None): # Replace a linear parameter units = default_expected_units.copy() with pm.Model() as model: - K = xu.with_unit(pm.Normal('K', 10, 0.5), - u.m/u.s) - units['K'] = u.m/u.s + K = xu.with_unit(pm.Normal("K", 10, 0.5), u.m / u.s) + units["K"] = u.m / u.s - default_nonlinear_prior(P_min=1*u.day, P_max=1*u.year) - default_linear_prior(sigma_v=10*u.km/u.s, - pars={'K': K}) + default_nonlinear_prior(P_min=1 * u.day, P_max=1 * u.year) + default_linear_prior(sigma_v=10 * u.km / u.s, pars={"K": K}) prior = JokerPrior(model=model) @@ -63,10 +66,10 @@ def get_prior(case=None): elif case == 3: # Pass pars instead of relying on model with pm.Model() as model: - nl_pars = default_nonlinear_prior(P_min=1*u.day, P_max=1*u.year) - l_pars = default_linear_prior(sigma_K0=25*u.km/u.s, - P0=1*u.year, - sigma_v=10*u.km/u.s) + nl_pars = default_nonlinear_prior(P_min=1 * u.day, P_max=1 * u.year) + l_pars = default_linear_prior( + sigma_K0=25 * u.km / u.s, P0=1 * u.year, sigma_v=10 * u.km / u.s + ) pars = {**nl_pars, **l_pars} prior = JokerPrior(pars=pars) @@ -76,104 +79,113 @@ def get_prior(case=None): # Try with more poly_trends units = default_expected_units.copy() with pm.Model() as model: - nl_pars = default_nonlinear_prior(P_min=1*u.day, P_max=1*u.year) - l_pars = default_linear_prior(sigma_K0=25*u.km/u.s, - P0=1*u.year, - sigma_v=[10*u.km/u.s, - 0.1*u.km/u.s/u.year], - poly_trend=2) + nl_pars = default_nonlinear_prior(P_min=1 * u.day, P_max=1 * u.year) + l_pars = default_linear_prior( + sigma_K0=25 * u.km / u.s, + P0=1 * u.year, + sigma_v=[10 * u.km / u.s, 0.1 * u.km / u.s / u.year], + poly_trend=2, + ) pars = {**nl_pars, **l_pars} prior = JokerPrior(pars=pars) - units['v1'] = u.km/u.s/u.year + units["v1"] = u.km / u.s / u.year return prior, units elif case == 5: # Default prior with .default() - prior = JokerPrior.default(P_min=1*u.day, P_max=10*u.year, - sigma_K0=25*u.km/u.s, - P0=1*u.year, - sigma_v=10*u.km/u.s) + prior = JokerPrior.default( + P_min=1 * u.day, + P_max=10 * u.year, + sigma_K0=25 * u.km / u.s, + P0=1 * u.year, + sigma_v=10 * u.km / u.s, + ) return prior, default_expected_units elif case == 6: # poly_trend with .default() units = default_expected_units.copy() - prior = JokerPrior.default(P_min=1*u.day, P_max=1*u.year, - sigma_K0=25*u.km/u.s, - P0=1*u.year, - sigma_v=[10*u.km/u.s, - 0.1*u.km/u.s/u.year], - poly_trend=2) - units['v1'] = u.km/u.s/u.year + prior = JokerPrior.default( + P_min=1 * u.day, + P_max=1 * u.year, + sigma_K0=25 * u.km / u.s, + P0=1 * u.year, + sigma_v=[10 * u.km / u.s, 0.1 * u.km / u.s / u.year], + poly_trend=2, + ) + units["v1"] = u.km / u.s / u.year return prior, units elif case == 7: # Replace a linear parameter with .default() units = default_expected_units.copy() with pm.Model() as model: - K = xu.with_unit(pm.Normal('K', 10, 0.5), - u.m/u.s) - units['K'] = u.m/u.s - - prior = JokerPrior.default(P_min=1*u.day, P_max=1*u.year, - sigma_v=100*u.km/u.s, - pars={'K': K}) + K = xu.with_unit(pm.Normal("K", 10, 0.5), u.m / u.s) + units["K"] = u.m / u.s + + prior = JokerPrior.default( + P_min=1 * u.day, + P_max=1 * u.year, + sigma_v=100 * u.km / u.s, + pars={"K": K}, + ) return prior, units elif case == 8: # Replace s with pymc3 var with .default() units = default_expected_units.copy() with pm.Model() as model: - s = xu.with_unit(pm.Normal('s', 10, 0.5), - u.m/u.s) - units['s'] = u.m/u.s - - prior = JokerPrior.default(P_min=1*u.day, P_max=1*u.year, - sigma_K0=25*u.km/u.s, - sigma_v=100*u.km/u.s, - s=s) + s = xu.with_unit(pm.Normal("s", 10, 0.5), u.m / u.s) + units["s"] = u.m / u.s + + prior = JokerPrior.default( + P_min=1 * u.day, + P_max=1 * u.year, + sigma_K0=25 * u.km / u.s, + sigma_v=100 * u.km / u.s, + s=s, + ) return prior, units return 9 # number of cases above -@pytest.mark.parametrize('case', range(get_prior())) +@pytest.mark.parametrize("case", range(get_prior())) def test_init_sample(case): prior, expected_units = get_prior(case) for k in expected_units.keys(): - if k == 'ln_prior': # skip + if k == "ln_prior": # skip continue assert k in prior.model.named_vars samples = prior.sample() for k in samples.par_names: - assert hasattr(samples[k], 'unit') + assert hasattr(samples[k], "unit") assert samples[k].unit == expected_units[k] samples = prior.sample(size=10) for k in samples.par_names: - assert hasattr(samples[k], 'unit') + assert hasattr(samples[k], "unit") assert samples[k].unit == expected_units[k] samples = prior.sample(size=10, generate_linear=True) for k in samples.par_names: - assert hasattr(samples[k], 'unit') + assert hasattr(samples[k], "unit") assert samples[k].unit == expected_units[k] samples = prior.sample(size=10, return_logprobs=True) - assert 'ln_prior' in samples.par_names + assert "ln_prior" in samples.par_names for k in samples.par_names: - assert hasattr(samples[k], 'unit') + assert hasattr(samples[k], "unit") assert samples[k].unit == expected_units[k] - samples = prior.sample(size=10, generate_linear=True, - return_logprobs=True) - assert 'ln_prior' in samples.par_names + samples = prior.sample(size=10, generate_linear=True, return_logprobs=True) + assert "ln_prior" in samples.par_names for k in samples.par_names: - assert hasattr(samples[k], 'unit') + assert hasattr(samples[k], "unit") assert samples[k].unit == expected_units[k] diff --git a/thejoker/tests/test_utils.py b/thejoker/tests/test_utils.py index e9a83749..207ae03a 100644 --- a/thejoker/tests/test_utils.py +++ b/thejoker/tests/test_utils.py @@ -1,17 +1,21 @@ # Third-party -from astropy.table import QTable import astropy.units as u -from astropy.io.misc.hdf5 import meta_path import h5py import numpy as np -import tables as tb +from astropy.io.misc.hdf5 import meta_path +from astropy.table import QTable # Package from ..samples import JokerSamples -from ..utils import (batch_tasks, table_header_to_units, - table_contains_column, - read_batch, read_batch_slice, read_batch_idx, - read_random_batch) +from ..utils import ( + batch_tasks, + read_batch, + read_batch_idx, + read_batch_slice, + read_random_batch, + table_contains_column, + table_header_to_units, +) def test_batch_tasks(): @@ -19,113 +23,113 @@ def test_batch_tasks(): start_idx = 1103 tasks = batch_tasks(N, n_batches=16, start_idx=start_idx) assert tasks[0][0][0] == start_idx - assert tasks[-1][0][1] == N+start_idx + assert tasks[-1][0][1] == N + start_idx # try with an array: - tasks = batch_tasks(N, n_batches=16, start_idx=start_idx, - arr=np.random.random(size=8*N)) + tasks = batch_tasks( + N, n_batches=16, start_idx=start_idx, arr=np.random.random(size=8 * N) + ) n_tasks = sum([tasks[i][0].size for i in range(len(tasks))]) assert n_tasks == N def test_table_header_to_units(tmpdir): - filename = str(tmpdir / 'test.hdf5') + filename = str(tmpdir / "test.hdf5") tbl = QTable() - tbl['a'] = np.arange(10) * u.kpc - tbl['b'] = np.arange(10) * u.km/u.s - tbl['c'] = np.arange(10) * u.day - tbl.write(filename, path='test', serialize_meta=True) + tbl["a"] = np.arange(10) * u.kpc + tbl["b"] = np.arange(10) * u.km / u.s + tbl["c"] = np.arange(10) * u.day + tbl.write(filename, path="test", serialize_meta=True) # TODO: pytables doesn't support variable length strings # with tb.open_file(filename, mode='r') as f: # units = table_header_to_units(f.root[meta_path('test')]) - with h5py.File(filename, mode='r') as f: - units = table_header_to_units(f[meta_path('test')]) + with h5py.File(filename, mode="r") as f: + units = table_header_to_units(f[meta_path("test")]) for col in tbl.colnames: assert tbl[col].unit == units[col] def test_table_contains_column(tmpdir): - filename = str(tmpdir / 'test.hdf5') + filename = str(tmpdir / "test.hdf5") tbl = QTable() - tbl['a'] = np.arange(10) * u.kpc - tbl['b'] = np.arange(10) * u.km/u.s - tbl['c'] = np.arange(10) * u.day + tbl["a"] = np.arange(10) * u.kpc + tbl["b"] = np.arange(10) * u.km / u.s + tbl["c"] = np.arange(10) * u.day tbl.write(filename, path=JokerSamples._hdf5_path, serialize_meta=True) # TODO: pytables doesn't support variable length strings # with tb.open_file(filename, mode='r') as f: - with h5py.File(filename, mode='r') as f: - assert table_contains_column(f, 'a') - assert table_contains_column(f, 'b') - assert table_contains_column(f, 'c') - assert not table_contains_column(f, 'd') + with h5py.File(filename, mode="r") as f: + assert table_contains_column(f, "a") + assert table_contains_column(f, "b") + assert table_contains_column(f, "c") + assert not table_contains_column(f, "d") def test_read_batch_slice(tmpdir): - filename = str(tmpdir / 'test.hdf5') + filename = str(tmpdir / "test.hdf5") tbl = QTable() - tbl['a'] = np.arange(100) * u.kpc - tbl['b'] = np.arange(100) * u.km/u.s - tbl['c'] = np.arange(100) * u.day + tbl["a"] = np.arange(100) * u.kpc + tbl["b"] = np.arange(100) * u.km / u.s + tbl["c"] = np.arange(100) * u.day tbl.write(filename, path=JokerSamples._hdf5_path, serialize_meta=True) for read_func in [read_batch_slice, read_batch]: - batch = read_func(filename, ['a', 'b'], slice(10, 20)) + batch = read_func(filename, ["a", "b"], slice(10, 20)) assert batch.shape == (10, 2) - assert np.allclose(batch[:, 0], tbl['a'].value[10:20]) - assert np.allclose(batch[:, 1], tbl['b'].value[10:20]) + assert np.allclose(batch[:, 0], tbl["a"].value[10:20]) + assert np.allclose(batch[:, 1], tbl["b"].value[10:20]) - batch = read_func(filename, ['b', 'c'], slice(0, 100), - units={'b': u.kpc/u.Myr}) + batch = read_func( + filename, ["b", "c"], slice(0, 100), units={"b": u.kpc / u.Myr} + ) assert batch.shape == (100, 2) - assert np.allclose(batch[:, 0], tbl['b'].to_value(u.kpc/u.Myr)) - assert np.allclose(batch[:, 1], tbl['c'].value) + assert np.allclose(batch[:, 0], tbl["b"].to_value(u.kpc / u.Myr)) + assert np.allclose(batch[:, 1], tbl["c"].value) - batch = read_func(filename, ['b', 'c'], slice(0, 100, 2)) + batch = read_func(filename, ["b", "c"], slice(0, 100, 2)) assert batch.shape == (50, 2) def test_read_batch_idx(tmpdir): - filename = str(tmpdir / 'test.hdf5') + filename = str(tmpdir / "test.hdf5") tbl = QTable() - tbl['a'] = np.arange(100) * u.kpc - tbl['b'] = np.arange(100) * u.km/u.s - tbl['c'] = np.arange(100) * u.day + tbl["a"] = np.arange(100) * u.kpc + tbl["b"] = np.arange(100) * u.km / u.s + tbl["c"] = np.arange(100) * u.day tbl.write(filename, path=JokerSamples._hdf5_path, serialize_meta=True) for read_func in [read_batch_idx, read_batch]: idx = np.arange(10, 20, 1) - batch = read_func(filename, ['a', 'b'], idx, units=None) + batch = read_func(filename, ["a", "b"], idx, units=None) assert batch.shape == (len(idx), 2) - assert np.allclose(batch[:, 0], tbl['a'].value[idx]) - assert np.allclose(batch[:, 1], tbl['b'].value[idx]) + assert np.allclose(batch[:, 0], tbl["a"].value[idx]) + assert np.allclose(batch[:, 1], tbl["b"].value[idx]) - batch = read_func(filename, ['b', 'c'], idx, - units={'b': u.kpc/u.Myr}) + batch = read_func(filename, ["b", "c"], idx, units={"b": u.kpc / u.Myr}) assert batch.shape == (len(idx), 2) - assert np.allclose(batch[:, 0], tbl['b'].to_value(u.kpc/u.Myr)[idx]) - assert np.allclose(batch[:, 1], tbl['c'].value[idx]) + assert np.allclose(batch[:, 0], tbl["b"].to_value(u.kpc / u.Myr)[idx]) + assert np.allclose(batch[:, 1], tbl["c"].value[idx]) def test_read_random_batch(tmpdir): - filename = str(tmpdir / 'test.hdf5') + filename = str(tmpdir / "test.hdf5") tbl = QTable() - tbl['a'] = np.arange(100) * u.kpc - tbl['b'] = np.arange(100) * u.km/u.s - tbl['c'] = np.arange(100) * u.day + tbl["a"] = np.arange(100) * u.kpc + tbl["b"] = np.arange(100) * u.km / u.s + tbl["c"] = np.arange(100) * u.day tbl.write(filename, path=JokerSamples._hdf5_path, serialize_meta=True) for read_func in [read_random_batch, read_batch]: - batch = read_func(filename, ['a', 'b'], 10, units=None) + batch = read_func(filename, ["a", "b"], 10, units=None) assert batch.shape == (10, 2) - batch = read_func(filename, ['b', 'c'], 100, - units={'b': u.kpc/u.Myr}) + batch = read_func(filename, ["b", "c"], 100, units={"b": u.kpc / u.Myr}) assert batch.shape == (100, 2) From 762ed6da4cad2449b44b9ddc316b01883c23941b Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Sun, 3 Mar 2024 18:33:59 -0500 Subject: [PATCH 15/50] working on getting tests to pass again --- thejoker/data.py | 273 +++++++++++---------- thejoker/exceptions.py | 6 - thejoker/samples.py | 262 +++++++++++--------- thejoker/src/tests/py_likelihood.py | 13 +- thejoker/src/tests/test_fast_likelihood.py | 20 +- thejoker/tests/test_data.py | 12 +- thejoker/tests/test_prior.py | 2 +- 7 files changed, 311 insertions(+), 277 deletions(-) delete mode 100644 thejoker/exceptions.py diff --git a/thejoker/data.py b/thejoker/data.py index 48fffffd..ea073bb1 100644 --- a/thejoker/data.py +++ b/thejoker/data.py @@ -1,18 +1,16 @@ -import warnings - # Third-party +import astropy.units as u +import numpy as np from astropy.table import Table from astropy.time import Time from astropy.utils.decorators import deprecated_renamed_argument -import astropy.units as u -import numpy as np + +from .data_helpers import guess_time_format # Project from .logging import logger -from .data_helpers import guess_time_format -from .exceptions import TheJokerDeprecationWarning -__all__ = ['RVData'] +__all__ = ["RVData"] class RVData: @@ -37,11 +35,9 @@ class RVData: Filter out any NaN or Inf data points. """ - @deprecated_renamed_argument('t0', 't_ref', since='v1.2', - warning_type=TheJokerDeprecationWarning) - @u.quantity_input(rv=u.km/u.s, rv_err=[u.km/u.s, (u.km/u.s)**2]) - def __init__(self, t, rv, rv_err, t_ref=None, clean=True): + @u.quantity_input(rv=u.km / u.s, rv_err=[u.km / u.s, (u.km / u.s) ** 2]) + def __init__(self, t, rv, rv_err, t_ref=None, clean=True): # For speed, time is saved internally as BMJD: if isinstance(t, Time): _t_bmjd = t.tcb.mjd @@ -59,22 +55,26 @@ def __init__(self, t, rv, rv_err, t_ref=None, clean=True): elif self.rv_err.ndim == 2: self._has_cov = True - if (self.rv_err.shape != (self.rv.size, self.rv.size) - and self.rv_err.shape != (self.rv.size, )): - raise ValueError(f"Invalid shape for input RV error " - f"{self.rv_err.shape}. Should either be " - f"({self.rv.size},) or " - f"({self.rv.size}, {self.rv.size})") + if self.rv_err.shape != (self.rv.size, self.rv.size) and self.rv_err.shape != ( + self.rv.size, + ): + raise ValueError( + f"Invalid shape for input RV error " + f"{self.rv_err.shape}. Should either be " + f"({self.rv.size},) or " + f"({self.rv.size}, {self.rv.size})" + ) # make sure shapes are consistent if self._t_bmjd.shape != self.rv.shape: - raise ValueError(f"Shape of input times and RVs must be consistent" - f" ({self._t_bmjd.shape} vs {self.rv.shape})") + raise ValueError( + f"Shape of input times and RVs must be consistent" + f" ({self._t_bmjd.shape} vs {self.rv.shape})" + ) if clean: # filter out NAN or INF data points - idx = (np.isfinite(self._t_bmjd) - & np.isfinite(self.rv)) + idx = np.isfinite(self._t_bmjd) & np.isfinite(self.rv) if self._has_cov: idx &= np.isfinite(self.rv_err).all(axis=0) @@ -106,26 +106,21 @@ def __init__(self, t, rv, rv_err, t_ref=None, clean=True): if t_ref is False: self.t_ref = None - self._t_ref_bmjd = 0. + self._t_ref_bmjd = 0.0 else: if t_ref is None: t_ref = self.t.min() if not isinstance(t_ref, Time): - raise TypeError('If a reference time t_ref is specified, it ' - 'must be an astropy.time.Time object.') + raise TypeError( + "If a reference time t_ref is specified, it " + "must be an astropy.time.Time object." + ) self.t_ref = t_ref self._t_ref_bmjd = self.t_ref.tcb.mjd - @property - def t0(self): - warnings.warn('The argument and attribute "t0" has been renamed ' - 'and should now be specified / accessed as "t_ref"', - TheJokerDeprecationWarning) - return self.t_ref - # ------------------------------------------------------------------------ # Computed or convenience properties @@ -138,7 +133,7 @@ def t(self): t : `~astropy.time.Time` An Astropy Time object for all times. """ - return Time(self._t_bmjd, scale='tcb', format='mjd') + return Time(self._t_bmjd, scale="tcb", format="mjd") @property def cov(self): @@ -154,16 +149,15 @@ def ivar(self): if self._has_cov: return np.linalg.inv(self.rv_err.value) / self.rv_err.unit else: - return 1 / self.rv_err ** 2 + return 1 / self.rv_err**2 # ------------------------------------------------------------------------ # Other initialization methods: @classmethod - @deprecated_renamed_argument('t0', 't_ref', since='v1.2', - warning_type=TheJokerDeprecationWarning) - def guess_from_table(cls, tbl, time_kwargs=None, rv_unit=None, - fuzzy=False, t_ref=None): + def guess_from_table( + cls, tbl, time_kwargs=None, rv_unit=None, fuzzy=False, t_ref=None + ): """ Try to construct an ``RVData`` instance by guessing column names from the input table. @@ -201,49 +195,57 @@ def guess_from_table(cls, tbl, time_kwargs=None, rv_unit=None, # First check for any of the valid astropy Time format names: # FUTURETODO: right now we only support jd and mjd (and b-preceding) - for fmt in ['jd', 'mjd']: + for fmt in ["jd", "mjd"]: if fmt in lwr_cols: - time_kwargs['format'] = time_kwargs.get('format', fmt) + time_kwargs["format"] = time_kwargs.get("format", fmt) time_data = tbl[lwr_to_col[fmt]] break - elif f'b{fmt}' in lwr_cols: - time_kwargs['format'] = time_kwargs.get('format', fmt) - time_kwargs['scale'] = time_kwargs.get('scale', 'tcb') - time_data = tbl[lwr_to_col[f'b{fmt}']] + elif f"b{fmt}" in lwr_cols: + time_kwargs["format"] = time_kwargs.get("format", fmt) + time_kwargs["scale"] = time_kwargs.get("scale", "tcb") + time_data = tbl[lwr_to_col[f"b{fmt}"]] _scale_specified = True break - time_info_msg = ("Assuming time scale is '{}' because it was not " - "specified. To change this, pass in: " - "time_kwargs=dict(scale='...') with whatever time " - "scale your data are in.") - _fmt_specified = 'format' in time_kwargs - _scale_specified = 'scale' in time_kwargs + time_info_msg = ( + "Assuming time scale is '{}' because it was not " + "specified. To change this, pass in: " + "time_kwargs=dict(scale='...') with whatever time " + "scale your data are in." + ) + _fmt_specified = "format" in time_kwargs + _scale_specified = "scale" in time_kwargs # check colnames for "t" or "time" - for name in ['t', 'time']: + for name in ["t", "time"]: if name in lwr_cols: time_data = tbl[lwr_to_col[name]] - time_kwargs['format'] = time_kwargs.get( - 'format', guess_time_format(tbl[lwr_to_col[name]])) + time_kwargs["format"] = time_kwargs.get( + "format", guess_time_format(tbl[lwr_to_col[name]]) + ) if not _fmt_specified: - logger.info("Guessed time format: '{}'. If this is " - "incorrect, try passing in " - "time_kwargs=dict(format='...') with the " - "correct format, and open an issue at " - "https://github.com/adrn/thejoker/issues" - .format(time_kwargs['format'])) + logger.info( + "Guessed time format: '{}'. If this is " + "incorrect, try passing in " + "time_kwargs=dict(format='...') with the " + "correct format, and open an issue at " + "https://github.com/adrn/thejoker/issues".format( + time_kwargs["format"] + ) + ) break if not _scale_specified: - logger.info(time_info_msg.format(time_kwargs.get('scale', 'utc'))) + logger.info(time_info_msg.format(time_kwargs.get("scale", "utc"))) if time_data is None: - raise RuntimeWarning("Failed to parse time data and format from " - "input table. Instead, try using the " - "initializer directly, and specify the time " - "as an astropy.time.Time instance.") + raise RuntimeWarning( + "Failed to parse time data and format from " + "input table. Instead, try using the " + "initializer directly, and specify the time " + "as an astropy.time.Time instance." + ) time = Time(time_data, **time_kwargs) @@ -251,16 +253,17 @@ def guess_from_table(cls, tbl, time_kwargs=None, rv_unit=None, # Now deal with RV data: # FUTURETODO: could make this customizable... - _valid_rv_names = ['rv', 'vr', 'radial_velocity', - 'vhelio', 'vrad', 'vlos'] + _valid_rv_names = ["rv", "vr", "radial_velocity", "vhelio", "vrad", "vlos"] if fuzzy: try: from fuzzywuzzy import process except ImportError: - raise ImportError("Fuzzy column name matching requires " - "`fuzzywuzzy`. Install with pip install " - "fuzzywuzzy.") + raise ImportError( + "Fuzzy column name matching requires " + "`fuzzywuzzy`. Install with pip install " + "fuzzywuzzy." + ) # FUTURETODO: could make this customizable too... score_thresh = 90 @@ -276,16 +279,19 @@ def guess_from_table(cls, tbl, time_kwargs=None, rv_unit=None, # error if the best match is below threshold if scores.max() < score_thresh: - raise RuntimeError("Failed to parse radial velocity data from " - "input table: No column names looked " - "good with fuzzy name matching.") + raise RuntimeError( + "Failed to parse radial velocity data from " + "input table: No column names looked " + "good with fuzzy name matching." + ) # check for multiple bests: if np.sum(scores == scores.max()) > 1: - raise RuntimeError("Failed to parse radial velocity data from " - "input table: Multiple column names looked " - "good with fuzzy matching {}." - .format(matches[scores == scores.max()])) + raise RuntimeError( + "Failed to parse radial velocity data from " + "input table: Multiple column names looked " + f"good with fuzzy matching {matches[scores == scores.max()]}." + ) best_rv_name = matches[scores.argmax()] @@ -295,25 +301,33 @@ def guess_from_table(cls, tbl, time_kwargs=None, rv_unit=None, best_rv_name = name break else: - raise RuntimeError("Failed to parse radial velocity data from " - "input table: no matches to input names: " - f"{_valid_rv_names}. Use fuzzy=True or " - "use the initializer directly.") + raise RuntimeError( + "Failed to parse radial velocity data from " + "input table: no matches to input names: " + f"{_valid_rv_names}. Use fuzzy=True or " + "use the initializer directly." + ) rv_data = u.Quantity(tbl[lwr_to_col[best_rv_name]]) # FUTURETODO: allow customizing? - _valid_err_names = [f'{best_rv_name}err', f'{best_rv_name}_err', - f'{best_rv_name}_e', f'e_{best_rv_name}'] + _valid_err_names = [ + f"{best_rv_name}err", + f"{best_rv_name}_err", + f"{best_rv_name}_e", + f"e_{best_rv_name}", + ] for err_name in _valid_err_names: if err_name in lwr_cols: err_data = u.Quantity(tbl[lwr_to_col[err_name]]) break else: - raise RuntimeError("Failed to parse radial velocity error data " - "from input table: no matches to input names: " - f"{_valid_err_names}. Try using the " - "initializer directly.") + raise RuntimeError( + "Failed to parse radial velocity error data " + "from input table: no matches to input names: " + f"{_valid_err_names}. Try using the " + "initializer directly." + ) if rv_unit is not None: if rv_data.unit is u.one: @@ -333,25 +347,23 @@ def to_timeseries(self): """ from astropy.timeseries import TimeSeries - ts = TimeSeries(time=self.t, data={'rv': self.rv, - 'rv_err': self.rv_err}) - ts.meta['t_ref'] = self.t_ref + ts = TimeSeries(time=self.t, data={"rv": self.rv, "rv_err": self.rv_err}) + ts.meta["t_ref"] = self.t_ref return ts @classmethod def from_timeseries(cls, f, path=None): from astropy.timeseries import TimeSeries + ts = TimeSeries.read(f, path=path) - t_ref = ts.meta.get('t_ref', None) - return cls(t=ts['time'], - rv=ts['rv'], - rv_err=ts['rv_err'], - t_ref=t_ref) + t_ref = ts.meta.get("t_ref", None) + return cls(t=ts["time"], rv=ts["rv"], rv_err=ts["rv_err"], t_ref=t_ref) # ------------------------------------------------------------------------ # Other methods - def phase(self, P, t0=None): + @deprecated_renamed_argument("t0", "t_ref", "v1.3") + def phase(self, P, t_ref=None): """ Convert time to a phase. @@ -362,7 +374,7 @@ def phase(self, P, t0=None): ---------- P : `~astropy.units.Quantity` [time] The period. - t0 : `~astropy.time.Time` (optional) + t_ref : `~astropy.time.Time` (optional) Default uses the internal reference epoch. Use this to compute the phase relative to some other epoch @@ -372,16 +384,21 @@ def phase(self, P, t0=None): The dimensionless phase of each observation. """ - if t0 is None: - t0 = self.t_ref - return ((self.t - t0) / P) % 1. - - @deprecated_renamed_argument('relative_to_t0', 'relative_to_t_ref', - since='v1.2', - warning_type=TheJokerDeprecationWarning) - def plot(self, ax=None, rv_unit=None, time_format='mjd', phase_fold=None, - relative_to_t_ref=False, add_labels=True, color_by=None, - **kwargs): + if t_ref is None: + t_ref = self.t_ref + return ((self.t - t_ref) / P) % 1.0 + + def plot( + self, + ax=None, + rv_unit=None, + time_format="mjd", + phase_fold=None, + relative_to_t_ref=False, + add_labels=True, + color_by=None, + **kwargs, + ): """ Plot the data points. @@ -414,6 +431,7 @@ def plot(self, ax=None, rv_unit=None, time_format='mjd', phase_fold=None, """ if ax is None: import matplotlib.pyplot as plt + ax = plt.gca() if rv_unit is None: @@ -421,14 +439,14 @@ def plot(self, ax=None, rv_unit=None, time_format='mjd', phase_fold=None, # some default stylings style = kwargs.copy() - style.setdefault('linestyle', 'none') - style.setdefault('alpha', 1.) - style.setdefault('marker', 'o') - style.setdefault('elinewidth', 1) + style.setdefault("linestyle", "none") + style.setdefault("alpha", 1.0) + style.setdefault("marker", "o") + style.setdefault("elinewidth", 1) - if style.get('color', 'k') is not None: - style.setdefault('color', 'k') - style.setdefault('ecolor', '#666666') + if style.get("color", "k") is not None: + style.setdefault("color", "k") + style.setdefault("ecolor", "#666666") if callable(time_format): t = time_format(self.t) @@ -446,36 +464,39 @@ def plot(self, ax=None, rv_unit=None, time_format='mjd', phase_fold=None, if self._has_cov: # FIXME: this is a bit of a hack diag_var = np.diag(self.rv_err.value) - err = np.sqrt(diag_var) * self.rv_err.unit ** 0.5 + err = np.sqrt(diag_var) * self.rv_err.unit**0.5 else: err = self.rv_err - ax.errorbar(t, self.rv.to(rv_unit).value, - err.to(rv_unit).value, **style) + ax.errorbar(t, self.rv.to(rv_unit).value, err.to(rv_unit).value, **style) if add_labels: - ax.set_xlabel('time [BMJD]') - ax.set_ylabel('RV [{:latex_inline}]'.format(rv_unit)) + ax.set_xlabel("time [BMJD]") + ax.set_ylabel(f"RV [{rv_unit:latex_inline}]") return ax def __copy__(self): - return self.__class__(t=self.t.copy(), - rv=self.rv.copy(), - rv_err=self.rv_err.copy()) + return self.__class__( + t=self.t.copy(), rv=self.rv.copy(), rv_err=self.rv_err.copy() + ) def copy(self): return self.__copy__() def __getitem__(self, slc): if self._has_cov: - return self.__class__(t=self.t.copy()[slc], - rv=self.rv.copy()[slc], - rv_err=self.rv_err.copy()[slc][:, slc]) + return self.__class__( + t=self.t.copy()[slc], + rv=self.rv.copy()[slc], + rv_err=self.rv_err.copy()[slc][:, slc], + ) else: - return self.__class__(t=self.t.copy()[slc], - rv=self.rv.copy()[slc], - rv_err=self.rv_err.copy()[slc]) + return self.__class__( + t=self.t.copy()[slc], + rv=self.rv.copy()[slc], + rv_err=self.rv_err.copy()[slc], + ) def __len__(self): return len(self.rv.value) diff --git a/thejoker/exceptions.py b/thejoker/exceptions.py deleted file mode 100644 index eba63390..00000000 --- a/thejoker/exceptions.py +++ /dev/null @@ -1,6 +0,0 @@ -__all__ = ['TheJokerDeprecationWarning'] - - -class TheJokerDeprecationWarning(Warning): - """Because standard deprecation warnings are ignored!""" - pass diff --git a/thejoker/samples.py b/thejoker/samples.py index 575ac288..ee026393 100644 --- a/thejoker/samples.py +++ b/thejoker/samples.py @@ -1,38 +1,41 @@ # Standard library -from collections import OrderedDict import copy import os import warnings +from collections import OrderedDict # Third-party import astropy.units as u -from astropy.time import Time -from astropy.table import Table, QTable, Row, meta, serialize -from astropy.utils.decorators import deprecated_renamed_argument import numpy as np +from astropy.table import QTable, Row, Table, meta, serialize +from astropy.time import Time from twobody import KeplerOrbit, PolynomialRVTrend # Project -from thejoker.src.fast_likelihood import (_nonlinear_packed_order, - _nonlinear_internal_units) -from .prior_helpers import (validate_poly_trend, validate_n_offsets, - get_nonlinear_equiv_units, - get_linear_equiv_units, - get_v0_offsets_equiv_units) -from .samples_helpers import write_table_hdf5 -from .exceptions import TheJokerDeprecationWarning +from thejoker.src.fast_likelihood import ( + _nonlinear_internal_units, + _nonlinear_packed_order, +) + from .likelihood_helpers import ln_normal +from .prior_helpers import ( + get_linear_equiv_units, + get_nonlinear_equiv_units, + get_v0_offsets_equiv_units, + validate_n_offsets, + validate_poly_trend, +) +from .samples_helpers import write_table_hdf5 -__all__ = ['JokerSamples'] +__all__ = ["JokerSamples"] class JokerSamples: - _hdf5_path = 'samples' + _hdf5_path = "samples" - @deprecated_renamed_argument('t0', 't_ref', since='v1.2', - warning_type=TheJokerDeprecationWarning) - def __init__(self, samples=None, t_ref=None, n_offsets=None, - poly_trend=None, **kwargs): + def __init__( + self, samples=None, t_ref=None, n_offsets=None, poly_trend=None, **kwargs + ): """ A dictionary-like object for storing prior or posterior samples from The Joker, with some extra functionality. @@ -61,10 +64,10 @@ def __init__(self, samples=None, t_ref=None, n_offsets=None, if n_offsets is None: n_offsets = 0 - if isinstance(samples, Table) or isinstance(samples, Row): - t_ref = samples.meta.pop('t_ref', t_ref) - poly_trend = samples.meta.pop('poly_trend', poly_trend) - n_offsets = samples.meta.pop('n_offsets', n_offsets) + if isinstance(samples, (Row, Table)): + t_ref = samples.meta.pop("t_ref", t_ref) + poly_trend = samples.meta.pop("poly_trend", poly_trend) + n_offsets = samples.meta.pop("n_offsets", n_offsets) kwargs.update(samples.meta) if isinstance(samples, Table): @@ -82,18 +85,18 @@ def __init__(self, samples=None, t_ref=None, n_offsets=None, valid_units = { **get_nonlinear_equiv_units(), **get_linear_equiv_units(poly_trend), - **get_v0_offsets_equiv_units(n_offsets) + **get_v0_offsets_equiv_units(n_offsets), } # log-prior and log-likelihood values are also valid: - valid_units['ln_prior'] = u.one - valid_units['ln_likelihood'] = u.one - valid_units['ln_posterior'] = u.one + valid_units["ln_prior"] = u.one + valid_units["ln_likelihood"] = u.one + valid_units["ln_posterior"] = u.one self._valid_units = valid_units - self.tbl.meta['t_ref'] = t_ref - self.tbl.meta['poly_trend'] = poly_trend - self.tbl.meta['n_offsets'] = n_offsets + self.tbl.meta["t_ref"] = t_ref + self.tbl.meta["poly_trend"] = poly_trend + self.tbl.meta["n_offsets"] = n_offsets for k, v in kwargs.items(): self.tbl.meta[k] = v @@ -118,32 +121,29 @@ def __getitem__(self, key): def __setitem__(self, key, val): if key not in self._valid_units: - raise ValueError(f"Invalid parameter name '{key}'. Must be one " - "of: {0}".format(list(self._valid_units.keys()))) + raise ValueError( + f"Invalid parameter name '{key}'. Must be one " "of: {0}".format( + list(self._valid_units.keys()) + ) + ) - if not hasattr(val, 'unit'): + if not hasattr(val, "unit"): val = val * u.one # eccentricity expected_unit = self._valid_units[key] if not val.unit.is_equivalent(expected_unit): - raise u.UnitsError(f"Units of '{key}' must be convertable to " - f"{expected_unit}") + raise u.UnitsError( + f"Units of '{key}' must be convertable to " f"{expected_unit}" + ) self.tbl[key] = val @property def t_ref(self): - return self.tbl.meta['t_ref'] - - @property - def t0(self): - warnings.warn('The argument and attribute "t0" has been renamed ' - 'and should now be specified / accessed as "t_ref"', - TheJokerDeprecationWarning) - return self.t_ref + return self.tbl.meta["t_ref"] @u.quantity_input(phase=u.rad) - def get_time_with_phase(self, phase=0*u.rad, t_ref=None): + def get_time_with_phase(self, phase=0 * u.rad, t_ref=None): """ Use the phase at reference time to convert to a time with given phase. @@ -157,34 +157,38 @@ def get_time_with_phase(self, phase=0*u.rad, t_ref=None): """ if t_ref is None: if self.t_ref is None: - raise ValueError("This samples object has no reference time " - "t_ref, so you must pass in the reference " - "time via t_ref") + raise ValueError( + "This samples object has no reference time " + "t_ref, so you must pass in the reference " + "time via t_ref" + ) else: t_ref = self.t_ref elif self.t_ref is not None: - raise ValueError("You passed in a reference time t_ref, but this " - "samples object already has a reference time.") + raise ValueError( + "You passed in a reference time t_ref, but this " + "samples object already has a reference time." + ) - dt = (self['P'] * self['M0'] / (2*np.pi)).to( - u.day, u.dimensionless_angles()) + dt = (self["P"] * self["M0"] / (2 * np.pi)).to(u.day, u.dimensionless_angles()) t0 = t_ref + dt - return t0 + (self['P'] * phase / (2*np.pi)).to( - u.day, u.dimensionless_angles()) + return t0 + (self["P"] * phase / (2 * np.pi)).to( + u.day, u.dimensionless_angles() + ) # TODO: make a property, after deprecation cycle, to replace .t0 def get_t0(self, t_ref=None): - return self.get_time_with_phase(phase=0*u.rad, t_ref=t_ref) + return self.get_time_with_phase(phase=0 * u.rad, t_ref=t_ref) @property def poly_trend(self): - return self.tbl.meta['poly_trend'] + return self.tbl.meta["poly_trend"] @property def n_offsets(self): - return self.tbl.meta['n_offsets'] + return self.tbl.meta["n_offsets"] @property def par_names(self): @@ -197,8 +201,7 @@ def __len__(self): return 1 def __repr__(self): - return (f'') + return f'' def __str__(self): return self.__repr__() @@ -235,28 +238,36 @@ def get_orbit(self, index=None, **kwargs): The samples converted to an orbit object. The barycenter position and distance are set to arbitrary values. """ - if 'orbit' not in self._cache: - self._cache['orbit'] = KeplerOrbit(P=1*u.yr, e=0., omega=0*u.deg, - Omega=0*u.deg, i=90*u.deg, - a=1*u.au, t0=self.t_ref) + if "orbit" not in self._cache: + self._cache["orbit"] = KeplerOrbit( + P=1 * u.yr, + e=0.0, + omega=0 * u.deg, + Omega=0 * u.deg, + i=90 * u.deg, + a=1 * u.au, + t0=self.t_ref, + ) # all of this to avoid the __init__ of KeplerOrbit / KeplerElements - orbit = copy.copy(self._cache['orbit']) + orbit = copy.copy(self._cache["orbit"]) - P = self['P'] - e = self['e'] - K = self['K'] - omega = self['omega'] - M0 = self['M0'] - a = kwargs.pop('a', P * K / (2*np.pi) * np.sqrt(1 - e**2)) + P = self["P"] + e = self["e"] + K = self["K"] + omega = self["omega"] + M0 = self["M0"] + a = kwargs.pop("a", P * K / (2 * np.pi) * np.sqrt(1 - e**2)) names = list(get_linear_equiv_units(self.poly_trend).keys()) trend_coeffs = [self[x] for x in names[1:]] # skip K if index is None or self.isscalar: if len(self) > 1: - raise ValueError("You must specify an index when the number " - f"of samples is >1 (here, it's {len(self)})") + raise ValueError( + "You must specify an index when the number " + f"of samples is >1 (here, it's {len(self)})" + ) else: P = P[index] @@ -272,14 +283,15 @@ def get_orbit(self, index=None, **kwargs): orbit.elements._a = a orbit.elements._omega = omega orbit.elements._M0 = M0 - orbit.elements._Omega = kwargs.pop('Omega', 0*u.deg) - orbit.elements._i = kwargs.pop('i', 90*u.deg) + orbit.elements._Omega = kwargs.pop("Omega", 0 * u.deg) + orbit.elements._i = kwargs.pop("i", 90 * u.deg) orbit._vtrend = PolynomialRVTrend(trend_coeffs, t0=self.t_ref) - orbit._barycenter = kwargs.pop('barycenter', None) + orbit._barycenter = kwargs.pop("barycenter", None) if kwargs: - raise ValueError("Unrecognized arguments {0}" - .format(', '.join(list(kwargs.keys())))) + raise ValueError( + "Unrecognized arguments {0}".format(", ".join(list(kwargs.keys()))) + ) return orbit @@ -313,11 +325,11 @@ def median(self): Return a new scalar object by taking the median in period, and returning the values for that sample """ - med_val = np.percentile(self['P'], 0.5, interpolation='nearest') - if hasattr(med_val, 'unit'): + med_val = np.percentile(self["P"], 0.5, interpolation="nearest") + if hasattr(med_val, "unit"): med_val = med_val.value - median_i, = np.where(self['P'].value == med_val) + (median_i,) = np.where(self["P"].value == med_val) return self[median_i] def std(self): @@ -329,11 +341,11 @@ def wrap_K(self): """ Change negative K values to positive K values and wrap omega to adjust """ - mask = self.tbl['K'] < 0 + mask = self.tbl["K"] < 0 if np.any(mask): - self.tbl['K'][mask] = np.abs(self.tbl['K'][mask]) - self.tbl['omega'][mask] = self.tbl['omega'][mask] + np.pi * u.rad - self.tbl['omega'][mask] = self.tbl['omega'][mask] % (2*np.pi*u.rad) + self.tbl["K"][mask] = np.abs(self.tbl["K"][mask]) + self.tbl["omega"][mask] = self.tbl["omega"][mask] + np.pi * u.rad + self.tbl["omega"][mask] = self.tbl["omega"][mask] % (2 * np.pi * u.rad) return self @@ -434,53 +446,61 @@ def write(self, output, overwrite=False, append=False): try: ext = os.path.splitext(output)[1] except Exception: - raise ValueError("Invalid file name to save samples to: " - f"{output}") - if ext not in ['.hdf5', '.h5', '.fits']: - raise NotImplementedError("We currently only support writing " - "to HDF5 files, with extension .hdf5 " - "or .h5, or FITS files.") + raise ValueError("Invalid file name to save samples to: " f"{output}") + if ext not in [".hdf5", ".h5", ".fits"]: + raise NotImplementedError( + "We currently only support writing " + "to HDF5 files, with extension .hdf5 " + "or .h5, or FITS files." + ) else: - ext = '' + ext = "" - if ext == '.fits': + if ext == ".fits": if append: raise NotImplementedError() t = self.tbl.copy() - if 't0' in t.meta: - warnings.warn('This data file was produced with a deprecated ' - 'version of The Joker and uses old naming ' - 'conventions for the reference time. This file ' - 'may not work with future versions of thejoker.', - TheJokerDeprecationWarning) - t.meta['t_ref'] = t.meta['t0'] - - if t.meta.get('t_ref', None) is not None: - t.meta['__t_ref_bmjd'] = t.meta.pop('t_ref').tcb.mjd + if "t0" in t.meta: + warnings.warn( + "This data file was produced with a deprecated " + "version of The Joker and uses old naming " + "conventions for the reference time. This file " + "may not work with future versions of thejoker.", + DeprecationWarning, + ) + t.meta["t_ref"] = t.meta["t0"] + + if t.meta.get("t_ref", None) is not None: + t.meta["__t_ref_bmjd"] = t.meta.pop("t_ref").tcb.mjd t.write(output, overwrite=overwrite) else: - write_table_hdf5(self.tbl, output, path=self._hdf5_path, - compression=False, - append=append, overwrite=overwrite, - serialize_meta=True, metadata_conflicts='error', - maxshape=(None, )) + write_table_hdf5( + self.tbl, + output, + path=self._hdf5_path, + compression=False, + append=append, + overwrite=overwrite, + serialize_meta=True, + metadata_conflicts="error", + maxshape=(None,), + ) @classmethod def _read_tables(cls, group, path=None): if path is None: path = cls._hdf5_path - samples = group[f'{path}'] - metadata = group[f'{path}.__table_column_meta__'] + samples = group[f"{path}"] + metadata = group[f"{path}.__table_column_meta__"] - header = meta.get_header_from_yaml( - h.decode('utf-8') for h in metadata.read()) + header = meta.get_header_from_yaml(h.decode("utf-8") for h in metadata.read()) table = Table(np.array(samples.read())) - if 'meta' in list(header.keys()): - table.meta = header['meta'] + if "meta" in list(header.keys()): + table.meta = header["meta"] table = serialize._construct_mixins_from_columns(table) @@ -500,6 +520,7 @@ def read(cls, filename, path=None): The output filename or HDF5 group. """ import tables as tb + if isinstance(filename, tb.group.Group): return cls._read_tables(filename) @@ -512,14 +533,15 @@ def read(cls, filename, path=None): except Exception: raise ValueError(f"Invalid file name {filename}") - if ext in ['.hdf5', '.h5']: + if ext in [".hdf5", ".h5"]: tbl = QTable.read(filename, path=path) else: tbl = QTable.read(filename) - if '__t_ref_bmjd' in tbl.meta.keys(): - tbl.meta['t_ref'] = Time(tbl.meta['__t_ref_bmjd'], - format='mjd', scale='tcb') + if "__t_ref_bmjd" in tbl.meta.keys(): + tbl.meta["t_ref"] = Time( + tbl.meta["__t_ref_bmjd"], format="mjd", scale="tcb" + ) else: tbl = QTable.read(filename, path=path) @@ -539,16 +561,16 @@ def ln_unmarginalized_likelihood(self, data): data_unit = data.rv.unit data_var = (data.rv_err.to_value(data_unit)) ** 2 - if 's' in self.tbl.colnames: - s_vars = self['s'].to_value(data_unit) ** 2 + if "s" in self.tbl.colnames: + s_vars = self["s"].to_value(data_unit) ** 2 else: s_vars = np.zeros(len(self)) lls = np.full(len(self), np.nan) for i, (orbit, s) in enumerate(zip(self.orbits, s_vars)): model_rv = orbit.radial_velocity(data.t) - lls[i] = ln_normal(model_rv.to_value(data_unit), - data_rv, - data_var + s).sum() + lls[i] = ln_normal( + model_rv.to_value(data_unit), data_rv, data_var + s + ).sum() return lls diff --git a/thejoker/src/tests/py_likelihood.py b/thejoker/src/tests/py_likelihood.py index 832751f9..6389b38c 100644 --- a/thejoker/src/tests/py_likelihood.py +++ b/thejoker/src/tests/py_likelihood.py @@ -8,7 +8,6 @@ from astroML.utils import log_multivariate_gaussian from twobody.wrap import cy_rv_from_elements -from ...distributions import FixedCompanionMass from ...samples import JokerSamples __all__ = ["get_ivar", "likelihood_worker", "marginal_ln_likelihood"] @@ -145,19 +144,21 @@ def get_M_Lambda_ivar(samples, prior, data): for i, k in enumerate(prior._linear_equiv_units.keys()): if k == "K": continue # set below - Lambda[i] = prior.pars[k].distribution.sd.eval() ** 2 + pars = prior.pars[k].owner.inputs[3:] + Lambda[i] = pars[1].eval() ** 2 - K_dist = prior.pars["K"].distribution - if isinstance(K_dist, FixedCompanionMass): + K_dist = prior.pars["K"] + K_pars = K_dist.owner.inputs[3:] + if K_dist.owner.op._print_name[0] == "FixedCompanionMass": sigma_K0 = K_dist._sigma_K0.to_value(v_unit) P0 = K_dist._P0.to_value(samples["P"].unit) max_K2 = K_dist._max_K.to_value(v_unit) ** 2 else: - Lambda[0] = K_dist.sd.eval() ** 2 + Lambda[0] = K_pars[1].eval() ** 2 for n in range(n_samples): M = design_matrix(packed_samples[n], data, prior) - if isinstance(K_dist, FixedCompanionMass): + if K_dist.owner.op._print_name[0] == "FixedCompanionMass": P = samples["P"][n].value e = samples["e"][n] Lambda[0] = sigma_K0**2 / (1 - e**2) * (P / P0) ** (-2 / 3) diff --git a/thejoker/src/tests/test_fast_likelihood.py b/thejoker/src/tests/test_fast_likelihood.py index 618a9691..b3ad595b 100644 --- a/thejoker/src/tests/test_fast_likelihood.py +++ b/thejoker/src/tests/test_fast_likelihood.py @@ -18,6 +18,8 @@ def test_against_py(): + rng = np.random.default_rng(seed=42) + with pm.Model(): K = xu.with_unit(pm.Normal("K", 0, 10.0), u.km / u.s) @@ -30,9 +32,9 @@ def test_against_py(): ) # t = np.random.uniform(0, 250, 16) + 56831.324 - t = np.sort(np.random.uniform(0, 250, 3)) + 56831.324 + t = np.sort(rng.uniform(0, 250, 3)) + 56831.324 rv = np.cos(t) - rv_err = np.random.uniform(0.1, 0.2, t.size) + rv_err = rng.uniform(0.1, 0.2, t.size) data = RVData(t=t, rv=rv * u.km / u.s, rv_err=rv_err * u.km / u.s) trend_M = get_constant_term_design_matrix(data) @@ -54,6 +56,8 @@ def test_against_py(): def test_scale_varK_against_py(): + rng = np.random.default_rng(seed=42) + prior = JokerPrior.default( P_min=8 * u.day, P_max=32768 * u.day, @@ -63,9 +67,9 @@ def test_scale_varK_against_py(): ) # t = np.random.uniform(0, 250, 16) + 56831.324 - t = np.sort(np.random.uniform(0, 250, 3)) + 56831.324 + t = np.sort(rng.uniform(0, 250, 3)) + 56831.324 rv = np.cos(t) - rv_err = np.random.uniform(0.1, 0.2, t.size) + rv_err = rng.uniform(0.1, 0.2, t.size) data = RVData(t=t, rv=rv * u.km / u.s, rv_err=rv_err * u.km / u.s) trend_M = get_constant_term_design_matrix(data) @@ -87,6 +91,8 @@ def test_scale_varK_against_py(): def test_likelihood_helpers(): + rng = np.random.default_rng(seed=42) + with pm.Model(): K = xu.with_unit(pm.Normal("K", 0, 1.0), u.km / u.s) @@ -98,10 +104,10 @@ def test_likelihood_helpers(): pars={"K": K}, ) - # t = np.random.uniform(0, 250, 16) + 56831.324 - t = np.sort(np.random.uniform(0, 250, 3)) + 56831.324 + # t = rng.uniform(0, 250, 16) + 56831.324 + t = np.sort(rng.uniform(0, 250, 3)) + 56831.324 rv = np.cos(t) - rv_err = np.random.uniform(0.1, 0.2, t.size) + rv_err = rng.uniform(0.1, 0.2, t.size) data = RVData(t=t, rv=rv * u.km / u.s, rv_err=rv_err * u.km / u.s) trend_M = get_constant_term_design_matrix(data) diff --git a/thejoker/tests/test_data.py b/thejoker/tests/test_data.py index c271de78..9f7c0a6e 100644 --- a/thejoker/tests/test_data.py +++ b/thejoker/tests/test_data.py @@ -121,11 +121,6 @@ def test_rvdata_init(): data = RVData(t_arr, rv, err, t_ref=t_obj[3]) assert np.isclose(data.t_ref.mjd, t_obj[3].mjd) - # deprecated: - with warnings.catch_warnings(record=True) as warns: - RVData(t_arr, rv, err, t0=t_obj[3]) - assert len(warns) != 0 - # ------------------------------------------------------------------------ # Test expected failures: @@ -196,17 +191,12 @@ def test_data_methods(tmpdir, inputs): assert u.allclose(data1.rv_err, data2.rv_err) assert u.allclose(data1.t_ref.mjd, data2.t_ref.mjd) - # deprecated: - with warnings.catch_warnings(record=True) as warns: - data1.t0 - assert len(warns) != 0 - # get phase from data object phase1 = data1.phase(P=15.0 * u.day) assert phase1.min() >= 0 assert phase1.max() <= 1 - phase2 = data1.phase(P=15.0 * u.day, t0=Time(58585.24, format="mjd")) + phase2 = data1.phase(P=15.0 * u.day, t_ref=Time(58585.24, format="mjd")) assert not np.allclose(phase1, phase2) # compute inverse variance diff --git a/thejoker/tests/test_prior.py b/thejoker/tests/test_prior.py index dd20b231..8536e834 100644 --- a/thejoker/tests/test_prior.py +++ b/thejoker/tests/test_prior.py @@ -134,7 +134,7 @@ def get_prior(case=None): return prior, units elif case == 8: - # Replace s with pymc3 var with .default() + # Replace s with pymc var with .default() units = default_expected_units.copy() with pm.Model() as model: s = xu.with_unit(pm.Normal("s", 10, 0.5), u.m / u.s) From 0135c12856b0ec493ba625d09ccd4f182e92bc24 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 4 Mar 2024 18:34:19 -0500 Subject: [PATCH 16/50] tests passing again --- pyproject.toml | 3 +- thejoker/plot.py | 223 +++++++++++++++++++-------------- thejoker/prior.py | 70 +++++------ thejoker/samples.py | 75 ++++++----- thejoker/tests/test_plot.py | 84 ++++++++----- thejoker/tests/test_prior.py | 15 ++- thejoker/tests/test_sampler.py | 20 +-- thejoker/tests/test_samples.py | 209 ++++++++++++++++-------------- thejoker/thejoker.py | 11 +- thejoker/utils.py | 2 +- 10 files changed, 395 insertions(+), 317 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index cf9b548e..67b0bd21 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,11 +24,12 @@ dependencies = [ "twobody>=0.9", "scipy", "h5py", - "schwimmbad>=0.3.1", + "schwimmbad>=0.4", "pymc>=5", "pymc_ext @ git+https://github.com/exoplanet-dev/pymc-ext", "exoplanet-core[pymc] @ git+https://github.com/exoplanet-dev/exoplanet-core", "tables", + "dill" ] [project.urls] diff --git a/thejoker/plot.py b/thejoker/plot.py index d3a07b59..c7bbda6c 100644 --- a/thejoker/plot.py +++ b/thejoker/plot.py @@ -1,18 +1,19 @@ # Standard library import warnings +import astropy.units as u +import numpy as np + # Third-party from astropy.time import Time -import astropy.units as u from astropy.utils.decorators import deprecated_renamed_argument -import numpy as np # Project from .data import RVData from .data_helpers import validate_prepare_data from .prior_helpers import get_v0_offsets_equiv_units -__all__ = ['plot_rv_curves', 'plot_phase_fold'] +__all__ = ["plot_rv_curves", "plot_phase_fold"] def get_t_grid(data, P, span_factor=0.1, max_t_grid=None): @@ -32,21 +33,34 @@ def get_t_grid(data, P, span_factor=0.1, max_t_grid=None): "Time grid has more than 10,000 grid points, so plotting orbits " "could be very slow! Set 't_grid' manually, or set 'max_t_grid' to " "decrease the number of time grid points.", - ResourceWarning) + ResourceWarning, + ) - t_grid = np.arange(data.t.mjd.min() - w*span_factor/2, - data.t.mjd.max() + w*span_factor/2 + dt, - dt) + t_grid = np.arange( + data.t.mjd.min() - w * span_factor / 2, + data.t.mjd.max() + w * span_factor / 2 + dt, + dt, + ) return t_grid -@deprecated_renamed_argument('relative_to_t0', 'relative_to_t_ref', - since='v1.2', warning_type=DeprecationWarning) -def plot_rv_curves(samples, t_grid=None, rv_unit=None, data=None, - ax=None, plot_kwargs=dict(), data_plot_kwargs=dict(), - add_labels=True, relative_to_t_ref=False, - apply_mean_v0_offset=True, max_t_grid=None): +@deprecated_renamed_argument( + "relative_to_t0", "relative_to_t_ref", since="v1.2", warning_type=DeprecationWarning +) +def plot_rv_curves( + samples, + t_grid=None, + rv_unit=None, + data=None, + ax=None, + plot_kwargs=dict(), + data_plot_kwargs=dict(), + add_labels=True, + relative_to_t_ref=False, + apply_mean_v0_offset=True, + max_t_grid=None, +): """ Plot radial velocity curves for the input set of orbital parameter samples over the input grid of times. @@ -85,40 +99,43 @@ def plot_rv_curves(samples, t_grid=None, rv_unit=None, data=None, if ax is None: import matplotlib.pyplot as plt + ax = plt.gca() fig = ax.figure if data is not None: - data, ids, _ = validate_prepare_data(data, samples.poly_trend, - samples.n_offsets) + data, ids, _ = validate_prepare_data( + data, samples.poly_trend, samples.n_offsets + ) if t_grid is None: if data is None: - raise ValueError('If data is not passed in, you must specify ' - 'the time grid.') + raise ValueError( + "If data is not passed in, you must specify " "the time grid." + ) - t_grid = get_t_grid(data, samples['P'].min(), max_t_grid=max_t_grid) + t_grid = get_t_grid(data, samples["P"].min(), max_t_grid=max_t_grid) if not isinstance(t_grid, Time): # Assume BMJD - t_grid = Time(t_grid, format='mjd', scale='tcb') + t_grid = Time(t_grid, format="mjd", scale="tcb") # scale the transparency of the lines n_plot = len(samples) - Q = 4. # HACK + Q = 4.0 # HACK line_alpha = 0.05 + Q / (n_plot + Q) if rv_unit is None: - rv_unit = u.km/u.s + rv_unit = u.km / u.s # default plotting style # TODO: move default style to global style config style = plot_kwargs.copy() - style.setdefault('linestyle', '-') - style.setdefault('linewidth', 0.5) - style.setdefault('alpha', line_alpha) - style.setdefault('marker', '') - style.setdefault('color', '#555555') - style.setdefault('rasterized', True) + style.setdefault("linestyle", "-") + style.setdefault("linewidth", 0.5) + style.setdefault("alpha", line_alpha) + style.setdefault("marker", "") + style.setdefault("color", "#555555") + style.setdefault("rasterized", True) # plot orbits over the data model_rv = np.zeros((n_plot, len(t_grid))) @@ -126,17 +143,19 @@ def plot_rv_curves(samples, t_grid=None, rv_unit=None, data=None, orbit = samples.get_orbit(i) if t_grid is None: - t_grid = get_t_grid(data, samples['P'][i], max_t_grid=max_t_grid) + t_grid = get_t_grid(data, samples["P"][i], max_t_grid=max_t_grid) model_rv[i] = orbit.radial_velocity(t_grid).to(rv_unit).value - model_ylim = (np.percentile(model_rv.min(axis=1), 5), - np.percentile(model_rv.max(axis=1), 95)) + model_ylim = ( + np.percentile(model_rv.min(axis=1), 5), + np.percentile(model_rv.max(axis=1), 95), + ) bmjd = t_grid.tcb.mjd if relative_to_t_ref: if samples.t_ref is None: - raise ValueError('Input samples object has no epoch .t_ref') + raise ValueError("Input samples object has no epoch .t_ref") bmjd = bmjd - samples.t_ref.tcb.mjd ax.plot(bmjd, model_rv.T, **style) @@ -149,40 +168,38 @@ def plot_rv_curves(samples, t_grid=None, rv_unit=None, data=None, unq_ids = np.unique(ids) dv0_names = get_v0_offsets_equiv_units(samples.n_offsets).keys() for i, name in enumerate(dv0_names): - mask = ids == unq_ids[i+1] + mask = ids == unq_ids[i + 1] offset_samples = samples[name].to_value(data.rv.unit) data_rv[mask] -= np.mean(offset_samples) - data_err[mask] = np.sqrt(data_err[mask]**2 + - np.var(offset_samples)) - data = RVData(t=data.t, - rv=data_rv * data.rv.unit, - rv_err=data_err * data.rv.unit) + data_err[mask] = np.sqrt(data_err[mask] ** 2 + np.var(offset_samples)) + data = RVData( + t=data.t, rv=data_rv * data.rv.unit, rv_err=data_err * data.rv.unit + ) data_style = data_plot_kwargs.copy() - data_style.setdefault('rv_unit', rv_unit) - data_style.setdefault('markersize', 4.) + data_style.setdefault("rv_unit", rv_unit) + data_style.setdefault("markersize", 4.0) - if data_style['rv_unit'] != rv_unit: + if data_style["rv_unit"] != rv_unit: raise u.UnitsError("Data plot units don't match rv_unit!") - data.plot(ax=ax, relative_to_t_ref=relative_to_t_ref, add_labels=False, - **data_style) + data.plot( + ax=ax, relative_to_t_ref=relative_to_t_ref, add_labels=False, **data_style + ) _rv = data.rv.to(rv_unit).value drv = _rv.max() - _rv.min() - data_ylim = (_rv.min() - 0.2*drv, _rv.max() + 0.2*drv) + data_ylim = (_rv.min() - 0.2 * drv, _rv.max() + 0.2 * drv) else: data_ylim = None ax.set_xlim(bmjd.min(), bmjd.max()) if add_labels: - ax.set_xlabel('BMJD') - ax.set_ylabel('RV [{}]' - .format(rv_unit.to_string(format='latex_inline'))) + ax.set_xlabel("BMJD") + ax.set_ylabel("RV [{}]".format(rv_unit.to_string(format="latex_inline"))) if data_ylim is not None: - ylim = (min(data_ylim[0], model_ylim[0]), - max(data_ylim[1], model_ylim[1])) + ylim = (min(data_ylim[0], model_ylim[0]), max(data_ylim[1], model_ylim[1])) else: ylim = model_ylim @@ -192,18 +209,27 @@ def plot_rv_curves(samples, t_grid=None, rv_unit=None, data=None, return fig -def plot_phase_fold(sample, data=None, ax=None, - with_time_unit=False, n_phase_samples=4096, - add_labels=True, show_s_errorbar=True, residual=False, - remove_trend=True, plot_kwargs=None, data_plot_kwargs=None): +def plot_phase_fold( + sample, + data=None, + ax=None, + with_time_unit=False, + n_phase_samples=4096, + add_labels=True, + show_s_errorbar=True, + residual=False, + remove_trend=True, + plot_kwargs=None, + data_plot_kwargs=None, +): """ Plot phase-folded radial velocity curves for the input orbital parameter sample, optionally with data phase-folded to the same period. Parameters ---------- - samples : :class:`~thejoker.sampler.JokerSamples` - Posterior samples from The Joker. + sample : :class:`~thejoker.sampler.JokerSamples` + One posterior sample from The Joker. data : `~thejoker.data.RVData`, optional Over-plot the data as well. ax : `~matplotlib.Axes`, optional @@ -235,42 +261,44 @@ def plot_phase_fold(sample, data=None, ax=None, if ax is None: import matplotlib.pyplot as plt + ax = plt.gca() fig = ax.figure - # TODO: what do if passing in multiple samples? + if len(sample) > 1: + msg = "You must pass in a single sample" + raise ValueError(msg) if data is not None: - data, ids, _ = validate_prepare_data(data, sample.poly_trend, - sample.n_offsets) + data, ids, _ = validate_prepare_data(data, sample.poly_trend, sample.n_offsets) rv_unit = data.rv.unit else: - rv_unit = sample['v0'].unit + rv_unit = sample["v0"].unit # plotting styles: if plot_kwargs is None: - plot_kwargs = dict() + plot_kwargs = {} if data_plot_kwargs is None: - data_plot_kwargs = dict() + data_plot_kwargs = {} # TODO: move default style to global style config orbit_style = plot_kwargs.copy() - orbit_style.setdefault('linestyle', '-') - orbit_style.setdefault('linewidth', 0.5) - orbit_style.setdefault('alpha', 1.) - orbit_style.setdefault('marker', '') - orbit_style.setdefault('color', '#555555') - orbit_style.setdefault('rasterized', True) + orbit_style.setdefault("linestyle", "-") + orbit_style.setdefault("linewidth", 0.5) + orbit_style.setdefault("alpha", 1.0) + orbit_style.setdefault("marker", "") + orbit_style.setdefault("color", "#555555") + orbit_style.setdefault("rasterized", True) data_style = data_plot_kwargs.copy() - data_style.setdefault('linestyle', 'none') - data_style.setdefault('marker', 'o') - data_style.setdefault('markersize', 4.) - data_style.setdefault('zorder', 10) + data_style.setdefault("linestyle", "none") + data_style.setdefault("marker", "o") + data_style.setdefault("markersize", 4.0) + data_style.setdefault("zorder", 10) # Get orbit from input sample orbit = sample.get_orbit() - P = sample['P'].item() + P = orbit.P t0 = sample.get_t0() if data is not None: @@ -279,12 +307,11 @@ def plot_phase_fold(sample, data=None, ax=None, if remove_trend: # HACK: trend = orbit._vtrend - orbit._vtrend = lambda t: 0. + orbit._vtrend = lambda t: 0.0 rv = rv - trend(data.t) v0_offset_names = get_v0_offsets_equiv_units(sample.n_offsets).keys() - for i, offset_name in zip(range(1, sample.n_offsets+1), - v0_offset_names): + for i, offset_name in zip(range(1, sample.n_offsets + 1), v0_offset_names): _tmp = sample[offset_name].item() rv[ids == i] -= _tmp @@ -301,17 +328,23 @@ def plot_phase_fold(sample, data=None, ax=None, rv = rv - orbit.radial_velocity(data.t) # plot the phase-folded data and orbit - ax.errorbar(phase, rv.to(rv_unit).value, - data.rv_err.to(rv_unit).value, - **data_style) - - if show_s_errorbar and 's' in sample.par_names: - ax.errorbar(phase, rv.to(rv_unit).value, - np.sqrt(data.rv_err**2 + - sample['s']**2).to(rv_unit).value, - linestyle='none', marker='', elinewidth=0., - color='#aaaaaa', alpha=0.9, capsize=0, - zorder=9) + ax.errorbar( + phase, rv.to(rv_unit).value, data.rv_err.to(rv_unit).value, **data_style + ) + + if show_s_errorbar and "s" in sample.par_names: + ax.errorbar( + phase, + rv.to(rv_unit).value, + np.sqrt(data.rv_err**2 + sample["s"] ** 2).to(rv_unit).value, + linestyle="none", + marker="", + elinewidth=0.0, + color="#aaaaaa", + alpha=0.9, + capsize=0, + zorder=9, + ) rv_unit = data.rv.unit @@ -319,7 +352,7 @@ def plot_phase_fold(sample, data=None, ax=None, raise ValueError("TODO: not allowed") else: - rv_unit = sample['K'].unit + rv_unit = sample["K"].unit # Set up the phase grid: unit_phase_grid = np.linspace(0, 1, n_phase_samples) @@ -329,17 +362,19 @@ def plot_phase_fold(sample, data=None, ax=None, phase_grid = unit_phase_grid if not residual: - ax.plot(phase_grid, - orbit.radial_velocity(t0 + - P * unit_phase_grid).to_value(rv_unit), - **orbit_style) + ax.plot( + phase_grid, + orbit.radial_velocity(t0 + P * unit_phase_grid).to_value(rv_unit), + **orbit_style, + ) if add_labels: if with_time_unit is not False: - ax.set_xlabel(r'phase, $(t-t_0)~{\rm mod}~P$ ' + - f'[{time_unit:latex_inline}]') + ax.set_xlabel( + r"phase, $(t-t_0)~{\rm mod}~P$ " + f"[{time_unit:latex_inline}]" + ) else: - ax.set_xlabel(r'phase, $\frac{t-t_0}{P}$') - ax.set_ylabel(f'RV [{rv_unit:latex_inline}]') + ax.set_xlabel(r"phase, $\frac{t-t_0}{P}$") + ax.set_ylabel(f"RV [{rv_unit:latex_inline}]") return fig diff --git a/thejoker/prior.py b/thejoker/prior.py index a10e8807..01253e51 100644 --- a/thejoker/prior.py +++ b/thejoker/prior.py @@ -32,9 +32,8 @@ def _validate_model(model): model = pm.Model() if not isinstance(model, pm.Model): - raise TypeError( - "Input model must be a pymc.Model instance, not " f"a {type(model)}" - ) + msg = f"Input model must be a pymc.Model instance, not a {type(model)}" + raise TypeError(msg) return model @@ -76,7 +75,7 @@ def __init__(self, pars=None, poly_trend=1, v0_offsets=None, model=None): # Parse and clean up the input pars if pars is None: - pars = dict() + pars = {} pars.update(model.named_vars) elif isinstance(pars, pt.TensorVariable): # a single variable @@ -91,13 +90,13 @@ def __init__(self, pars=None, poly_trend=1, v0_offsets=None, model=None): # if that fails, assume it is an iterable, like a list or tuple try: pars = {p.name: p for p in pars} - except Exception: - raise ValueError( - "Invalid input parameters: The input " - "`pars` must either be a dictionary, " - "list, or a single pymc variable, not a " + except Exception as e: + msg = ( + "Invalid input parameters: The input `pars` must either be a " + "dictionary, list, or a single pymc variable, not a " f"'{type(pars)}'." ) + raise ValueError(msg) from e # Set the number of polynomial trend parameters self.poly_trend, self._v_trend_names = validate_poly_trend(poly_trend) @@ -108,12 +107,12 @@ def __init__(self, pars=None, poly_trend=1, v0_offsets=None, model=None): try: v0_offsets = list(v0_offsets) - except Exception: - raise TypeError( - "Constant velocity offsets must be an iterable " - "of pymc variables that define the priors on " - "each offset term." + except Exception as e: + msg = ( + "Constant velocity offsets must be an iterable of pymc variables that " + "define the priors on each offset term." ) + raise TypeError(msg) from e self.v0_offsets = v0_offsets pars.update({p.name: p for p in self.v0_offsets}) @@ -267,9 +266,7 @@ def default( ) pars = {**nl_pars, **l_pars} - obj = cls(pars=pars, model=model, poly_trend=poly_trend, v0_offsets=v0_offsets) - - return obj + return cls(pars=pars, model=model, poly_trend=poly_trend, v0_offsets=v0_offsets) @property def par_names(self): @@ -349,7 +346,7 @@ def sample( par_names = list(self._nonlinear_equiv_units.keys()) # MAJOR HACK RELATED TO UPSTREAM ISSUES WITH pymc3: - # init_shapes = dict() + # init_shapes = {} # for name, par in sub_pars.items(): # if hasattr(par, "distribution"): # init_shapes[name] = par.distribution.shape @@ -366,23 +363,14 @@ def sample( } if return_logprobs: + # raise NotImplementedError("This feature has been disabled in v1.3") logp = [] - for name, par in sub_pars.items(): + for par in sub_pars.values(): try: - _logp = par.distribution.logp(raw_samples[par.name]).eval() - except AttributeError: - logger.warning( - "Cannot auto-compute log-prior value for " - f"parameter {par} because it is defined " - "as a transformation from another " - "variable." - ) - continue + _logp = pm.logp(par, raw_samples[par.name]).eval() except Exception: logger.warning( - "Cannot auto-compute log-prior value for " - f"parameter {par} because it depends on " - "other variables." + f"Cannot auto-compute log-prior value for parameter {par}" ) continue @@ -402,7 +390,7 @@ def sample( p = sub_pars[name] unit = getattr(p, xu.UNIT_ATTR_NAME, u.one) - if name not in prior_samples._valid_units.keys(): + if name not in prior_samples._valid_units: continue prior_samples[name] = np.atleast_1d(raw_samples[name]) * unit @@ -455,7 +443,7 @@ def default_nonlinear_prior(P_min=None, P_max=None, s=None, model=None, pars=Non model = pm.modelcontext(model) if pars is None: - pars = dict() + pars = {} if s is None: s = 0 * u.m / u.s @@ -467,7 +455,7 @@ def default_nonlinear_prior(P_min=None, P_max=None, s=None, model=None, pars=Non raise u.UnitsError("Invalid unit for s: must be equivalent to km/s") # dictionary of parameters to return - out_pars = dict() + out_pars = {} with model: # Set up the default priors for parameters with defaults @@ -492,16 +480,16 @@ def default_nonlinear_prior(P_min=None, P_max=None, s=None, model=None, pars=Non if "P" not in pars: if P_min is None or P_max is None: - raise ValueError( - "If you are using the default period prior, " - "you must pass in both P_min and P_max to set " - "the period prior domain." + msg = ( + "If you are using the default period prior, you must pass in both " + "P_min and P_max to set the period prior domain." ) + raise ValueError(msg) out_pars["P"] = xu.with_unit( UniformLog("P", P_min.value, P_max.to_value(P_min.unit)), P_min.unit ) - for k in pars.keys(): + for k in pars: out_pars[k] = pars[k] return out_pars @@ -537,10 +525,10 @@ def default_linear_prior( model = pm.modelcontext(model) if pars is None: - pars = dict() + pars = {} # dictionary of parameters to return - out_pars = dict() + out_pars = {} # set up poly. trend names: poly_trend, v_names = validate_poly_trend(poly_trend) diff --git a/thejoker/samples.py b/thejoker/samples.py index ee026393..8ff66b0f 100644 --- a/thejoker/samples.py +++ b/thejoker/samples.py @@ -64,19 +64,13 @@ def __init__( if n_offsets is None: n_offsets = 0 - if isinstance(samples, (Row, Table)): - t_ref = samples.meta.pop("t_ref", t_ref) - poly_trend = samples.meta.pop("poly_trend", poly_trend) - n_offsets = samples.meta.pop("n_offsets", n_offsets) - kwargs.update(samples.meta) - - if isinstance(samples, Table): - self.tbl = QTable() - elif isinstance(samples, Row): - self.tbl = samples - - else: - self.tbl = QTable() + self.tbl = QTable() + if isinstance(samples, (Row, Table, QTable)): + meta = samples.meta.copy() + t_ref = meta.pop("t_ref", t_ref) + poly_trend = meta.pop("poly_trend", poly_trend) + n_offsets = meta.pop("n_offsets", n_offsets) + kwargs.update(meta) # Validate input poly_trend / n_offsets: poly_trend, _ = validate_poly_trend(poly_trend) @@ -104,20 +98,19 @@ def __init__( if samples is not None: _tbl = QTable(samples) for colname in _tbl.colnames: - self[colname] = _tbl[colname] + self[colname] = np.atleast_1d(_tbl[colname]) # used for speed-ups below - self._cache = dict() + self._cache = {} def __getitem__(self, key): if isinstance(key, int): return self.__class__(samples=self.tbl[key]) - elif isinstance(key, str) and key in self.par_names: + if isinstance(key, str) and key in self.par_names: return self.tbl[key] - else: - return self.__class__(samples=self.tbl[key]) + return self.__class__(samples=self.tbl[key]) def __setitem__(self, key, val): if key not in self._valid_units: @@ -174,8 +167,8 @@ def get_time_with_phase(self, phase=0 * u.rad, t_ref=None): dt = (self["P"] * self["M0"] / (2 * np.pi)).to(u.day, u.dimensionless_angles()) t0 = t_ref + dt - return t0 + (self["P"] * phase / (2 * np.pi)).to( - u.day, u.dimensionless_angles() + return np.squeeze( + t0 + (self["P"] * phase / (2 * np.pi)).to(u.day, u.dimensionless_angles()) ) # TODO: make a property, after deprecation cycle, to replace .t0 @@ -262,20 +255,21 @@ def get_orbit(self, index=None, **kwargs): names = list(get_linear_equiv_units(self.poly_trend).keys()) trend_coeffs = [self[x] for x in names[1:]] # skip K - if index is None or self.isscalar: + if index is None: if len(self) > 1: - raise ValueError( - "You must specify an index when the number " - f"of samples is >1 (here, it's {len(self)})" + msg = ( + "You must specify an index when the number of samples is >1 (here, " + f"it's {len(self)})" ) + raise ValueError(msg) + index = 0 - else: - P = P[index] - e = e[index] - a = a[index] - omega = omega[index] - M0 = M0[index] - trend_coeffs = [x[index] for x in trend_coeffs] + P = P[index] + e = e[index] + a = a[index] + omega = omega[index] + M0 = M0[index] + trend_coeffs = [x[index] for x in trend_coeffs] orbit.elements.t0 = self.t_ref orbit.elements._P = P @@ -310,7 +304,7 @@ def orbits(self): def _apply(self, func): cls = self.__class__ - new_samples = dict() + new_samples = {} for k in self.tbl.colnames: new_samples[k] = np.atleast_1d(func(self[k])) @@ -320,17 +314,13 @@ def mean(self): """Return a new scalar object by taking the mean across all samples""" return self._apply(np.mean) - def median(self): + def median_period(self): """ Return a new scalar object by taking the median in period, and returning the values for that sample """ - med_val = np.percentile(self["P"], 0.5, interpolation="nearest") - if hasattr(med_val, "unit"): - med_val = med_val.value - - (median_i,) = np.where(self["P"].value == med_val) - return self[median_i] + idx = np.argpartition(self["P"], len(self["P"]) // 2)[len(self["P"]) // 2] + return self[idx] def std(self): """Return a new scalar object by taking the standard deviation across @@ -457,6 +447,8 @@ def write(self, output, overwrite=False, append=False): ext = "" if ext == ".fits": + from astropy.io import fits + if append: raise NotImplementedError() @@ -474,7 +466,10 @@ def write(self, output, overwrite=False, append=False): if t.meta.get("t_ref", None) is not None: t.meta["__t_ref_bmjd"] = t.meta.pop("t_ref").tcb.mjd - t.write(output, overwrite=overwrite) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=fits.verify.VerifyWarning) + t.write(output, overwrite=overwrite) else: write_table_hdf5( self.tbl, diff --git a/thejoker/tests/test_plot.py b/thejoker/tests/test_plot.py index 1a9cd965..2730968e 100644 --- a/thejoker/tests/test_plot.py +++ b/thejoker/tests/test_plot.py @@ -1,36 +1,47 @@ # Third-party -from astropy.time import Time import astropy.units as u import numpy as np import pytest +from astropy.time import Time try: import matplotlib.pyplot as plt + HAS_MPL = True except ImportError: HAS_MPL = False # Package +from ..plot import plot_phase_fold, plot_rv_curves from ..prior import JokerPrior from ..samples import JokerSamples -from ..plot import plot_rv_curves, plot_phase_fold from .test_sampler import make_data -@pytest.mark.skipif(not HAS_MPL, reason='matplotlib not installed') -@pytest.mark.parametrize('prior', [ - JokerPrior.default(10*u.day, 20*u.day, - 25*u.km/u.s, sigma_v=100*u.km/u.s), - JokerPrior.default(10*u.day, 20*u.day, - 25*u.km/u.s, poly_trend=2, - sigma_v=[100*u.km/u.s, 0.2*u.km/u.s/u.day]) -]) +@pytest.mark.skipif(not HAS_MPL, reason="matplotlib not installed") +@pytest.mark.parametrize( + "prior", + [ + JokerPrior.default( + 10 * u.day, 20 * u.day, 25 * u.km / u.s, sigma_v=100 * u.km / u.s + ), + JokerPrior.default( + 10 * u.day, + 20 * u.day, + 25 * u.km / u.s, + poly_trend=2, + sigma_v=[100 * u.km / u.s, 0.2 * u.km / u.s / u.day], + ), + ], +) def test_plot_rv_curves(prior): - data, _ = make_data() - samples = prior.sample(100, generate_linear=True, t_ref=Time('J2000')) + samples = prior.sample(100, generate_linear=True, t_ref=Time("J2000")) + + samples[0] - t_grid = np.random.uniform(56000, 56500, 1024) + rng = np.random.default_rng(42) + t_grid = rng.uniform(56000, 56500, 1024) t_grid.sort() plot_rv_curves(samples, t_grid) @@ -41,35 +52,42 @@ def test_plot_rv_curves(prior): plot_rv_curves(samples, t_grid, ax=ax) -@pytest.mark.skipif(not HAS_MPL, reason='matplotlib not installed') -@pytest.mark.parametrize('prior', [ - JokerPrior.default(10*u.day, 20*u.day, - 25*u.km/u.s, sigma_v=100*u.km/u.s), - JokerPrior.default(10*u.day, 20*u.day, - 25*u.km/u.s, poly_trend=2, - sigma_v=[100*u.km/u.s, 0.2*u.km/u.s/u.day]) -]) +@pytest.mark.skipif(not HAS_MPL, reason="matplotlib not installed") +@pytest.mark.parametrize( + "prior", + [ + JokerPrior.default( + 10 * u.day, 20 * u.day, 25 * u.km / u.s, sigma_v=100 * u.km / u.s + ), + JokerPrior.default( + 10 * u.day, + 20 * u.day, + 25 * u.km / u.s, + poly_trend=2, + sigma_v=[100 * u.km / u.s, 0.2 * u.km / u.s / u.day], + ), + ], +) def test_plot_phase_fold(prior): - data, _ = make_data() - samples = prior.sample(100, generate_linear=True, t_ref=Time('J2000')) + samples = prior.sample(100, generate_linear=True, t_ref=Time("J2000")) - plot_phase_fold(samples.median(), data) - plot_phase_fold(samples[0:1], data) - plot_phase_fold(samples[0:1], data=None) + plot_phase_fold(samples.median_period(), data) + plot_phase_fold(samples[0], data) + plot_phase_fold(samples[0], data=None) def test_big_grid_warning(): data, _ = make_data() samples = JokerSamples(t_ref=data.t_ref) - samples['P'] = [0.001] * u.day - samples['e'] = [0.3] - samples['omega'] = [0.5] * u.rad - samples['M0'] = [0.34] * u.rad - samples['s'] = [0] * u.km/u.s - samples['v0'] = [10.] * u.km/u.s - samples['K'] = [5.] * u.km/u.s + samples["P"] = [0.001] * u.day + samples["e"] = [0.3] + samples["omega"] = [0.5] * u.rad + samples["M0"] = [0.34] * u.rad + samples["s"] = [0] * u.km / u.s + samples["v0"] = [10.0] * u.km / u.s + samples["K"] = [5.0] * u.km / u.s with pytest.warns(ResourceWarning, match="10,000"): plot_rv_curves(samples, data=data, max_t_grid=20000) diff --git a/thejoker/tests/test_prior.py b/thejoker/tests/test_prior.py index 8536e834..8dc62fdf 100644 --- a/thejoker/tests/test_prior.py +++ b/thejoker/tests/test_prior.py @@ -1,5 +1,7 @@ -# Third-party +import io + import astropy.units as u +import dill import numpy as np import pymc as pm import pytest @@ -189,6 +191,17 @@ def test_init_sample(case): assert samples[k].unit == expected_units[k] +@pytest.mark.parametrize("case", range(get_prior())) +def test_pickle(case): + prior, _ = get_prior(case) + + with io.BytesIO() as f: + dill.dump(prior, f) + + f.seek(0) + dill.load(f) + + def test_dtype(): # Things that don't need to be checked for all cases above prior, expected_units = get_prior(0) diff --git a/thejoker/tests/test_sampler.py b/thejoker/tests/test_sampler.py index 48446742..1d6cd8b8 100644 --- a/thejoker/tests/test_sampler.py +++ b/thejoker/tests/test_sampler.py @@ -1,7 +1,5 @@ -# Standard library import os -# Third-party import astropy.units as u import numpy as np import pytest @@ -9,12 +7,10 @@ from schwimmbad import MultiPool, SerialPool from twobody import KeplerOrbit -from ..data import RVData - -# Package -from ..prior import JokerPrior -from ..thejoker import TheJoker -from .test_prior import get_prior +from thejoker.data import RVData +from thejoker.prior import JokerPrior +from thejoker.tests.test_prior import get_prior +from thejoker.thejoker import TheJoker def make_data(n_times=8, rng=None, v1=None, K=None): @@ -69,7 +65,7 @@ def test_init(case): TheJoker(prior, pool="sdfks") # Random state: - rnd = np.random.default_rng(42) + rng = np.random.default_rng(42) TheJoker(prior, rng=rng) # fail when random state is invalid: @@ -229,3 +225,9 @@ def test_iterative_rejection_sample(tmpdir, prior): for i in range(1, len(all_Ps)): assert u.allclose(all_Ps[0], all_Ps[i]) assert u.allclose(all_Ks[0], all_Ks[i]) + + +if __name__ == "__main__": + import pathlib + + test_marginal_ln_likelihood(pathlib.Path("/tmp/"), 0) diff --git a/thejoker/tests/test_samples.py b/thejoker/tests/test_samples.py index c303966c..1a2fe543 100644 --- a/thejoker/tests/test_samples.py +++ b/thejoker/tests/test_samples.py @@ -1,10 +1,10 @@ # Third-party -from astropy.time import Time import astropy.units as u -from astropy.tests.helper import quantity_allclose import h5py import numpy as np import pytest +from astropy.tests.helper import quantity_allclose +from astropy.time import Time # Project from ..samples import JokerSamples @@ -15,8 +15,8 @@ def test_joker_samples(tmpdir): N = 100 # Generate some fake samples - samples = JokerSamples(dict(P=np.random.random(size=N)*u.day)) - assert 'P' in samples.par_names + samples = JokerSamples(dict(P=np.random.random(size=N) * u.day)) + assert "P" in samples.par_names assert len(samples) == N # Test slicing with a number @@ -33,134 +33,135 @@ def test_joker_samples(tmpdir): # Invalid name with pytest.raises(ValueError): - JokerSamples(dict(derp=[15.]*u.kpc)) + JokerSamples(dict(derp=[15.0] * u.kpc)) # Length conflicts with previous length - samples = JokerSamples({'P': np.random.random(size=N)*u.day}) + samples = JokerSamples({"P": np.random.random(size=N) * u.day}) with pytest.raises(ValueError): - samples['v0'] = np.random.random(size=10)*u.km/u.s + samples["v0"] = np.random.random(size=10) * u.km / u.s # Write to HDF5 file samples = JokerSamples() - samples['P'] = np.random.uniform(800, 1000, size=N)*u.day - samples['M0'] = 2*np.pi*np.random.random(size=N)*u.radian - samples['e'] = np.random.random(size=N) - samples['omega'] = 2*np.pi*np.random.random(size=N)*u.radian + samples["P"] = np.random.uniform(800, 1000, size=N) * u.day + samples["M0"] = 2 * np.pi * np.random.random(size=N) * u.radian + samples["e"] = np.random.random(size=N) + samples["omega"] = 2 * np.pi * np.random.random(size=N) * u.radian - fn = str(tmpdir / 'test.hdf5') + fn = str(tmpdir / "test.hdf5") samples.write(fn) samples2 = JokerSamples.read(fn) for k in samples.par_names: assert quantity_allclose(samples[k], samples2[k]) - new_samples = samples[samples['P'].argmin()] - assert quantity_allclose(new_samples['P'], - samples['P'][samples['P'].argmin()]) + new_samples = samples[samples["P"].argmin()] + assert quantity_allclose(new_samples["P"], samples["P"][samples["P"].argmin()]) # Check that t_ref writes / gets loaded - samples = JokerSamples(t_ref=Time('J2000')) - samples['P'] = np.random.uniform(800, 1000, size=N)*u.day - samples['M0'] = 2*np.pi*np.random.random(size=N)*u.radian - samples['e'] = np.random.random(size=N) - samples['omega'] = 2*np.pi*np.random.random(size=N)*u.radian + samples = JokerSamples(t_ref=Time("J2000")) + samples["P"] = np.random.uniform(800, 1000, size=N) * u.day + samples["M0"] = 2 * np.pi * np.random.random(size=N) * u.radian + samples["e"] = np.random.random(size=N) + samples["omega"] = 2 * np.pi * np.random.random(size=N) * u.radian - fn = str(tmpdir / 'test2.hdf5') + fn = str(tmpdir / "test2.hdf5") samples.write(fn) samples.write(fn, overwrite=True) samples2 = JokerSamples.read(fn) assert samples2.t_ref is not None - assert np.isclose(samples2.t_ref.mjd, Time('J2000').mjd) + assert np.isclose(samples2.t_ref.mjd, Time("J2000").mjd) # Check that scalar samples are supported: - samples = JokerSamples(t_ref=Time('J2000')) - samples['P'] = np.random.uniform(800, 1000, size=N)*u.day - samples['M0'] = 2*np.pi*np.random.random(size=N)*u.radian - samples['e'] = np.random.random(size=N) - samples['omega'] = 2*np.pi*np.random.random(size=N)*u.radian + samples = JokerSamples(t_ref=Time("J2000")) + samples["P"] = np.random.uniform(800, 1000, size=N) * u.day + samples["M0"] = 2 * np.pi * np.random.random(size=N) * u.radian + samples["e"] = np.random.random(size=N) + samples["omega"] = 2 * np.pi * np.random.random(size=N) * u.radian new_samples = samples[0] assert len(new_samples) == 1 # Check that polynomial trends work - samples = JokerSamples(t_ref=Time('J2015.5'), - poly_trend=3) - samples['P'] = np.random.uniform(800, 1000, size=N)*u.day - samples['M0'] = 2*np.pi*np.random.random(size=N)*u.radian - samples['e'] = np.random.random(size=N) - samples['omega'] = 2*np.pi*np.random.random(size=N)*u.radian - samples['K'] = 100 * np.random.normal(size=N) * u.km/u.s - samples['v0'] = np.random.uniform(0, 10, size=N) * u.km/u.s - samples['v1'] = np.random.uniform(0, 1, size=N) * u.km/u.s/u.day - samples['v2'] = np.random.uniform(0, 1e-2, size=N) * u.km/u.s/u.day**2 + samples = JokerSamples(t_ref=Time("J2015.5"), poly_trend=3) + samples["P"] = np.random.uniform(800, 1000, size=N) * u.day + samples["M0"] = 2 * np.pi * np.random.random(size=N) * u.radian + samples["e"] = np.random.random(size=N) + samples["omega"] = 2 * np.pi * np.random.random(size=N) * u.radian + samples["K"] = 100 * np.random.normal(size=N) * u.km / u.s + samples["v0"] = np.random.uniform(0, 10, size=N) * u.km / u.s + samples["v1"] = np.random.uniform(0, 1, size=N) * u.km / u.s / u.day + samples["v2"] = np.random.uniform(0, 1e-2, size=N) * u.km / u.s / u.day**2 new_samples = samples[0] orb = samples.get_orbit(0) - orb.radial_velocity(Time('J2015.6')) + orb.radial_velocity(Time("J2015.6")) + + with pytest.raises(ValueError): + samples.get_orbit() orb = new_samples.get_orbit() - orb.radial_velocity(Time('J2015.6')) + orb.radial_velocity(Time("J2015.6")) def test_append_write(tmpdir): N = 64 - samples = JokerSamples(t_ref=Time('J2000')) - samples['P'] = np.random.uniform(800, 1000, size=N)*u.day - samples['M0'] = 2*np.pi*np.random.random(size=N)*u.radian - samples['e'] = np.random.random(size=N) - samples['omega'] = 2*np.pi*np.random.random(size=N)*u.radian + samples = JokerSamples(t_ref=Time("J2000")) + samples["P"] = np.random.uniform(800, 1000, size=N) * u.day + samples["M0"] = 2 * np.pi * np.random.random(size=N) * u.radian + samples["e"] = np.random.random(size=N) + samples["omega"] = 2 * np.pi * np.random.random(size=N) * u.radian - fn = str(tmpdir / 'test-tbl.hdf5') + fn = str(tmpdir / "test-tbl.hdf5") samples.write(fn) samples.write(fn, append=True) - fn2 = str(tmpdir / 'test-tbl2.hdf5') - with h5py.File(fn2, 'w') as f: - g = f.create_group('2M00') + fn2 = str(tmpdir / "test-tbl2.hdf5") + with h5py.File(fn2, "w") as f: + g = f.create_group("2M00") samples.write(g) samples2 = JokerSamples.read(fn) - assert np.all(samples2['P'][:N] == samples2['P'][N:]) + assert np.all(samples2["P"][:N] == samples2["P"][N:]) def test_apply_methods(): N = 100 # Test that samples objects reduce properly - samples = JokerSamples(t_ref=Time('J2000')) - samples['P'] = np.random.uniform(800, 1000, size=N)*u.day - samples['M0'] = 2*np.pi*np.random.random(size=N)*u.radian - samples['e'] = np.random.random(size=N) - samples['omega'] = 2*np.pi*np.random.random(size=N)*u.radian + samples = JokerSamples(t_ref=Time("J2000")) + samples["P"] = np.random.uniform(800, 1000, size=N) * u.day + samples["M0"] = 2 * np.pi * np.random.random(size=N) * u.radian + samples["e"] = np.random.random(size=N) + samples["omega"] = 2 * np.pi * np.random.random(size=N) * u.radian new_samples = samples.mean() - assert quantity_allclose(new_samples['P'], np.mean(samples['P'])) + assert quantity_allclose(new_samples["P"], np.mean(samples["P"])) # try just executing others: - new_samples = samples.median() + new_samples = samples.median_period() new_samples = samples.std() -@pytest.mark.parametrize("t_ref", [None, Time('J2015.5')]) +@pytest.mark.parametrize("t_ref", [None, Time("J2015.5")]) @pytest.mark.parametrize("poly_trend", [1, 3]) -@pytest.mark.parametrize("ext", ['.hdf5', '.fits']) +@pytest.mark.parametrize("ext", [".hdf5", ".fits"]) def test_table(tmp_path, t_ref, poly_trend, ext): N = 16 samples = JokerSamples(t_ref=t_ref, poly_trend=poly_trend) - samples['P'] = np.random.uniform(800, 1000, size=N)*u.day - samples['M0'] = 2*np.pi*np.random.random(size=N)*u.radian - samples['e'] = np.random.random(size=N) - samples['omega'] = 2*np.pi*np.random.random(size=N)*u.radian - samples['K'] = 100 * np.random.normal(size=N) * u.km/u.s - samples['v0'] = np.random.uniform(0, 10, size=N) * u.km/u.s + samples["P"] = np.random.uniform(800, 1000, size=N) * u.day + samples["M0"] = 2 * np.pi * np.random.random(size=N) * u.radian + samples["e"] = np.random.random(size=N) + samples["omega"] = 2 * np.pi * np.random.random(size=N) * u.radian + samples["K"] = 100 * np.random.normal(size=N) * u.km / u.s + samples["v0"] = np.random.uniform(0, 10, size=N) * u.km / u.s if poly_trend > 1: - samples['v1'] = np.random.uniform(0, 1, size=N) * u.km/u.s/u.day - samples['v2'] = np.random.uniform(0, 1e-2, size=N) * u.km/u.s/u.day**2 + samples["v1"] = np.random.uniform(0, 1, size=N) * u.km / u.s / u.day + samples["v2"] = np.random.uniform(0, 1e-2, size=N) * u.km / u.s / u.day**2 d = tmp_path / "table" d.mkdir() - path = str(d / f"t_{str(t_ref)}_{poly_trend}{ext}") + path = str(d / f"t_{t_ref!s}_{poly_trend}{ext}") samples.write(path) samples2 = JokerSamples.read(path) @@ -178,52 +179,72 @@ def _get_dtype_compare_cases(): evals = [] # True - d1 = [{'name': 'P', 'unit': 'd', 'datatype': 'float64'}, - {'name': 'e', 'datatype': 'float64'}, - {'name': 'omega', 'unit': 'rad', 'datatype': 'float64'}, - {'name': 'M0', 'unit': 'rad', 'datatype': 'float64'}, - {'name': 's', 'unit': 'm / s', 'datatype': 'float64'}] - d2 = [{'name': 'P', 'unit': 'd', 'datatype': 'float64'}, - {'name': 'e', 'unit': '', 'datatype': 'float64'}, - {'name': 'omega', 'unit': 'rad', 'datatype': 'float64'}, - {'name': 'M0', 'unit': 'rad', 'datatype': 'float64'}, - {'name': 's', 'unit': 'm / s', 'datatype': 'float64'}] + d1 = [ + {"name": "P", "unit": "d", "datatype": "float64"}, + {"name": "e", "datatype": "float64"}, + {"name": "omega", "unit": "rad", "datatype": "float64"}, + {"name": "M0", "unit": "rad", "datatype": "float64"}, + {"name": "s", "unit": "m / s", "datatype": "float64"}, + ] + d2 = [ + {"name": "P", "unit": "d", "datatype": "float64"}, + {"name": "e", "unit": "", "datatype": "float64"}, + {"name": "omega", "unit": "rad", "datatype": "float64"}, + {"name": "M0", "unit": "rad", "datatype": "float64"}, + {"name": "s", "unit": "m / s", "datatype": "float64"}, + ] d1s.append(d1) d2s.append(d2) evals.append(True) # True - d1 = [{'name': 'P', 'unit': 'd', 'datatype': 'float64'}, - {'name': 'e', 'datatype': 'float64'}] - d2 = [{'name': 'P', 'unit': 'd', 'datatype': 'float64'}, - {'name': 'e', 'unit': '', 'datatype': 'float64'}] + d1 = [ + {"name": "P", "unit": "d", "datatype": "float64"}, + {"name": "e", "datatype": "float64"}, + ] + d2 = [ + {"name": "P", "unit": "d", "datatype": "float64"}, + {"name": "e", "unit": "", "datatype": "float64"}, + ] d1s.append(d1) d2s.append(d2) evals.append(True) # True - d1 = [{'name': 'P', 'unit': 'd', 'datatype': 'float64'}, - {'name': 'e', 'unit': '', 'datatype': 'float64'}] - d2 = [{'name': 'P', 'unit': 'd', 'datatype': 'float64'}, - {'name': 'e', 'unit': '', 'datatype': 'float64'}] + d1 = [ + {"name": "P", "unit": "d", "datatype": "float64"}, + {"name": "e", "unit": "", "datatype": "float64"}, + ] + d2 = [ + {"name": "P", "unit": "d", "datatype": "float64"}, + {"name": "e", "unit": "", "datatype": "float64"}, + ] d1s.append(d1) d2s.append(d2) evals.append(True) # False - d1 = [{'name': 'P', 'unit': 'd', 'datatype': 'float64'}, - {'name': 'e', 'unit': 'a', 'datatype': 'float64'}] - d2 = [{'name': 'P', 'unit': 'd', 'datatype': 'float64'}, - {'name': 'e', 'unit': '', 'datatype': 'float64'}] + d1 = [ + {"name": "P", "unit": "d", "datatype": "float64"}, + {"name": "e", "unit": "a", "datatype": "float64"}, + ] + d2 = [ + {"name": "P", "unit": "d", "datatype": "float64"}, + {"name": "e", "unit": "", "datatype": "float64"}, + ] d1s.append(d1) d2s.append(d2) evals.append(False) # False - d1 = [{'name': 'P', 'unit': 'd', 'datatype': 'float64'}, - {'name': 'e', 'unit': 'a', 'datatype': 'float64'}] - d2 = [{'name': 'P', 'unit': 'd', 'datatype': 'float64'}, - {'name': 'e', 'unit': 'b', 'datatype': 'float64'}] + d1 = [ + {"name": "P", "unit": "d", "datatype": "float64"}, + {"name": "e", "unit": "a", "datatype": "float64"}, + ] + d2 = [ + {"name": "P", "unit": "d", "datatype": "float64"}, + {"name": "e", "unit": "b", "datatype": "float64"}, + ] d1s.append(d1) d2s.append(d2) evals.append(False) @@ -231,6 +252,6 @@ def _get_dtype_compare_cases(): return zip(d1s, d2s, evals) -@pytest.mark.parametrize('d1,d2,equal', _get_dtype_compare_cases()) +@pytest.mark.parametrize("d1,d2,equal", _get_dtype_compare_cases()) def test_table_dtype_compare(d1, d2, equal): assert _custom_tbl_dtype_compare(d1, d2) == equal diff --git a/thejoker/thejoker.py b/thejoker/thejoker.py index 9dab5f82..ea0aa595 100644 --- a/thejoker/thejoker.py +++ b/thejoker/thejoker.py @@ -61,6 +61,12 @@ def __init__(self, prior, pool=None, rng=None, tempfile_path=None): # based on the parent if rng is None: rng = np.random.default_rng() + elif not isinstance(rng, np.random.Generator): + msg = ( + "The input random number generator must be a numpy.random.Generator " + "instance." + ) + raise TypeError(msg) self.rng = rng # check if a JokerParams instance was passed in to specify the state @@ -82,8 +88,7 @@ def _make_joker_helper(self, data): all_data, ids, trend_M = validate_prepare_data( data, self.prior.poly_trend, self.prior.n_offsets ) - joker_helper = CJokerHelper(all_data, self.prior, trend_M) - return joker_helper + return CJokerHelper(all_data, self.prior, trend_M) def marginal_ln_likelihood( self, data, prior_samples, n_batches=None, in_memory=False @@ -417,7 +422,7 @@ def setup_mcmc(self, data, joker_samples, model=None, custom_func=None): if not is_P_unimodal(joker_samples, data): logger.warn("TODO: samples ain't unimodal") - MAP_sample = joker_samples.median() + MAP_sample = joker_samples.median_period() else: MAP_sample = joker_samples diff --git a/thejoker/utils.py b/thejoker/utils.py index efd587f9..4fd6b226 100644 --- a/thejoker/utils.py +++ b/thejoker/utils.py @@ -240,7 +240,7 @@ def read_random_batch(prior_samples_file, columns, size, units=None, rng=None): path = JokerSamples._hdf5_path with tb.open_file(prior_samples_file, mode="r") as f: - idx = rng.integers(rng, 0, f.root[path].shape[0], size=size) + idx = rng.choice(f.root[path].shape[0], size=size, replace=False) return read_batch_idx(prior_samples_file, columns, idx=idx, units=units) From 67c2f82e42e5650b64bd21596d20fe2a299deb18 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 4 Mar 2024 18:47:38 -0500 Subject: [PATCH 17/50] remove some old packaging stuff --- .gitattributes | 2 -- .readthedocs.yml | 13 +++++-------- MANIFEST.in | 20 -------------------- environment.yml | 6 ------ setup.py | 4 +--- thejoker/conftest.py | 28 ---------------------------- 6 files changed, 6 insertions(+), 67 deletions(-) delete mode 100644 .gitattributes delete mode 100644 MANIFEST.in delete mode 100644 environment.yml delete mode 100644 thejoker/conftest.py diff --git a/.gitattributes b/.gitattributes deleted file mode 100644 index 305a9b04..00000000 --- a/.gitattributes +++ /dev/null @@ -1,2 +0,0 @@ -*.ipynb filter=nbstripout -*.ipynb diff=ipynb diff --git a/.readthedocs.yml b/.readthedocs.yml index 46905bc1..d65e9c1b 100644 --- a/.readthedocs.yml +++ b/.readthedocs.yml @@ -1,21 +1,18 @@ # Read the Docs configuration file # See https://docs.readthedocs.io/en/stable/config-file/v2.html for details -# Required version: 2 -# Build documentation in the docs/ directory with Sphinx +build: + os: ubuntu-22.04 + tools: + python: "3.11" sphinx: configuration: docs/conf.py -# Have to use conda because of weird lazylinker error -conda: - environment: environment.yml - python: - version: 3.8 install: - method: pip path: . extra_requirements: - - docs + - docs \ No newline at end of file diff --git a/MANIFEST.in b/MANIFEST.in deleted file mode 100644 index e5b87c22..00000000 --- a/MANIFEST.in +++ /dev/null @@ -1,20 +0,0 @@ -include README.rst -include CHANGES.rst - -include setup.cfg -include pyproject.toml -include LICENSE - -recursive-include thejoker *.pyx *.c *.pxd -recursive-include docs * -recursive-include licenses * -recursive-include cextern * -recursive-include scripts * - -prune build -prune docs/_build -prune docs/api -prune docs/examples - - -global-exclude *.pyc *.o diff --git a/environment.yml b/environment.yml deleted file mode 100644 index 02ec2a0d..00000000 --- a/environment.yml +++ /dev/null @@ -1,6 +0,0 @@ -dependencies: - - python=3.8 - - mkl - - pip - - pip: - - -e .[docs] diff --git a/setup.py b/setup.py index 1190b3ae..30594e80 100644 --- a/setup.py +++ b/setup.py @@ -24,6 +24,4 @@ cfg["sources"].append("thejoker/src/fast_likelihood.pyx") exts.append(Extension("thejoker.src.fast_likelihood", **cfg)) -setup( - ext_modules=exts, -) +setup(ext_modules=exts) diff --git a/thejoker/conftest.py b/thejoker/conftest.py deleted file mode 100644 index d3596c01..00000000 --- a/thejoker/conftest.py +++ /dev/null @@ -1,28 +0,0 @@ -import os - -try: - from pytest_astropy_header.display import PYTEST_HEADER_MODULES, TESTED_VERSIONS - ASTROPY_HEADER = True -except ImportError: - ASTROPY_HEADER = False - - -def pytest_configure(config): - """Configure Pytest with Astropy. - Parameters - ---------- - config : pytest configuration - """ - if ASTROPY_HEADER: - - config.option.astropy_header = True - - # Customize the following lines to add/remove entries from the list of - # packages for which version numbers are displayed when running the tests. - PYTEST_HEADER_MODULES.pop('Pandas', None) - PYTEST_HEADER_MODULES['scikit-image'] = 'skimage' - PYTEST_HEADER_MODULES['twobody'] = 'twobody' - - from . import __version__ - packagename = os.path.basename(os.path.dirname(__file__)) - TESTED_VERSIONS[packagename] = __version__ From b43f2038ff18df7119e9e496b7119ea8fb8e6511 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 4 Mar 2024 18:55:25 -0500 Subject: [PATCH 18/50] redo ci and cd --- .github/workflows/cd.yml | 73 ++++++++++++++++++++++++++ .github/workflows/ci.yml | 42 +++++++++++++++ .github/workflows/packaging.yml | 85 ------------------------------ .github/workflows/tests.yml | 91 --------------------------------- 4 files changed, 115 insertions(+), 176 deletions(-) create mode 100644 .github/workflows/cd.yml create mode 100644 .github/workflows/ci.yml delete mode 100644 .github/workflows/packaging.yml delete mode 100644 .github/workflows/tests.yml diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml new file mode 100644 index 00000000..6e0e4cdb --- /dev/null +++ b/.github/workflows/cd.yml @@ -0,0 +1,73 @@ +name: Continuous Deployment + +on: + workflow_dispatch: + pull_request: + push: + branches: + - main + release: + types: + - published + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +env: + FORCE_COLOR: 3 + +jobs: + dist: + name: Distribution build + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - uses: hynek/build-and-inspect-python-package@v2 + + # Upload to Test PyPI on every commit on main. + test-publish: + needs: [dist] + name: Test Publish to TestPyPI + environment: pypi + permissions: + id-token: write + runs-on: ubuntu-latest + if: + github.repository_owner == 'adrn' && github.event_name == 'push' && + github.ref == 'refs/heads/main' + + steps: + - name: Download packages built by build-and-inspect-python-package + uses: actions/download-artifact@v4 + with: + name: Packages + path: dist + + - name: Upload package to Test PyPI + uses: pypa/gh-action-pypi-publish@release/v1 + with: + repository-url: https://test.pypi.org/legacy/ + + publish: + needs: [dist] + name: Publish to PyPI + environment: pypi + permissions: + id-token: write + runs-on: ubuntu-latest + if: github.event_name == 'release' && github.event.action == 'published' + + steps: + - name: Download packages built by build-and-inspect-python-package + uses: actions/download-artifact@v4 + with: + name: Packages + path: dist + + - uses: pypa/gh-action-pypi-publish@release/v1 + if: github.event_name == 'release' && github.event.action == 'published' \ No newline at end of file diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 00000000..1379a942 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,42 @@ +name: Tests + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +env: + FORCE_COLOR: 3 + +jobs: + test: + runs-on: ${{ matrix.os }} + strategy: + matrix: + python-version: ["3.9", "3.10", "3.11", "3.12"] + os: ["ubuntu-latest", "macos-latest"] + steps: + + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: | + python -m pip install -U pip + python -m pip install -e .[test] + + - name: Test package + run: >- + python -m pytest -ra --cov --cov-report=xml --cov-report=term --durations=20 + + - name: Upload coverage report + uses: codecov/codecov-action@v4.1.0 \ No newline at end of file diff --git a/.github/workflows/packaging.yml b/.github/workflows/packaging.yml deleted file mode 100644 index 6ad1d482..00000000 --- a/.github/workflows/packaging.yml +++ /dev/null @@ -1,85 +0,0 @@ -name: Packaging -on: - release: - types: - - published - pull_request: - branches: - - main - -env: - CIBW_BUILD: "cp38-* cp39-*" - CIBW_SKIP: "*-win32 *musllinux* *i686*" - CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014 - -jobs: - # build_wheels: - # name: Build ${{ matrix.python-version }} wheels on ${{ matrix.os }} - # runs-on: ${{ matrix.os }} - # strategy: - # fail-fast: false - # matrix: - # os: [ubuntu-latest, macos-latest] # , windows-latest] - - # steps: - # - uses: actions/checkout@v4 - # with: - # fetch-depth: 0 - - # - uses: actions/setup-python@v3 - # name: Install Python - # with: - # # Note: cibuildwheel builds for many Python versions beyond this one - # python-version: "3.9" - - # - name: Install MSVC / Visual C++ - # if: runner.os == 'Windows' - # uses: ilammy/msvc-dev-cmd@v1 - - # - name: Build wheels - # run: | - # python -m pip install cibuildwheel - # python -m cibuildwheel --output-dir wheelhouse - - # - uses: actions/upload-artifact@v2 - # with: - # path: ./wheelhouse/*.whl - - build_sdist: - name: Build source distribution - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - with: - fetch-depth: 0 - - - uses: actions/setup-python@v5 - name: Install Python - with: - python-version: "3.9" - - - name: Build sdist - run: | - python -m pip install build - python -m build --sdist - - - uses: actions/upload-artifact@v2 - with: - path: dist/*.tar.gz - - upload_pypi: - # needs: [build_wheels, build_sdist] - needs: [build_sdist] - runs-on: ubuntu-latest - if: github.event_name == 'release' && github.event.action == 'published' - steps: - - uses: actions/download-artifact@v4 - with: - name: artifact - path: dist - - - uses: pypa/gh-action-pypi-publish@master - with: - user: __token__ - password: ${{ secrets.pypi_password }} - # To test: repository_url: https://test.pypi.org/legacy/ \ No newline at end of file diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml deleted file mode 100644 index 5852e0f8..00000000 --- a/.github/workflows/tests.yml +++ /dev/null @@ -1,91 +0,0 @@ -name: Mac/Linux tests -on: - push: - branches: [main] - pull_request: - branches: [main] - schedule: - - cron: "0 11 * * 1" # Mondays @ 7AM Eastern - -jobs: - tests: - name: ${{ matrix.name }} - runs-on: ${{ matrix.os }} - strategy: - fail-fast: true - matrix: - include: - - - name: Python 3.9 with minimal dependencies and coverage - os: ubuntu-latest - python: "3.9" - toxenv: py39-test-cov - - - name: Python 3.9 - os: ubuntu-latest - python: '3.9' - toxenv: py39-test - - - name: Python 3.9 dev dependencies (allowed failure! check logs) - os: ubuntu-latest - python: 3.9 - toxenv: py39-test-devdeps - toxposargs: --durations=50 || true # override exit code - - # Mac: - - name: Python 3.9 standard tests (macOS) - os: macos-latest - python: "3.9" - toxenv: py39-test - - # Older Python versions: - - name: Python 3.9 with oldest supported version of all dependencies - os: ubuntu-latest - python: 3.9 - toxenv: py39-test-oldestdeps - - steps: - - uses: actions/checkout@v4 - with: - fetch-depth: 0 - - - name: Set up Python ${{ matrix.python-version }} on ${{ matrix.os }} - if: "!startsWith(matrix.os, 'windows')" - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python }} - - # Mac: - - name: Setup Mac - hdf5 - if: startsWith(matrix.os, 'mac') - run: | - brew install hdf5 - brew install c-blosc # See: https://github.com/PyTables/PyTables/issues/828 - brew link c-blosc - - # Any *nix: - - name: Install Python dependencies - nix - if: "!startsWith(matrix.os, 'windows')" - run: python -m pip install --upgrade tox - - - name: Run tests - nix - if: "!startsWith(matrix.os, 'windows')" - run: tox ${{ matrix.toxargs }} -e ${{ matrix.toxenv }} -- ${{ matrix.toxposargs }} - - - name: Check coverage.yml existence - if: "endsWith(matrix.name, 'coverage')" - id: check_files - uses: andstor/file-existence-action@v1 - with: - files: "coverage.xml" - - - name: Upload coverage report to codecov - uses: codecov/codecov-action@v1 - if: steps.check_files.outputs.files_exists == 'true' && endsWith(matrix.name, 'coverage') - with: - file: ./coverage.xml # optional - - - uses: actions/upload-artifact@v4 - if: ${{ failure() }} - with: - path: ./result_images From b4ac621ead1781f5c6ee4baf958dff5d1d42f13b Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 4 Mar 2024 19:01:44 -0500 Subject: [PATCH 19/50] try getting run notebooks working --- .github/workflows/tutorials.yml | 51 +++++++++++++++++++++++++++++++++ docs/run_notebooks.py | 15 ++++------ 2 files changed, 57 insertions(+), 9 deletions(-) create mode 100644 .github/workflows/tutorials.yml diff --git a/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml new file mode 100644 index 00000000..d46bc207 --- /dev/null +++ b/.github/workflows/tutorials.yml @@ -0,0 +1,51 @@ +name: Tutorials +on: + push: + branches: + - main + pull_request: + branches: + - main + release: + types: + - published + +jobs: + notebooks: + name: "Build the notebooks for the docs" + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + submodules: true + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install dependencies + run: | + python -m pip install -U pip wheel + python -m pip install ".[tutorials]" + + - name: Execute the notebooks + run: | + cd docs + CACHEDIR=$pwd/cache + rm -rf $CACHEDIR + export PYTENSOR_FLAGS=base_compiledir=$CACHEDIR + python run_notebooks.py + + - uses: actions/upload-artifact@v3 + with: + name: notebooks-for-${{ github.sha }} + path: docs/examples + + - name: Trigger RTDs build + if: ${{ github.event_name != 'pull_request' }} + uses: dfm/rtds-action@v1.1.0 + with: + webhook_url: ${{ secrets.RTDS_WEBHOOK_URL }} + webhook_token: ${{ secrets.RTDS_WEBHOOK_TOKEN }} + commit_ref: ${{ github.ref }} \ No newline at end of file diff --git a/docs/run_notebooks.py b/docs/run_notebooks.py index b25045af..7a174b3f 100644 --- a/docs/run_notebooks.py +++ b/docs/run_notebooks.py @@ -2,21 +2,18 @@ # Standard library import glob +import logging import os import sys -import logging def process_notebook(filename, kernel_name=None): import nbformat from nbconvert.preprocessors import CellExecutionError, ExecutePreprocessor - path = os.path.join( - os.path.abspath("theano_cache"), "p{0}".format(os.getpid()) - ) + path = os.path.join(os.path.abspath("cache"), f"p{os.getpid()}") os.makedirs(path, exist_ok=True) - os.environ["THEANO_FLAGS"] = "base_compiledir={0}".format(path) - os.environ["AESARA_FLAGS"] = os.environ["THEANO_FLAGS"] + os.environ["PYTENSOR_FLAGS"] = f"base_compiledir={path}" with open(filename) as f: notebook = nbformat.read(f, as_version=4) @@ -28,7 +25,7 @@ def process_notebook(filename, kernel_name=None): try: ep.preprocess(notebook, {"metadata": {"path": "examples/"}}) except CellExecutionError as e: - msg = "error while running: {0}\n\n".format(filename) + msg = f"error while running: {filename}\n\n" msg += e.traceback print(msg) finally: @@ -36,14 +33,14 @@ def process_notebook(filename, kernel_name=None): nbformat.write(notebook, f) -if __name__ == '__main__': +if __name__ == "__main__": if len(sys.argv) == 2: pattern = sys.argv[1] else: pattern = "examples/*.ipynb" - nbsphinx_kernel_name = os.environ.get('NBSPHINX_KERNEL', 'python3') + nbsphinx_kernel_name = os.environ.get("NBSPHINX_KERNEL", "python3") for filename in sorted(glob.glob(pattern)): process_notebook(filename, kernel_name=nbsphinx_kernel_name) From d1221b8ec4cf84a4b9adb1115b46c57f5fedf2c2 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 4 Mar 2024 19:02:08 -0500 Subject: [PATCH 20/50] remove circleci stuff --- .circleci/branch_name_check.sh | 12 ----- .circleci/ci_skip_check.sh | 7 --- .circleci/config.yml | 95 ---------------------------------- .circleci/execute_notebooks.sh | 25 --------- 4 files changed, 139 deletions(-) delete mode 100644 .circleci/branch_name_check.sh delete mode 100644 .circleci/ci_skip_check.sh delete mode 100644 .circleci/config.yml delete mode 100644 .circleci/execute_notebooks.sh diff --git a/.circleci/branch_name_check.sh b/.circleci/branch_name_check.sh deleted file mode 100644 index 9bb9bcba..00000000 --- a/.circleci/branch_name_check.sh +++ /dev/null @@ -1,12 +0,0 @@ -#!/bin/bash - -if [ -z "${CIRCLE_PULL_REQUEST}" ] && [ "${CIRCLE_BRANCH}" != "main" ] && [[ ! "$CIRCLE_BRANCH" =~ ^v[0-9\.]+ ]]; -then - echo "HALTING" - circleci step halt -else - echo "Continuing with tests" -fi - -echo $CIRCLE_PULL_REQUEST -echo $CIRCLE_BRANCH diff --git a/.circleci/ci_skip_check.sh b/.circleci/ci_skip_check.sh deleted file mode 100644 index 6890bc29..00000000 --- a/.circleci/ci_skip_check.sh +++ /dev/null @@ -1,7 +0,0 @@ -#!/bin/bash - -commitmessage=$(git log --pretty=%B -n 1) -if [[ $commitmessage = *"[ci skip]"* ]] || [[ $commitmessage = *"[skip ci]"* ]]; then - echo "Skipping build because [ci skip] found in commit message" - circleci step halt -fi diff --git a/.circleci/config.yml b/.circleci/config.yml deleted file mode 100644 index ac9756b3..00000000 --- a/.circleci/config.yml +++ /dev/null @@ -1,95 +0,0 @@ -# Python CircleCI 2.0 configuration file -# Check https://circleci.com/docs/2.0/language-python/ for more details -version: 2.1 - -apt-install: &apt-install - name: Install apt packages - command: | - sudo apt-get update - sudo apt-get install build-essential pandoc - -pip-install: &pip-install - name: Install Python dependencies - command: | - python3 -m venv venv - . venv/bin/activate - pip install -U pip - -pip-install-dev: &pip-install-dev - name: Install Python development dependencies - command: | - python3 -m venv venv - . venv/bin/activate - pip install -U pip - pip install --progress-bar off tox - -main-or-pr-check: &main-or-pr-check - name: Only build branches that are "main", or a PR - command: bash .circleci/branch_name_check.sh - -ci-skip-check: &ci-skip-check - name: Check for [ci skip] - command: bash .circleci/ci_skip_check.sh - -jobs: - - test-39: - docker: - - image: circleci/python:3.9 - steps: - - checkout - - run: *ci-skip-check - - run: *main-or-pr-check - - run: *apt-install - - run: *pip-install-dev - - run: - name: Run test - command: | - . venv/bin/activate - tox -e py39 - - # docs build - build-docs: - docker: - - image: circleci/python:3.9 - steps: - - checkout - - add_ssh_keys: # add GitHub SSH keys - fingerprints: - - 21:3a:31:0c:f7:f4:94:ed:1c:4f:16:cb:67:60:61:0a - - run: *ci-skip-check - - run: *main-or-pr-check - - run: *apt-install - - run: *pip-install-dev - - run: - name: Install - command: | - . venv/bin/activate - pip install -e .[docs] - - run: - name: Execute notebooks - command: bash .circleci/execute_notebooks.sh - - run: - name: Build documentation - no_output_timeout: 30m - command: | - . venv/bin/activate - tox -e build_docs - - store_artifacts: - path: docs/_build/html - - run: - name: Built documentation is available at - command: | - DOCS_URL="${CIRCLE_BUILD_URL}/artifacts/${CIRCLE_NODE_INDEX}/${CIRCLE_WORKING_DIRECTORY/#\~/$HOME}/docs/_build/html/index.html"; echo $DOCS_URL - -workflows: - thejoker: - jobs: - - test-39 - - build-docs: - requires: - - test-39 - -notify: - webhooks: - - url: https://giles.cadair.dev/circleci diff --git a/.circleci/execute_notebooks.sh b/.circleci/execute_notebooks.sh deleted file mode 100644 index 97982e59..00000000 --- a/.circleci/execute_notebooks.sh +++ /dev/null @@ -1,25 +0,0 @@ -#!/bin/bash -eux - -if [[ -z $CIRCLE_PULL_REQUEST ]]; then - CACHEDIR=`pwd`/theano_cache - rm -rf $CACHEDIR - export THEANO_FLAGS=base_compiledir=$CACHEDIR - export AESARA_FLAGS=base_compiledir=$CACHEDIR - - git branch -D executed-notebooks || true - git checkout -b executed-notebooks - - . venv/bin/activate - cd docs - python run_notebooks.py - - git add examples/*.ipynb - git -c user.name='circleci' -c user.email='circleci' commit -m "now with executed tutorials" - - git push -f origin executed-notebooks - - echo "Not a pull request: pushing run notebooks branch." -else - echo $CIRCLE_PULL_REQUEST - echo "This is a pull request: not pushing executed notebooks." -fi \ No newline at end of file From b76395366e9e554b0b1b662b937b4b66471265e2 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 4 Mar 2024 19:06:31 -0500 Subject: [PATCH 21/50] remove python 3.12 --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1379a942..51d6f288 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -18,7 +18,7 @@ jobs: runs-on: ${{ matrix.os }} strategy: matrix: - python-version: ["3.9", "3.10", "3.11", "3.12"] + python-version: ["3.9", "3.10", "3.11"] os: ["ubuntu-latest", "macos-latest"] steps: From e65ce4a40243e427c5f063d169426c537993c08d Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 4 Mar 2024 20:31:24 -0500 Subject: [PATCH 22/50] old nb setup --- docs/examples/notebook_setup.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/docs/examples/notebook_setup.py b/docs/examples/notebook_setup.py index 376401b5..e81a2a48 100644 --- a/docs/examples/notebook_setup.py +++ b/docs/examples/notebook_setup.py @@ -1,7 +1,3 @@ -from IPython.display import set_matplotlib_formats -set_matplotlib_formats('retina') - -import logging import warnings import matplotlib.pyplot as plt @@ -10,8 +6,8 @@ warnings.filterwarnings("ignore", category=FutureWarning) -logger = logging.getLogger("theano.gof.compilelock") -logger.setLevel(logging.ERROR) +# logger = logging.getLogger("pytensor.gof.compilelock") +# logger.setLevel(logging.ERROR) plt.style.use("default") From 992b8d73fa2a34f6de9f8fd095fe2ae567303833 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 4 Mar 2024 20:34:16 -0500 Subject: [PATCH 23/50] tutorials install flag --- pyproject.toml | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 67b0bd21..f98a534d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,7 +53,12 @@ docs = [ "nbsphinx", "nbconvert", "nbformat", - "ipykernel" + "ipykernel", + "matplotlib", +] +tutorials = [ + "thejoker[docs]", + "jupyter-client" ] [tool.setuptools.packages.find] From a870eae8c8ba26fed4566a156f9c220325149b6f Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 4 Mar 2024 21:42:18 -0500 Subject: [PATCH 24/50] fix mcmc sampling --- .pre-commit-config.yaml | 2 +- thejoker/_keplerian_orbit.py | 14 +++++-- thejoker/distributions.py | 3 ++ thejoker/tests/test_sampler.py | 69 +++++++++++++++------------------- thejoker/thejoker.py | 8 ++-- 5 files changed, 49 insertions(+), 47 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 431d7c88..a924eb7f 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,6 @@ repos: - repo: https://github.com/kynan/nbstripout - rev: main + rev: 0.7.1 hooks: - id: nbstripout files: ".ipynb" diff --git a/thejoker/_keplerian_orbit.py b/thejoker/_keplerian_orbit.py index be0c2bab..8a326757 100644 --- a/thejoker/_keplerian_orbit.py +++ b/thejoker/_keplerian_orbit.py @@ -6,9 +6,11 @@ import numpy as np import pytensor.tensor as pt from astropy import units as u +from astropy.constants import G from exoplanet_core.pymc import ops tt = pt +G_grav = G.to(u.R_sun**3 / u.M_sun / u.day**2).value class KeplerianOrbit: @@ -67,10 +69,16 @@ def __init__( model=None, **kwargs, ): - self.a = a self.period = period self.m_planet = tt.zeros_like(period) - # self.m_total = self.m_star + self.m_planet + self.m_star = tt.ones_like(period) + self.m_total = self.m_star + self.m_planet + + if a is None: + a = ( + G_grav * (self.m_star + self.m_planet) * self.period**2 / (4 * np.pi**2) + ) ** (1.0 / 3) + self.a = a self.n = 2 * np.pi / self.period self.K0 = self.n * self.a / self.m_total @@ -115,7 +123,7 @@ def __init__( ome2 = 1 - self.ecc**2 self.K0 /= tt.sqrt(ome2) - zla = tt.zeros_like(self.P) + zla = tt.zeros_like(self.period) self.incl = 0.5 * np.pi + zla self.cos_incl = zla self.b = zla diff --git a/thejoker/distributions.py b/thejoker/distributions.py index 07f22cf6..b8441b02 100644 --- a/thejoker/distributions.py +++ b/thejoker/distributions.py @@ -40,6 +40,9 @@ def support_point(rv, size, a, b): a, b = pt.broadcast_arrays(a, b) return 0.5 * (a + b) + # TODO: remove this once new pymc version is released + moment = support_point + def logp(value, a, b): _fac = pt.log(b) - pt.log(a) res = -pt.as_tensor_variable(value) - pt.log(_fac) diff --git a/thejoker/tests/test_sampler.py b/thejoker/tests/test_sampler.py index 1d6cd8b8..2f70f228 100644 --- a/thejoker/tests/test_sampler.py +++ b/thejoker/tests/test_sampler.py @@ -2,6 +2,7 @@ import astropy.units as u import numpy as np +import pymc as pm import pytest from astropy.time import Time from schwimmbad import MultiPool, SerialPool @@ -107,24 +108,24 @@ def test_marginal_ln_likelihood(tmpdir, case): assert len(ll) == len(prior_samples) -@pytest.mark.parametrize( - "prior", - [ - JokerPrior.default( - P_min=5 * u.day, - P_max=500 * u.day, - sigma_K0=25 * u.km / u.s, - sigma_v=100 * u.km / u.s, - ), - JokerPrior.default( - P_min=5 * u.day, - P_max=500 * u.day, - sigma_K0=25 * u.km / u.s, - poly_trend=2, - sigma_v=[100 * u.km / u.s, 0.5 * u.km / u.s / u.day], - ), - ], -) +priors = [ + JokerPrior.default( + P_min=5 * u.day, + P_max=500 * u.day, + sigma_K0=25 * u.km / u.s, + sigma_v=100 * u.km / u.s, + ), + JokerPrior.default( + P_min=5 * u.day, + P_max=500 * u.day, + sigma_K0=25 * u.km / u.s, + poly_trend=2, + sigma_v=[100 * u.km / u.s, 0.5 * u.km / u.s / u.day], + ), +] + + +@pytest.mark.parametrize("prior", priors) def test_rejection_sample(tmpdir, prior): data, orbit = make_data() flat_data, orbit = make_data(K=0.1 * u.m / u.s) @@ -172,24 +173,7 @@ def test_rejection_sample(tmpdir, prior): assert u.allclose(all_Ks[0], all_Ks[i]) -@pytest.mark.parametrize( - "prior", - [ - JokerPrior.default( - P_min=5 * u.day, - P_max=500 * u.day, - sigma_K0=25 * u.km / u.s, - sigma_v=100 * u.km / u.s, - ), - JokerPrior.default( - P_min=5 * u.day, - P_max=500 * u.day, - sigma_K0=25 * u.km / u.s, - poly_trend=2, - sigma_v=[100 * u.km / u.s, 0.5 * u.km / u.s / u.day], - ), - ], -) +@pytest.mark.parametrize("prior", priors) def test_iterative_rejection_sample(tmpdir, prior): data, orbit = make_data(n_times=3) @@ -227,7 +211,14 @@ def test_iterative_rejection_sample(tmpdir, prior): assert u.allclose(all_Ks[0], all_Ks[i]) -if __name__ == "__main__": - import pathlib +@pytest.mark.parametrize("prior", priors) +def test_continue_mcmc(prior): + data, orbit = make_data(n_times=10) - test_marginal_ln_likelihood(pathlib.Path("/tmp/"), 0) + prior_samples = prior.sample(size=16384, return_logprobs=True) + joker = TheJoker(prior) + joker_samples = joker.rejection_sample(data, prior_samples) + + with prior.model: + mcmc_init = joker.setup_mcmc(data, joker_samples) + trace = pm.sample(tune=500, draws=500, initvals=mcmc_init, cores=1, chains=1) diff --git a/thejoker/thejoker.py b/thejoker/thejoker.py index ea0aa595..3c082e2b 100644 --- a/thejoker/thejoker.py +++ b/thejoker/thejoker.py @@ -473,15 +473,15 @@ def setup_mcmc(self, data, joker_samples, model=None, custom_func=None): pm.Deterministic("model_rv", rv_model) err = pt.sqrt(err**2 + p["s"] ** 2) - pm.Normal("obs", mu=rv_model, sd=err, observed=y) + pm.Normal("obs", mu=rv_model, sigma=err, observed=y) - pm.Deterministic("logp", model.logpt) + pm.Deterministic("logp", model.logp()) dist = pm.Normal.dist(model.model_rv, data.rv_err.value) lnlike = pm.Deterministic( - "ln_likelihood", dist.logp(data.rv.value).sum(axis=-1) + "ln_likelihood", pm.logp(dist, data.rv.value).sum(axis=-1) ) - pm.Deterministic("ln_prior", model.logpt - lnlike) + pm.Deterministic("ln_prior", model.logp() - lnlike) return mcmc_init From 8a82c22bb9724edef482b6a774806e0dba0217e2 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Tue, 5 Mar 2024 18:33:48 -0500 Subject: [PATCH 25/50] use more of keplerianorbit and deprecate phase_fold argument --- thejoker/_keplerian_orbit.py | 671 +++++++++++++++++++++++++++++-- thejoker/data.py | 8 +- thejoker/samples.py | 61 +++ thejoker/samples_helpers.py | 99 ----- thejoker/src/fast_likelihood.pyx | 14 +- thejoker/thejoker.py | 9 +- 6 files changed, 722 insertions(+), 140 deletions(-) diff --git a/thejoker/_keplerian_orbit.py b/thejoker/_keplerian_orbit.py index 8a326757..1b584334 100644 --- a/thejoker/_keplerian_orbit.py +++ b/thejoker/_keplerian_orbit.py @@ -1,16 +1,31 @@ -"""A port from exoplanet, to support pymc recent versions""" - -__all__ = ["KeplerianOrbit"] +__all__ = [ + "KeplerianOrbit", + "get_true_anomaly", + "get_aor_from_transit_duration", +] +import warnings +from collections import defaultdict +import astropy.constants as c import numpy as np -import pytensor.tensor as pt +import pytensor.tensor as tt from astropy import units as u -from astropy.constants import G from exoplanet_core.pymc import ops +from pytensor import ifelse +from pytensor.tensor import as_tensor_variable + +from thejoker.units import has_unit, to_unit, with_unit + +G_grav = c.G.to(u.R_sun**3 / u.M_sun / u.day**2).value +gcc_per_sun = (u.M_sun / u.R_sun**3).to(u.g / u.cm**3) +au_per_R_sun = u.R_sun.to(u.au) +c_light = c.c.to(u.R_sun / u.day).value -tt = pt -G_grav = G.to(u.R_sun**3 / u.M_sun / u.day**2).value +day_per_yr_over_2pi = ( + (1.0 * u.au) ** (3 / 2) + / (np.sqrt(c.G.to(u.au**3 / (u.M_sun * u.day**2)) * (1.0 * u.M_sun))) +).value class KeplerianOrbit: @@ -52,41 +67,110 @@ class KeplerianOrbit: m_star: The mass of the star in ``M_sun``. r_star: The radius of the star in ``R_sun``. rho_star: The density of the star in units of ``rho_star_units``. + m_planet_units: An ``astropy.units`` compatible unit object giving the + units of the planet masses. If not given, the default is ``M_sun``. + rho_star_units: An ``astropy.units`` compatible unit object giving the + units of the stellar density. If not given, the default is + ``g / cm^3``. """ + __citations__ = ("astropy",) + def __init__( self, period=None, a=None, t0=None, t_periastron=None, + incl=None, + b=None, + duration=None, ecc=None, omega=None, sin_omega=None, cos_omega=None, Omega=None, + m_planet=0.0, + m_star=None, + r_star=None, + rho_star=None, + ror=None, model=None, **kwargs, ): - self.period = period - self.m_planet = tt.zeros_like(period) - self.m_star = tt.ones_like(period) + if "m_planet_units" in kwargs: + # deprecation_warning( + # "'m_planet_units' is deprecated; Use `with_unit` instead" + # ) + m_planet = with_unit(m_planet, kwargs.pop("m_planet_units")) + if "rho_star_units" in kwargs: + # deprecation_warning( + # "'rho_star_units' is deprecated; Use `with_unit` instead" + # ) + rho_star = with_unit(rho_star, kwargs.pop("rho_star_units")) + + self.jacobians = defaultdict(lambda: defaultdict(None)) + + daordtau = None + if ecc is None and duration is not None: + if r_star is None: + r_star = as_tensor_variable(1.0) + if b is None: + raise ValueError( + "'b' must be provided for a circular orbit with a " "'duration'" + ) + if ror is None: + warnings.warn( + "When using the 'duration' parameter in KeplerianOrbit, " + "the 'ror' parameter should also be provided.", + UserWarning, + ) + aor, daordtau = get_aor_from_transit_duration(duration, period, b, ror=ror) + a = r_star * aor + duration = None + + inputs = _get_consistent_inputs(a, period, rho_star, r_star, m_star, m_planet) + ( + self.a, + self.period, + self.rho_star, + self.r_star, + self.m_star, + self.m_planet, + ) = inputs self.m_total = self.m_star + self.m_planet - if a is None: - a = ( - G_grav * (self.m_star + self.m_planet) * self.period**2 / (4 * np.pi**2) - ) ** (1.0 / 3) - self.a = a - self.n = 2 * np.pi / self.period + self.a_star = self.a * self.m_planet / self.m_total + self.a_planet = -self.a * self.m_star / self.m_total + + # Track the Jacobian between the duration and a + if daordtau is not None: + dadtau = self.r_star * daordtau + self.jacobians["duration"]["a"] = dadtau + self.jacobians["duration"]["a_star"] = dadtau * self.m_planet / self.m_total + self.jacobians["duration"]["a_planet"] = ( + -dadtau * self.m_star / self.m_total + ) + + # rho = 3 * pi * (a/R)**3 / (G * P**2) + # -> drho / d(a/R) = 9 * pi * (a/R)**2 / (G * P**2) + self.jacobians["duration"]["rho_star"] = ( + 9 + * np.pi + * (self.a / self.r_star) ** 2 + * daordtau + * gcc_per_sun + / (G_grav * self.period**2) + ) + self.K0 = self.n * self.a / self.m_total if Omega is None: self.Omega = None else: - self.Omega = pt.as_tensor_variable(Omega) + self.Omega = as_tensor_variable(Omega) self.cos_Omega = tt.cos(self.Omega) self.sin_Omega = tt.sin(self.Omega) @@ -94,20 +178,21 @@ def __init__( if ecc is None: self.ecc = None self.M0 = 0.5 * np.pi + tt.zeros_like(self.n) + incl_factor = 1 else: - self.ecc = pt.as_tensor_variable(ecc) + self.ecc = as_tensor_variable(ecc) if omega is not None: if sin_omega is not None and cos_omega is not None: raise ValueError( "either 'omega' or 'sin_omega' and 'cos_omega' can be " "provided" ) - self.omega = pt.as_tensor_variable(omega) + self.omega = as_tensor_variable(omega) self.cos_omega = tt.cos(self.omega) self.sin_omega = tt.sin(self.omega) elif sin_omega is not None and cos_omega is not None: - self.cos_omega = pt.as_tensor_variable(cos_omega) - self.sin_omega = pt.as_tensor_variable(sin_omega) + self.cos_omega = as_tensor_variable(cos_omega) + self.sin_omega = as_tensor_variable(sin_omega) self.omega = tt.arctan2(self.sin_omega, self.cos_omega) else: @@ -122,11 +207,54 @@ def __init__( ome2 = 1 - self.ecc**2 self.K0 /= tt.sqrt(ome2) - - zla = tt.zeros_like(self.period) - self.incl = 0.5 * np.pi + zla - self.cos_incl = zla - self.b = zla + incl_factor = (1 + self.ecc * self.sin_omega) / ome2 + + # The Jacobian for the transform cos(i) -> b + self.dcosidb = self.jacobians["b"]["cos_incl"] = ( + incl_factor * self.r_star / self.a + ) + + if b is not None: + if incl is not None or duration is not None: + raise ValueError("only one of 'incl', 'b', and 'duration' can be given") + self.b = as_tensor_variable(b) + self.cos_incl = self.dcosidb * self.b + self.incl = tt.arccos(self.cos_incl) + elif incl is not None: + if duration is not None: + raise ValueError("only one of 'incl', 'b', and 'duration' can be given") + self.incl = as_tensor_variable(incl) + self.cos_incl = tt.cos(self.incl) + self.b = self.cos_incl / self.dcosidb + elif duration is not None: + # This assertion should never be hit because of the first + # conditional in this method, but let's keep it here anyways + assert self.ecc is not None + + self.duration = as_tensor_variable(to_unit(duration, u.day)) + c = tt.sin(np.pi * self.duration * incl_factor / self.period) + c2 = c * c + aor = self.a_planet / self.r_star + esinw = self.ecc * self.sin_omega + self.b = tt.sqrt( + (aor**2 * c2 - 1) + / ( + c2 * esinw**2 + + 2 * c2 * esinw + + c2 + - self.ecc**4 + + 2 * self.ecc**2 + - 1 + ) + ) + self.b *= 1 - self.ecc**2 + self.cos_incl = self.dcosidb * self.b + self.incl = tt.arccos(self.cos_incl) + else: + zla = tt.zeros_like(self.a) + self.incl = 0.5 * np.pi + zla + self.cos_incl = zla + self.b = zla if t0 is not None and t_periastron is not None: raise ValueError("you can't define both t0 and t_periastron") @@ -134,10 +262,10 @@ def __init__( t0 = tt.zeros_like(self.period) if t0 is None: - self.t_periastron = pt.as_tensor_variable(t_periastron) + self.t_periastron = as_tensor_variable(t_periastron) self.t0 = self.t_periastron + self.M0 / self.n else: - self.t0 = pt.as_tensor_variable(t0) + self.t0 = as_tensor_variable(t0) self.t_periastron = self.t0 - self.M0 / self.n self.tref = self.t_periastron - self.t0 @@ -197,6 +325,236 @@ def _get_true_anomaly(self, t, _pad=True): sinf, cosf = ops.kepler(M, self.ecc + tt.zeros_like(M)) return sinf, cosf + def _get_position_and_velocity(self, t, parallax=None): + sinf, cosf = self._get_true_anomaly(t) + + if self.ecc is None: + r = 1.0 + vx, vy, vz = self._rotate_vector(-self.K0 * sinf, self.K0 * cosf) + else: + r = (1.0 - self.ecc**2) / (1 + self.ecc * cosf) + vx, vy, vz = self._rotate_vector( + -self.K0 * sinf, self.K0 * (cosf + self.ecc) + ) + + x, y, z = self._rotate_vector(r * cosf, r * sinf) + + pos = tt.stack((x, y, z), axis=-1) + pos = tt.concatenate( + ( + tt.sum(tt.shape_padright(self.a_star) * pos, axis=0, keepdims=True), + tt.shape_padright(self.a_planet) * pos, + ), + axis=0, + ) + vel = tt.stack((vx, vy, vz), axis=-1) + vel = tt.concatenate( + ( + tt.sum( + tt.shape_padright(self.m_planet) * vel, + axis=0, + keepdims=True, + ), + -tt.shape_padright(self.m_star) * vel, + ), + axis=0, + ) + + if parallax is not None: + # convert r into arcseconds + pos = pos * (parallax * au_per_R_sun) + vel = vel * (parallax * au_per_R_sun) + + return pos, vel + + def _get_position(self, a, t, parallax=None, light_delay=False, _pad=True): + """Get the position of a body. + + Args: + a: the semi-major axis of the orbit. + t: the time (or tensor of times) to calculate the position. + parallax: (arcseconds) if provided, return the position in + units of arcseconds. + light_delay: account for the light travel time delay? Default is + False. + + Returns: + The position of the body in the observer frame. Default is in units + of R_sun, but if parallax is provided, then in units of arcseconds. + + """ + if light_delay: + return self._get_retarded_position(a, t, parallax=None, _pad=_pad) + + sinf, cosf = self._get_true_anomaly(t, _pad=_pad) + if self.ecc is None: + r = a + else: + r = a * (1.0 - self.ecc**2) / (1 + self.ecc * cosf) + + if parallax is not None: + # convert r into arcseconds + r = r * parallax * au_per_R_sun + + return self._rotate_vector(r * cosf, r * sinf) + + def _get_retarded_position(self, a, t, parallax=None, z0=0.0, _pad=True): + """Get the retarded position of a body, accounting for light delay. + + Args: + a: the semi-major axis of the orbit. + t: the time (or tensor of times) to calculate the position. + parallax: (arcseconds) if provided, return the position in + units of arcseconds. + z0: the reference point along the z axis whose light travel time + delay is taken to be zero. Default is the origin. + + Returns: + The position of the body in the observer frame. Default is in units + of R_sun, but if parallax is provided, then in units of arcseconds. + + """ + sinf, cosf = self._get_true_anomaly(t, _pad=_pad) + + # Compute the orbital radius and the component of the velocity + # in the z direction + angvel = 2 * np.pi / self.period + if self.ecc is None: + r = a + vamp = angvel * a + vz = vamp * self.sin_incl * cosf + else: + r = a * (1.0 - self.ecc**2) / (1 + self.ecc * cosf) + vamp = angvel * a / tt.sqrt(1 - self.ecc**2) + cwf = self.cos_omega * cosf - self.sin_omega * sinf + vz = vamp * self.sin_incl * (self.ecc * self.cos_omega + cwf) + + # True position of the body + x, y, z = self._rotate_vector(r * cosf, r * sinf) + + # Component of the acceleration in the z direction + az = -(angvel**2) * (a / r) ** 3 * z + + # Compute the time delay at the **retarded** position, accounting + # for the instantaneous velocity and acceleration of the body. + # See the derivation at https://github.com/rodluger/starry/issues/66 + delay = tt.switch( + tt.lt(tt.abs_(az), 1.0e-10), + (z0 - z) / (c_light + vz), + (c_light / az) + * ( + (1 + vz / c_light) + - tt.sqrt( + (1 + vz / c_light) * (1 + vz / c_light) + - 2 * az * (z0 - z) / c_light**2 + ) + ), + ) + + if _pad: + new_t = tt.shape_padright(t) - delay + else: + new_t = t - delay + + # Re-compute Kepler's equation, this time at the **retarded** position + return self._get_position(a, new_t, parallax, _pad=False) + + def get_planet_position(self, t, parallax=None, light_delay=False): + """The planets' positions in the barycentric frame + + Args: + t: The times where the position should be evaluated. + light_delay: account for the light travel time delay? Default is + False. + + Returns: + The components of the position vector at ``t`` in units of + ``R_sun``. + + """ + return tuple( + tt.squeeze(x) + for x in self._get_position( + self.a_planet, t, parallax, light_delay=light_delay + ) + ) + + def get_star_position(self, t, parallax=None, light_delay=False): + """The star's position in the barycentric frame + + .. note:: If there are multiple planets in the system, this will + return one column per planet with each planet's contribution to + the motion. The star's full position can be computed by summing + over the last axis. + + Args: + t: The times where the position should be evaluated. + light_delay: account for the light travel time delay? Default is + False. + + Returns: + The components of the position vector at ``t`` in units of + ``R_sun``. + + """ + return tuple( + tt.squeeze(x) + for x in self._get_position( + self.a_star, t, parallax, light_delay=light_delay + ) + ) + + def get_relative_position(self, t, parallax=None, light_delay=False): + """The planets' positions relative to the star in the X,Y,Z frame. + + .. note:: This treats each planet independently and does not take the + other planets into account when computing the position of the + star. This is fine as long as the planet masses are small. In + other words, the reflex motion of the star caused by the other + planets is neglected when computing the relative coordinates of + one of the planets. + + Args: + t: The times where the position should be evaluated. + light_delay: account for the light travel time delay? Default is + False. + + Returns: + The components of the position vector at ``t`` in units of + ``R_sun``. + + """ + return tuple( + tt.squeeze(x) + for x in self._get_position(-self.a, t, parallax, light_delay=light_delay) + ) + + def get_relative_angles(self, t, parallax=None, light_delay=False): + """The planets' relative position to the star in the sky plane, in + separation, position angle coordinates. + + .. note:: This treats each planet independently and does not take the + other planets into account when computing the position of the + star. This is fine as long as the planet masses are small. + + Args: + t: The times where the position should be evaluated. + light_delay: account for the light travel time delay? Default is + False. + + Returns: + The separation (arcseconds) and position angle (radians, + measured east of north) of the planet relative to the star. + + """ + X, Y, Z = self._get_position(-self.a, t, parallax, light_delay=light_delay) + + # calculate rho and theta + rho = tt.squeeze(tt.sqrt(X**2 + Y**2)) # arcsec + theta = tt.squeeze(tt.arctan2(Y, X)) # radians between [-pi, pi] + + return (rho, theta) + def _get_velocity(self, m, t): """Get the velocity vector of a body in the observer frame""" sinf, cosf = self._get_true_anomaly(t) @@ -205,6 +563,19 @@ def _get_velocity(self, m, t): return self._rotate_vector(-K * sinf, K * cosf) return self._rotate_vector(-K * sinf, K * (cosf + self.ecc)) + def get_planet_velocity(self, t): + """Get the planets' velocity vectors + + Args: + t: The times where the velocity should be evaluated. + + Returns: + The components of the velocity vector at ``t`` in units of + ``M_sun/day``. + + """ + return tuple(tt.squeeze(x) for x in self._get_velocity(-self.m_star, t)) + def get_star_velocity(self, t): """Get the star's velocity vector @@ -222,6 +593,23 @@ def get_star_velocity(self, t): """ return tuple(tt.squeeze(x) for x in self._get_velocity(self.m_planet, t)) + def get_relative_velocity(self, t): + """The planets' velocity relative to the star + + .. note:: This treats each planet independently and does not take the + other planets into account when computing the position of the + star. This is fine as long as the planet masses are small. + + Args: + t: The times where the velocity should be evaluated. + + Returns: + The components of the velocity vector at ``t`` in units of + ``R_sun/day``. + + """ + return tuple(tt.squeeze(x) for x in self._get_velocity(-self.m_total, t)) + def get_radial_velocity(self, t, K=None, output_units=None): """Get the radial velocity of the star @@ -268,6 +656,127 @@ def get_radial_velocity(self, t, K=None, output_units=None): v = self.get_star_velocity(t) return -conv * v[2] + def _get_acceleration(self, a, m, t): + sinf, cosf = self._get_true_anomaly(t) + K = self.K0 * m + if self.ecc is None: + factor = -(K**2) / a + else: + factor = K**2 * (self.ecc * cosf + 1) ** 2 / (a * (self.ecc**2 - 1)) + return self._rotate_vector(factor * cosf, factor * sinf) + + def get_planet_acceleration(self, t): + return tuple( + tt.squeeze(x) + for x in self._get_acceleration(self.a_planet, -self.m_star, t) + ) + + def get_star_acceleration(self, t): + return tuple( + tt.squeeze(x) for x in self._get_acceleration(self.a_star, self.m_planet, t) + ) + + def get_relative_acceleration(self, t): + return tuple( + tt.squeeze(x) for x in self._get_acceleration(-self.a, -self.m_total, t) + ) + + def in_transit(self, t, r=0.0, texp=None, light_delay=False): + """Get a list of timestamps that are in transit + + Args: + t (vector): A vector of timestamps to be evaluated. + r (Optional): The radii of the planets. + texp (Optional[float]): The exposure time. + + Returns: + The indices of the timestamps that are in transit. + + """ + if light_delay: + raise NotImplementedError( + "Light travel time delay not yet implemented for `in_transit`" + ) + + z = tt.zeros_like(self.a) + r = as_tensor_variable(r) + z + R = self.r_star + z + + # Wrap the times into time since transit + hp = 0.5 * self.period + dt = tt.mod(self._warp_times(t) + hp, self.period) - hp + + if self.ecc is None: + # Equation 14 from Winn (2010) + k = r / R + arg = tt.square(1 + k) - tt.square(self.b) + factor = R / (self.a * self.sin_incl) + hdur = hp * tt.arcsin(factor * tt.sqrt(arg)) / np.pi + t_start = -hdur + t_end = hdur + flag = z + + else: + M_contact = ops.contact_points( + self.a, + self.ecc, + self.cos_omega, + self.sin_omega, + self.cos_incl + z, + self.sin_incl + z, + R + r, + ) + flag = M_contact[2] + + t_start = (M_contact[0] - self.M0) / self.n + t_start = tt.mod(t_start + hp, self.period) - hp + t_end = (M_contact[1] - self.M0) / self.n + t_end = tt.mod(t_end + hp, self.period) - hp + + t_start = tt.switch(tt.gt(t_start, 0.0), t_start - self.period, t_start) + t_end = tt.switch(tt.lt(t_end, 0.0), t_end + self.period, t_end) + + if texp is not None: + t_start -= 0.5 * texp + t_end += 0.5 * texp + + mask = tt.any(tt.and_(dt >= t_start, dt <= t_end), axis=-1) + + result = ifelse( + tt.all(tt.eq(flag, 0)), + tt.arange(t.shape[0])[mask], + tt.arange(t.shape[0]), + ) + + return result + + def _flip(self, r_planet, model=None): + if self.ecc is None: + orbit = type(self)( + period=self.period, + t_periastron=self.t_periastron + 0.5 * self.period, + incl=self.incl, + Omega=self.Omega, + m_star=self.m_planet, + m_planet=self.m_star, + r_star=r_planet, + model=model, + ) + else: + orbit = type(self)( + period=self.period, + t_periastron=self.t_periastron, + incl=self.incl, + ecc=self.ecc, + omega=self.omega - np.pi, + Omega=self.Omega, + m_star=self.m_planet, + m_planet=self.m_star, + r_star=r_planet, + model=model, + ) + return orbit + def get_true_anomaly(M, e, **kwargs): """Get the true anomaly for a tensor of mean anomalies and eccentricities @@ -282,3 +791,109 @@ def get_true_anomaly(M, e, **kwargs): """ sinf, cosf = ops.kepler(M, e) return tt.arctan2(sinf, cosf) + + +def get_aor_from_transit_duration(duration, period, b, ror=None): + """Get the semimajor axis implied by a circular orbit and duration + + Args: + duration: The transit duration + period: The orbital period + b: The impact parameter of the transit + ror: The radius ratio of the planet to the star + + Returns: + The semimajor axis in units of the stellar radius and the Jacobian + ``d a / d duration`` + + """ + if ror is None: + ror = as_tensor_variable(0.0) + b2 = b**2 + opk2 = (1 + ror) ** 2 + phi = np.pi * duration / period + sinp = tt.sin(phi) + cosp = tt.cos(phi) + num = tt.sqrt(opk2 - b2 * cosp**2) + aor = num / sinp + grad = np.pi * cosp * (b2 - opk2) / (num * period * sinp**2) + return aor, grad + + +def _get_consistent_inputs(a, period, rho_star, r_star, m_star, m_planet): + if a is None and period is None: + raise ValueError("values must be provided for at least one of a " "and period") + + if m_planet is not None: + m_planet = as_tensor_variable(to_unit(m_planet, u.M_sun)) + + if a is not None: + a = as_tensor_variable(to_unit(a, u.R_sun)) + if m_planet is None: + m_planet = tt.zeros_like(a) + if period is not None: + period = as_tensor_variable(to_unit(period, u.day)) + if m_planet is None: + m_planet = tt.zeros_like(period) + + # Compute the implied density if a and period are given + implied_rho_star = False + if a is not None and period is not None: + if rho_star is not None or m_star is not None: + raise ValueError( + "if both a and period are given, you can't " + "also define rho_star or m_star" + ) + + # Default to a stellar radius of 1 if not provided + if r_star is None: + r_star = as_tensor_variable(1.0) + else: + r_star = as_tensor_variable(to_unit(r_star, u.R_sun)) + + # Compute the implied mass via Kepler's 3rd law + m_tot = 4 * np.pi * np.pi * a**3 / (G_grav * period**2) + + # Compute the implied density + m_star = m_tot - m_planet + vol_star = 4 * np.pi * r_star**3 / 3.0 + rho_star = m_star / vol_star + implied_rho_star = True + + # Make sure that the right combination of stellar parameters are given + if r_star is None and m_star is None: + r_star = 1.0 + if rho_star is None: + m_star = 1.0 + if (not implied_rho_star) and sum( + arg is None for arg in (rho_star, r_star, m_star) + ) != 1: + raise ValueError( + "values must be provided for exactly two of " "rho_star, m_star, and r_star" + ) + + if rho_star is not None and not implied_rho_star: + if has_unit(rho_star): + rho_star = as_tensor_variable(to_unit(rho_star, u.M_sun / u.R_sun**3)) + else: + rho_star = as_tensor_variable(rho_star) / gcc_per_sun + if r_star is not None: + r_star = as_tensor_variable(to_unit(r_star, u.R_sun)) + if m_star is not None: + m_star = as_tensor_variable(to_unit(m_star, u.M_sun)) + + # Work out the stellar parameters + if rho_star is None: + rho_star = 3 * m_star / (4 * np.pi * r_star**3) + elif r_star is None: + r_star = (3 * m_star / (4 * np.pi * rho_star)) ** (1 / 3) + elif m_star is None: + m_star = 4 * np.pi * r_star**3 * rho_star / 3.0 + + # Work out the planet parameters + if a is None: + a = (G_grav * (m_star + m_planet) * period**2 / (4 * np.pi**2)) ** (1.0 / 3) + elif period is None: + period = 2 * np.pi * a ** (3 / 2) / (tt.sqrt(G_grav * (m_star + m_planet))) + + return a, period, rho_star * gcc_per_sun, r_star, m_star, m_planet diff --git a/thejoker/data.py b/thejoker/data.py index ea073bb1..5dc82ca8 100644 --- a/thejoker/data.py +++ b/thejoker/data.py @@ -388,15 +388,15 @@ def phase(self, P, t_ref=None): t_ref = self.t_ref return ((self.t - t_ref) / P) % 1.0 + @deprecated_renamed_argument("phase_fold", "phase_fold_period", "v1.3") def plot( self, ax=None, rv_unit=None, time_format="mjd", - phase_fold=None, + phase_fold_period=None, relative_to_t_ref=False, add_labels=True, - color_by=None, **kwargs, ): """ @@ -458,8 +458,8 @@ def plot( if relative_to_t_ref: t = t - t0 - if phase_fold: - t = (t / phase_fold.to(u.day).value) % 1 + if phase_fold_period is not None: + t = (t / phase_fold_period.to(u.day).value) % 1 if self._has_cov: # FIXME: this is a bit of a hack diff --git a/thejoker/samples.py b/thejoker/samples.py index 8ff66b0f..289db48f 100644 --- a/thejoker/samples.py +++ b/thejoker/samples.py @@ -103,6 +103,67 @@ def __init__( # used for speed-ups below self._cache = {} + @classmethod + def from_inference_data(cls, prior, idata, data, prune_divergences=True): + """ + Create a ``JokerSamples`` instance from an arviz object. + + Parameters + ---------- + prior : `thejoker.JokerPrior` + idata : `arviz.InferenceData` + data : `thejoker.RVData` + prune_divergences : bool (optional) + + """ + import thejoker.units as xu + from thejoker.thejoker import validate_prepare_data + + if hasattr(idata, "posterior"): + posterior = idata.posterior + + else: + posterior = idata + idata = None + + data, *_ = validate_prepare_data(data, prior.poly_trend, prior.n_offsets) + + samples = cls( + poly_trend=prior.poly_trend, + n_offsets=prior.n_offsets, + t_ref=data.t_ref, + ) + + names = prior.par_names + + for name in names: + if name in prior.pars: + par = prior.pars[name] + unit = getattr(par, xu.UNIT_ATTR_NAME) + samples[name] = posterior[name].to_numpy().ravel() * unit + else: + samples[name] = posterior[name].to_numpy().ravel() + + if hasattr(posterior, "logp"): + samples["ln_posterior"] = posterior.logp.to_numpy().ravel() + + for name in ["ln_likelihood", "ln_prior"]: + if hasattr(posterior, name): + samples[name] = getattr(posterior, name).to_numpy().ravel() + + if prune_divergences: + if idata is None: + msg = ( + "If you want to remove divergences, you must pass in the root level" + " inferencedata object (instead of, e.g., inferencedata.posterior" + ) + raise ValueError(msg) + + divergences = idata.sample_stats.diverging.to_numpy().ravel() + samples = samples[~divergences] + + return samples + def __getitem__(self, key): if isinstance(key, int): return self.__class__(samples=self.tbl[key]) diff --git a/thejoker/samples_helpers.py b/thejoker/samples_helpers.py index 94e32551..d408aa1d 100644 --- a/thejoker/samples_helpers.py +++ b/thejoker/samples_helpers.py @@ -270,102 +270,3 @@ def write_table_hdf5( current_size = len(output_group[name]) output_group[name].resize((current_size + len(table),)) output_group[name][current_size:] = table.as_array() - - -def inferencedata_to_samples(joker_prior, inferencedata, data, prune_divergences=True): - """ - Create a ``JokerSamples`` instance from an arviz object. - - Parameters - ---------- - joker_prior : `thejoker.JokerPrior` - inferencedata : `arviz.InferenceData` - data : `thejoker.RVData` - prune_divergences : bool (optional) - - """ - import thejoker.units as xu - from thejoker.samples import JokerSamples - from thejoker.thejoker import validate_prepare_data - - if hasattr(inferencedata, "posterior"): - posterior = inferencedata.posterior - - else: - posterior = inferencedata - inferencedata = None - - data, *_ = validate_prepare_data( - data, joker_prior.poly_trend, joker_prior.n_offsets - ) - - samples = JokerSamples( - poly_trend=joker_prior.poly_trend, - n_offsets=joker_prior.n_offsets, - t_ref=data.t_ref, - ) - - names = joker_prior.par_names - - for name in names: - if name in joker_prior.pars: - par = joker_prior.pars[name] - unit = getattr(par, xu.UNIT_ATTR_NAME) - samples[name] = posterior[name].values.ravel() * unit - else: - samples[name] = posterior[name].values.ravel() - - if hasattr(posterior, "logp"): - samples["ln_posterior"] = posterior.logp.values.ravel() - - for name in ["ln_likelihood", "ln_prior"]: - if hasattr(posterior, name): - samples[name] = getattr(posterior, name).values.ravel() - - if prune_divergences: - if inferencedata is None: - raise ValueError( - "If you want to remove divergences, you must pass in the root " - "level inferencedata object (instead of, e.g., inferencedata. " - "posterior" - ) - - divergences = inferencedata.sample_stats.diverging.values.ravel() - samples = samples[~divergences] - - return samples - - -def trace_to_samples(self, trace, data, names=None): - """ - Create a ``JokerSamples`` instance from a pymc trace object. - - Parameters - ---------- - trace : `~pymc.backends.base.MultiTrace` - """ - import pymc as pm - - import thejoker.units as xu - from thejoker.samples import JokerSamples - from thejoker.thejoker import validate_prepare_data - - df = pm.trace_to_dataframe(trace) - - data, *_ = validate_prepare_data(data, self.prior.poly_trend, self.prior.n_offsets) - - samples = JokerSamples( - poly_trend=self.prior.poly_trend, - n_offsets=self.prior.n_offsets, - t_ref=data.t_ref, - ) - - if names is None: - names = self.prior.par_names - - for name in names: - par = self.prior.pars[name] - unit = getattr(par, xu.UNIT_ATTR_NAME) - samples[name] = df[name].values * unit - - return samples diff --git a/thejoker/src/fast_likelihood.pyx b/thejoker/src/fast_likelihood.pyx index f3021e3a..2c1e9f47 100644 --- a/thejoker/src/fast_likelihood.pyx +++ b/thejoker/src/fast_likelihood.pyx @@ -205,13 +205,19 @@ cdef class CJokerHelper: # - validated to be Normal() in JokerPrior for i in range(self.n_offsets): name = prior.v0_offsets[i].name - dist = prior.v0_offsets[i].distribution - _unit = getattr(prior.model[name], xu.UNIT_ATTR_NAME) + # dist = prior.v0_offsets[i].distribution + dist = prior.model[name] + _unit = getattr(dist, xu.UNIT_ATTR_NAME) to_unit = self.internal_units[name] + # mean is par 0, stddev par 1 + pars = dist.owner.inputs[3:] + mu = (pars[0].eval() * _unit).to_value(to_unit) + std = (pars[1].eval() * _unit).to_value(to_unit) + # K, v0 = 2 - start at index 2 - self.mu[2+i] = (dist.mean.eval() * _unit).to_value(to_unit) - self.Lambda[2+i] = (dist.sd.eval() * _unit).to_value(to_unit) ** 2 + self.mu[2+i] = mu + self.Lambda[2+i] = std ** 2 # --------------------------------------------------------------------- # TODO: This is a bit of a hack: diff --git a/thejoker/thejoker.py b/thejoker/thejoker.py index 3c082e2b..b0ec2433 100644 --- a/thejoker/thejoker.py +++ b/thejoker/thejoker.py @@ -133,10 +133,9 @@ def marginal_ln_likelihood( ) return marginal_ln_likelihood_inmem(joker_helper, prior_samples) - else: - return marginal_ln_likelihood_helper( - joker_helper, prior_samples, pool=self.pool, n_batches=n_batches - ) + return marginal_ln_likelihood_helper( + joker_helper, prior_samples, pool=self.pool, n_batches=n_batches + ) def rejection_sample( self, @@ -427,7 +426,7 @@ def setup_mcmc(self, data, joker_samples, model=None, custom_func=None): else: MAP_sample = joker_samples - mcmc_init = dict() + mcmc_init = {} for name in self.prior.par_names: unit = getattr(self.prior.pars[name], xu.UNIT_ATTR_NAME) mcmc_init[name] = MAP_sample[name].to_value(unit) From c96873360f4dcc34713688cc31ed4d50b0bc29d8 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Tue, 5 Mar 2024 18:34:12 -0500 Subject: [PATCH 26/50] update examples in docs and make sure they run --- docs/conf.py | 272 ++++-------------- docs/examples/1-Getting-started.ipynb | 91 +++--- docs/examples/2-Customize-prior.ipynb | 75 ++--- .../3-Polynomial-velocity-trend.ipynb | 52 ++-- docs/examples/4-Continue-sampling-mcmc.ipynb | 49 ++-- docs/examples/5-Calibration-offsets.ipynb | 54 ++-- docs/examples/Strader-circular-only.ipynb | 125 ++++---- docs/examples/Thompson-black-hole.ipynb | 172 ++++++----- docs/examples/notebook_setup.py | 3 - 9 files changed, 355 insertions(+), 538 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index 2851ae96..0955fe0f 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -1,224 +1,62 @@ -# Licensed under a 3-clause BSD style license - see LICENSE.rst -# -# Astropy documentation build configuration file. -# -# This file is execfile()d with the current directory set to its containing dir. -# -# Note that not all possible configuration values are present in this file. -# -# All configuration values have a default. Some values are defined in -# the global Astropy configuration which is loaded here before anything else. -# See astropy.sphinx.conf for which values are set there. - -# If extensions (or modules to document with autodoc) are in another directory, -# add these directories to sys.path here. If the directory is relative to the -# documentation root, use os.path.abspath to make it absolute, like shown here. -# sys.path.insert(0, os.path.abspath('..')) -# IMPORTANT: the above commented section was generated by sphinx-quickstart, but -# is *NOT* appropriate for astropy or Astropy affiliated packages. It is left -# commented out with this explanation to make it clear why this should not be -# done. If the sys.path entry above is added, when the astropy.sphinx.conf -# import occurs, it will import the *source* version of astropy instead of the -# version installed (if invoked as "make html" or directly with sphinx), or the -# version in the build directory (if "python setup.py build_sphinx" is used). -# Thus, any C-extensions that are needed to build the documentation will *not* -# be accessible, and the documentation will not build correctly. - -import datetime -import os -import sys -from importlib import import_module - -try: - from sphinx_astropy.conf.v1 import * # noqa -except ImportError: - print( - "ERROR: the documentation requires the sphinx-astropy package to be installed" - ) - sys.exit(1) - -# Get configuration information from setup.cfg -from configparser import ConfigParser - -conf = ConfigParser() - -conf.read([os.path.join(os.path.dirname(__file__), "..", "setup.cfg")]) -setup_cfg = dict(conf.items("metadata")) - -# -- General configuration ---------------------------------------------------- - -# By default, highlight as Python 3. -highlight_language = "python3" - -# If your documentation needs a minimal Sphinx version, state it here. -# needs_sphinx = '1.2' - -# To perform a Sphinx version check that needs to be more specific than -# major.minor, call `check_sphinx_version("x.y.z")` here. -# check_sphinx_version("1.2.1") - -# List of patterns, relative to source directory, that match files and -# directories to ignore when looking for source files. -exclude_patterns.append("_templates") -exclude_patterns.append("**.ipynb_checkpoints") - -# This is added to the end of RST files - a good place to put substitutions to -# be used globally. -rst_epilog += """ -.. |thejoker| replace:: *The Joker* -""" - -# -- Project information ------------------------------------------------------ - -# This does not *have* to match the package name, but typically does -project = setup_cfg["name"] -author = setup_cfg["author"] -copyright = "{0}, {1}".format(datetime.datetime.now().year, setup_cfg["author"]) - -# The version info for the project you're documenting, acts as replacement for -# |version| and |release|, also used in various other places throughout the -# built documents. - -import_module(setup_cfg["name"]) -package = sys.modules[setup_cfg["name"]] - -# The short X.Y version. -version = package.__version__.split("-", 1)[0] -# The full version, including alpha/beta/rc tags. -release = package.__version__ - - -# -- Options for HTML output -------------------------------------------------- - -# A NOTE ON HTML THEMES -# The global astropy configuration uses a custom theme, 'bootstrap-astropy', -# which is installed along with astropy. A different theme can be used or -# the options for this theme can be modified by overriding some of the -# variables set in the global configuration. The variables set in the -# global configuration are listed below, commented out. - - -# Add any paths that contain custom themes here, relative to this directory. -# To use a different custom theme, add the directory containing the theme. -# html_theme_path = [] - -# The theme to use for HTML and HTML Help pages. See the documentation for -# a list of builtin themes. To override the custom theme, set this to the -# name of a builtin theme or the name of a custom theme in html_theme_path. -# html_theme = None - - -# Please update these texts to match the name of your package. -html_theme_options = { - "logotext1": "The", # white, semi-bold - "logotext2": "Joker", # orange, light - "logotext3": ":docs", # white, light -} - -# Custom sidebar templates, maps document names to template names. -# html_sidebars = {} - -# The name of an image file (relative to this directory) to place at the top -# of the sidebar. -# html_logo = '' - -# The name of an image file (within the static path) to use as favicon of the -# docs. This file should be a Windows icon file (.ico) being 16x16 or 32x32 -# pixels large. -path = os.path.abspath(os.path.join(os.path.dirname(__file__), "_static")) -html_favicon = os.path.join(path, "icon.ico") - -# If not '', a 'Last updated on:' timestamp is inserted at every page bottom, -# using the given strftime format. -# html_last_updated_fmt = '' - -# The name for this set of Sphinx documents. If None, it defaults to -# " v documentation". -html_title = f"{project} v{release}" +import importlib.metadata + +project = "thejoker" +copyright = "2024, Adrian Price-Whelan" +author = "Adrian Price-Whelan" +version = release = importlib.metadata.version("thejoker") + +extensions = [ + "sphinx.ext.autodoc", + "sphinx.ext.intersphinx", + "sphinx.ext.mathjax", + "sphinx.ext.napoleon", + "sphinx_autodoc_typehints", + "sphinx_copybutton", + "sphinx_automodapi.automodapi", + "sphinx_automodapi.smart_resolver", +] -# Output file base name for HTML help builder. -htmlhelp_basename = project + "doc" +source_suffix = [".rst", ".md"] +exclude_patterns = [ + "_build", + "**.ipynb_checkpoints", + "Thumbs.db", + ".DS_Store", + ".env", + ".venv", +] +html_theme = "pydata_sphinx_theme" html_static_path = ["_static"] -html_style = "thejoker.css" - -# -- Options for LaTeX output ------------------------------------------------- - -# Grouping the document tree into LaTeX files. List of tuples -# (source start file, target name, title, author, documentclass [howto/manual]). -latex_documents = [ - ("index", project + ".tex", project + " Documentation", author, "manual") +html_css_files = [ + "custom.css", ] +intersphinx_mapping = { + "python": ("https://docs.python.org/3", None), + "pandas": ("https://pandas.pydata.org/pandas-docs/stable/", None), + "pymc": ("https://docs.pymc.io/", None), + "h5py": ("http://docs.h5py.org/en/latest/", None), + "twobody": ("https://twobody.readthedocs.io/en/latest/", None), + "schwimmbad": ("https://schwimmbad.readthedocs.io/en/latest/", None), + "numpy": ( + "https://numpy.org/doc/stable/", + (None, "http://data.astropy.org/intersphinx/numpy.inv"), + ), + "scipy": ( + "https://docs.scipy.org/doc/scipy/", + (None, "http://data.astropy.org/intersphinx/scipy.inv"), + ), + "matplotlib": ( + "https://matplotlib.org/stable/", + (None, "http://data.astropy.org/intersphinx/matplotlib.inv"), + ), + "astropy": ("https://docs.astropy.org/en/stable/", None), +} -# -- Options for manual page output ------------------------------------------- - -# One entry per manual page. List of tuples -# (source start file, name, description, authors, manual section). -man_pages = [("index", project.lower(), project + " Documentation", [author], 1)] - - -# -- Options for the edit_on_github extension --------------------------------- - -if eval(setup_cfg.get("edit_on_github")): - extensions += ["sphinx_astropy.ext.edit_on_github"] - - versionmod = __import__(setup_cfg["package_name"] + ".version") - edit_on_github_project = setup_cfg["github_project"] - if versionmod.version.release: - edit_on_github_branch = "v" + versionmod.version.version - else: - edit_on_github_branch = "main" - - edit_on_github_source_root = "" - edit_on_github_doc_root = "docs" - -# -- Resolving issue number to links in changelog ----------------------------- -github_issues_url = "https://github.com/{0}/issues/".format(setup_cfg["github_project"]) - -intersphinx_mapping["h5py"] = ("http://docs.h5py.org/en/latest/", None) -intersphinx_mapping["pymc"] = ("https://docs.pymc.io/", None) -intersphinx_mapping["twobody"] = ("https://twobody.readthedocs.io/en/latest/", None) -intersphinx_mapping["scwhimmbad"] = ( - "https://schwimmbad.readthedocs.io/en/latest/", - None, -) - -# see if we're running on CI: -ON_CI = os.environ.get("CI", False) -PR = os.environ.get("CIRCLE_PULL_REQUEST", False) - -# Use astropy plot style -plot_rcparams = dict() -if not ON_CI: - plot_rcparams["text.usetex"] = True -plot_rcparams["savefig.facecolor"] = "none" -plot_rcparams["savefig.bbox"] = "tight" -plot_apply_rcparams = True -plot_formats = [("png", 512)] - - -# nbsphinx config: -exclude_patterns.append("make-data.*") -exclude_patterns.append("*/make-data.*") - -extensions += ["nbsphinx"] -extensions += ["IPython.sphinxext.ipython_console_highlighting"] - -extensions += ["sphinx.ext.mathjax"] -mathjax_path = "https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_HTMLorMML" - -# nbsphinx_execute_arguments = [ -# "--InlineBackend.rc={'figure.dpi': 250}", -# ] - -nbsphinx_timeout = 600 - -if PR: - nbsphinx_execute = "never" - -if ON_CI: - nbsphinx_kernel_name = "python3" +nitpick_ignore = [ + ("py:class", "_io.StringIO"), + ("py:class", "_io.BytesIO"), +] -else: - nbsphinx_kernel_name = os.environ.get("NBSPHINX_KERNEL", "python3") +always_document_param_types = True diff --git a/docs/examples/1-Getting-started.ipynb b/docs/examples/1-Getting-started.ipynb index ee49ee1b..4e1c23c5 100644 --- a/docs/examples/1-Getting-started.ipynb +++ b/docs/examples/1-Getting-started.ipynb @@ -42,9 +42,9 @@ "from astropy.visualization.units import quantity_support\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", - "%matplotlib inline\n", + "import thejoker as tj\n", "\n", - "import thejoker as tj" + "%matplotlib inline" ] }, { @@ -79,7 +79,7 @@ "metadata": {}, "outputs": [], "source": [ - "data_tbl = at.QTable.read('data.ecsv')\n", + "data_tbl = at.QTable.read(\"data.ecsv\")\n", "data_tbl[:2]" ] }, @@ -121,8 +121,8 @@ "metadata": {}, "outputs": [], "source": [ - "t = Time(sub_tbl['bjd'], format='jd', scale='tcb') \n", - "data = tj.RVData(t=t, rv=sub_tbl['rv'], rv_err=sub_tbl['rv_err'])" + "t = Time(sub_tbl[\"bjd\"], format=\"jd\", scale=\"tcb\")\n", + "data = tj.RVData(t=t, rv=sub_tbl[\"rv\"], rv_err=sub_tbl[\"rv_err\"])" ] }, { @@ -196,9 +196,11 @@ "outputs": [], "source": [ "prior = tj.JokerPrior.default(\n", - " P_min=2*u.day, P_max=1e3*u.day,\n", - " sigma_K0=30*u.km/u.s,\n", - " sigma_v=100*u.km/u.s)" + " P_min=2 * u.day,\n", + " P_max=1e3 * u.day,\n", + " sigma_K0=30 * u.km / u.s,\n", + " sigma_v=100 * u.km / u.s,\n", + ")" ] }, { @@ -214,8 +216,7 @@ "metadata": {}, "outputs": [], "source": [ - "prior_samples = prior.sample(size=250_000,\n", - " rng=rnd)\n", + "prior_samples = prior.sample(size=250_000, rng=rnd)\n", "prior_samples" ] }, @@ -232,7 +233,7 @@ "metadata": {}, "outputs": [], "source": [ - "prior_samples['P']" + "prior_samples[\"P\"]" ] }, { @@ -241,7 +242,7 @@ "metadata": {}, "outputs": [], "source": [ - "prior_samples['e']" + "prior_samples[\"e\"]" ] }, { @@ -292,8 +293,7 @@ "outputs": [], "source": [ "joker = tj.TheJoker(prior, rng=rnd)\n", - "joker_samples = joker.rejection_sample(data, prior_samples, \n", - " max_posterior_samples=256)" + "joker_samples = joker.rejection_sample(data, prior_samples, max_posterior_samples=256)" ] }, { @@ -309,8 +309,9 @@ "metadata": {}, "outputs": [], "source": [ - "joker_samples = joker.rejection_sample(data, \"prior_samples.hdf5\", \n", - " max_posterior_samples=256)" + "joker_samples = joker.rejection_sample(\n", + " data, \"prior_samples.hdf5\", max_posterior_samples=256\n", + ")" ] }, { @@ -363,11 +364,15 @@ "outputs": [], "source": [ "fig, ax = plt.subplots(1, 1, figsize=(8, 4))\n", - "_ = tj.plot_rv_curves(joker_samples, data=data, \n", - " plot_kwargs=dict(color='tab:blue'),\n", - " data_plot_kwargs=dict(color='tab:red'),\n", - " relative_to_t_ref=True, ax=ax)\n", - "ax.set_xlabel(f'BMJD$ - {data.t.tcb.mjd.min():.3f}$')" + "_ = tj.plot_rv_curves(\n", + " joker_samples,\n", + " data=data,\n", + " plot_kwargs=dict(color=\"tab:blue\"),\n", + " data_plot_kwargs=dict(color=\"tab:red\"),\n", + " relative_to_t_ref=True,\n", + " ax=ax,\n", + ")\n", + "ax.set_xlabel(f\"BMJD$ - {data.t.tcb.mjd.min():.3f}$\")" ] }, { @@ -386,17 +391,14 @@ "fig, ax = plt.subplots(1, 1, figsize=(8, 5))\n", "\n", "with quantity_support():\n", - " ax.scatter(joker_samples['P'], \n", - " joker_samples['e'],\n", - " s=20, lw=0, alpha=0.5)\n", - " \n", - "ax.set_xscale('log')\n", - "ax.set_xlim(prior.pars['P'].distribution.a,\n", - " prior.pars['P'].distribution.b)\n", + " ax.scatter(joker_samples[\"P\"], joker_samples[\"e\"], s=20, lw=0, alpha=0.5)\n", + "\n", + "ax.set_xscale(\"log\")\n", + "ax.set_xlim(prior.pars[\"P\"].distribution.a, prior.pars[\"P\"].distribution.b)\n", "ax.set_ylim(0, 1)\n", "\n", - "ax.set_xlabel('$P$ [day]')\n", - "ax.set_ylabel('$e$')" + "ax.set_xlabel(\"$P$ [day]\")\n", + "ax.set_ylabel(\"$e$\")" ] }, { @@ -413,7 +415,8 @@ "outputs": [], "source": [ "import pickle\n", - "with open('true-orbit.pkl', 'rb') as f:\n", + "\n", + "with open(\"true-orbit.pkl\", \"rb\") as f:\n", " truth = pickle.load(f)" ] }, @@ -426,22 +429,20 @@ "fig, ax = plt.subplots(1, 1, figsize=(8, 5))\n", "\n", "with quantity_support():\n", - " ax.scatter(joker_samples['P'], \n", - " joker_samples['e'],\n", - " s=20, lw=0, alpha=0.5)\n", - " \n", - " ax.axvline(truth['P'], zorder=-1, color='tab:green')\n", - " ax.axhline(truth['e'], zorder=-1, color='tab:green')\n", - " ax.text(truth['P'], 0.95, 'truth', fontsize=20, \n", - " va='top', ha='left', color='tab:green')\n", - " \n", - "ax.set_xscale('log')\n", - "ax.set_xlim(prior.pars['P'].distribution.a,\n", - " prior.pars['P'].distribution.b)\n", + " ax.scatter(joker_samples[\"P\"], joker_samples[\"e\"], s=20, lw=0, alpha=0.5)\n", + "\n", + " ax.axvline(truth[\"P\"], zorder=-1, color=\"tab:green\")\n", + " ax.axhline(truth[\"e\"], zorder=-1, color=\"tab:green\")\n", + " ax.text(\n", + " truth[\"P\"], 0.95, \"truth\", fontsize=20, va=\"top\", ha=\"left\", color=\"tab:green\"\n", + " )\n", + "\n", + "ax.set_xscale(\"log\")\n", + "ax.set_xlim(prior.pars[\"P\"].distribution.a, prior.pars[\"P\"].distribution.b)\n", "ax.set_ylim(0, 1)\n", "\n", - "ax.set_xlabel('$P$ [day]')\n", - "ax.set_ylabel('$e$')" + "ax.set_xlabel(\"$P$ [day]\")\n", + "ax.set_ylabel(\"$e$\")" ] }, { diff --git a/docs/examples/2-Customize-prior.ipynb b/docs/examples/2-Customize-prior.ipynb index 09014a3b..74a92b49 100644 --- a/docs/examples/2-Customize-prior.ipynb +++ b/docs/examples/2-Customize-prior.ipynb @@ -39,11 +39,11 @@ "from astropy.visualization.units import quantity_support\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", - "%matplotlib inline\n", - "import pymc3 as pm\n", - "import exoplanet.units as xu\n", + "import pymc as pm\n", + "import thejoker.units as xu\n", + "import thejoker as tj\n", "\n", - "import thejoker as tj" + "%matplotlib inline" ] }, { @@ -72,13 +72,11 @@ "outputs": [], "source": [ "with pm.Model() as model:\n", - " P = xu.with_unit(pm.Normal('P', 50., 1),\n", - " u.day)\n", + " P = xu.with_unit(pm.Normal(\"P\", 50.0, 1), u.day)\n", "\n", " prior = tj.JokerPrior.default(\n", - " sigma_K0=30*u.km/u.s,\n", - " sigma_v=100*u.km/u.s,\n", - " pars={'P': P})\n", + " sigma_K0=30 * u.km / u.s, sigma_v=100 * u.km / u.s, pars={\"P\": P}\n", + " )\n", "\n", "samples1 = prior.sample(size=100_000, rng=rnd)" ] @@ -96,8 +94,8 @@ "metadata": {}, "outputs": [], "source": [ - "plt.hist(samples1['P'].to_value(u.day), bins=64);\n", - "plt.xlabel('$P$ [day]')" + "plt.hist(samples1[\"P\"].to_value(u.day), bins=64)\n", + "plt.xlabel(\"$P$ [day]\")" ] }, { @@ -116,14 +114,10 @@ "outputs": [], "source": [ "with pm.Model() as model:\n", - " P = xu.with_unit(pm.Normal('P', 50., 1),\n", - " u.day)\n", - " K = xu.with_unit(pm.Normal('K', 0., 15),\n", - " u.km/u.s)\n", + " P = xu.with_unit(pm.Normal(\"P\", 50.0, 1), u.day)\n", + " K = xu.with_unit(pm.Normal(\"K\", 0.0, 15), u.km / u.s)\n", "\n", - " prior = tj.JokerPrior.default(\n", - " sigma_v=100*u.km/u.s,\n", - " pars={'P': P, 'K': K})\n", + " prior = tj.JokerPrior.default(sigma_v=100 * u.km / u.s, pars={\"P\": P, \"K\": K})\n", "\n", "samples2 = prior.sample(size=100_000, rng=rnd)\n", "samples2" @@ -142,8 +136,7 @@ "metadata": {}, "outputs": [], "source": [ - "samples3 = prior.sample(size=100_000, generate_linear=True,\n", - " rng=rnd)\n", + "samples3 = prior.sample(size=100_000, generate_linear=True, rng=rnd)\n", "samples3" ] }, @@ -170,13 +163,14 @@ "outputs": [], "source": [ "default_prior = tj.JokerPrior.default(\n", - " P_min=1e1*u.day,\n", - " P_max=1e3*u.day,\n", - " sigma_K0=30*u.km/u.s,\n", - " sigma_v=75*u.km/u.s)\n", - "default_samples = default_prior.sample(size=20, generate_linear=True,\n", - " rng=rnd,\n", - " t_ref=Time('J2000')) # set arbitrary time zero-point" + " P_min=1e1 * u.day,\n", + " P_max=1e3 * u.day,\n", + " sigma_K0=30 * u.km / u.s,\n", + " sigma_v=75 * u.km / u.s,\n", + ")\n", + "default_samples = default_prior.sample(\n", + " size=20, generate_linear=True, rng=rnd, t_ref=Time(\"J2000\")\n", + ") # set arbitrary time zero-point" ] }, { @@ -186,18 +180,14 @@ "outputs": [], "source": [ "with pm.Model() as model:\n", - " K = xu.with_unit(pm.Normal('K', 0., 30),\n", - " u.km/u.s)\n", + " K = xu.with_unit(pm.Normal(\"K\", 0.0, 30), u.km / u.s)\n", " custom_prior = tj.JokerPrior.default(\n", - " P_min=1e1*u.day,\n", - " P_max=1e3*u.day,\n", - " sigma_v=75*u.km/u.s,\n", - " pars={'K': K})\n", + " P_min=1e1 * u.day, P_max=1e3 * u.day, sigma_v=75 * u.km / u.s, pars={\"K\": K}\n", + " )\n", "\n", - "custom_samples = custom_prior.sample(size=len(default_samples),\n", - " generate_linear=True,\n", - " rng=rnd,\n", - " t_ref=Time('J2000')) # set arbitrary time zero-point" + "custom_samples = custom_prior.sample(\n", + " size=len(default_samples), generate_linear=True, rng=rnd, t_ref=Time(\"J2000\")\n", + ") # set arbitrary time zero-point" ] }, { @@ -207,14 +197,11 @@ "outputs": [], "source": [ "now_mjd = Time.now().mjd\n", - "t_grid = Time(np.linspace(now_mjd - 1000, now_mjd + 1000, 16384),\n", - " format='mjd')\n", + "t_grid = Time(np.linspace(now_mjd - 1000, now_mjd + 1000, 16384), format=\"mjd\")\n", "\n", "fig, axes = plt.subplots(2, 1, sharex=True, sharey=True, figsize=(8, 8))\n", - "_ = tj.plot_rv_curves(default_samples, t_grid=t_grid,\n", - " ax=axes[0], add_labels=False)\n", - "_ = tj.plot_rv_curves(custom_samples, t_grid=t_grid,\n", - " ax=axes[1])\n", + "_ = tj.plot_rv_curves(default_samples, t_grid=t_grid, ax=axes[0], add_labels=False)\n", + "_ = tj.plot_rv_curves(custom_samples, t_grid=t_grid, ax=axes[1])\n", "axes[0].set_ylim(-200, 200)\n", "fig.tight_layout()" ] @@ -236,7 +223,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.3" + "version": "3.11.8" }, "toc": { "base_numbering": 1, diff --git a/docs/examples/3-Polynomial-velocity-trend.ipynb b/docs/examples/3-Polynomial-velocity-trend.ipynb index f86e1cf0..04942af4 100644 --- a/docs/examples/3-Polynomial-velocity-trend.ipynb +++ b/docs/examples/3-Polynomial-velocity-trend.ipynb @@ -35,12 +35,10 @@ "source": [ "import astropy.table as at\n", "import astropy.units as u\n", - "from astropy.visualization.units import quantity_support\n", - "import matplotlib.pyplot as plt\n", "import numpy as np\n", - "%matplotlib inline\n", + "import thejoker as tj\n", "\n", - "import thejoker as tj" + "%matplotlib inline" ] }, { @@ -50,7 +48,7 @@ "outputs": [], "source": [ "# set up a random number generator to ensure reproducibility\n", - "rnd = np.random.default_rng(seed=42)" + "rnd = np.random.default_rng(seed=123)" ] }, { @@ -66,7 +64,7 @@ "metadata": {}, "outputs": [], "source": [ - "data = tj.RVData.guess_from_table(at.QTable.read('data-triple.ecsv'))\n", + "data = tj.RVData.guess_from_table(at.QTable.read(\"data-triple.ecsv\"))\n", "data = data[rnd.choice(len(data), size=16, replace=False)] # downsample data" ] }, @@ -93,9 +91,11 @@ "outputs": [], "source": [ "prior = tj.JokerPrior.default(\n", - " P_min=2*u.day, P_max=1e3*u.day,\n", - " sigma_K0=30*u.km/u.s,\n", - " sigma_v=100*u.km/u.s)\n", + " P_min=2 * u.day,\n", + " P_max=1e3 * u.day,\n", + " sigma_K0=30 * u.km / u.s,\n", + " sigma_v=100 * u.km / u.s,\n", + ")\n", "prior_samples = prior.sample(size=250_000, rng=rnd)" ] }, @@ -113,8 +113,7 @@ "outputs": [], "source": [ "joker = tj.TheJoker(prior, rng=rnd)\n", - "samples = joker.rejection_sample(data, prior_samples,\n", - " max_posterior_samples=128)\n", + "samples = joker.rejection_sample(data, prior_samples, max_posterior_samples=128)\n", "samples" ] }, @@ -141,12 +140,12 @@ "outputs": [], "source": [ "prior_trend = tj.JokerPrior.default(\n", - " P_min=2*u.day, P_max=1e3*u.day,\n", - " sigma_K0=30*u.km/u.s,\n", - " sigma_v=[100*u.km/u.s,\n", - " 0.5*u.km/u.s/u.day,\n", - " 1e-2*u.km/u.s/u.day**2],\n", - " poly_trend=3)" + " P_min=2 * u.day,\n", + " P_max=1e3 * u.day,\n", + " sigma_K0=30 * u.km / u.s,\n", + " sigma_v=[100 * u.km / u.s, 0.5 * u.km / u.s / u.day, 1e-2 * u.km / u.s / u.day**2],\n", + " poly_trend=3,\n", + ")" ] }, { @@ -178,11 +177,11 @@ "metadata": {}, "outputs": [], "source": [ - "prior_samples_trend = prior_trend.sample(size=250_000,\n", - " rng=rnd)\n", + "prior_samples_trend = prior_trend.sample(size=250_000, rng=rnd)\n", "joker_trend = tj.TheJoker(prior_trend, rng=rnd)\n", - "samples_trend = joker_trend.rejection_sample(data, prior_samples_trend,\n", - " max_posterior_samples=128)\n", + "samples_trend = joker_trend.rejection_sample(\n", + " data, prior_samples_trend, max_posterior_samples=128\n", + ")\n", "samples_trend" ] }, @@ -211,7 +210,8 @@ "outputs": [], "source": [ "import pickle\n", - "with open('true-orbit-triple.pkl', 'rb') as f:\n", + "\n", + "with open(\"true-orbit-triple.pkl\", \"rb\") as f:\n", " truth = pickle.load(f)" ] }, @@ -228,7 +228,7 @@ "metadata": {}, "outputs": [], "source": [ - "truth['P'], truth['e'], truth['K']" + "truth[\"P\"], truth[\"e\"], truth[\"K\"]" ] }, { @@ -244,7 +244,7 @@ "metadata": {}, "outputs": [], "source": [ - "samples['P'], samples['e'], samples['K']" + "samples[\"P\"], samples[\"e\"], samples[\"K\"]" ] }, { @@ -260,7 +260,7 @@ "metadata": {}, "outputs": [], "source": [ - "samples_trend.mean()['P'], samples_trend.mean()['e'], samples_trend.mean()['K']" + "samples_trend.mean()[\"P\"], samples_trend.mean()[\"e\"], samples_trend.mean()[\"K\"]" ] }, { @@ -287,7 +287,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.7.5" + "version": "3.11.8" }, "toc": { "base_numbering": 1, diff --git a/docs/examples/4-Continue-sampling-mcmc.ipynb b/docs/examples/4-Continue-sampling-mcmc.ipynb index 3aae6b32..ce004c56 100644 --- a/docs/examples/4-Continue-sampling-mcmc.ipynb +++ b/docs/examples/4-Continue-sampling-mcmc.ipynb @@ -40,18 +40,14 @@ "source": [ "import astropy.coordinates as coord\n", "import astropy.table as at\n", - "from astropy.time import Time\n", "import astropy.units as u\n", - "import matplotlib.pyplot as plt\n", "import numpy as np\n", - "%matplotlib inline\n", "import corner\n", - "import pymc3 as pm\n", - "import pymc3_ext as pmx\n", - "import exoplanet as xo\n", + "import pymc as pm\n", "import arviz as az\n", + "import thejoker as tj\n", "\n", - "import thejoker as tj" + "%matplotlib inline" ] }, { @@ -61,7 +57,7 @@ "outputs": [], "source": [ "# set up a random number generator to ensure reproducibility\n", - "rnd = np.random.default_rng(seed=42)" + "rnd = np.random.default_rng(seed=8675309)" ] }, { @@ -77,9 +73,9 @@ "metadata": {}, "outputs": [], "source": [ - "data_tbl = at.QTable.read('data.ecsv')\n", + "data_tbl = at.QTable.read(\"data.ecsv\")\n", "sub_tbl = data_tbl[rnd.choice(len(data_tbl), size=18, replace=False)] # downsample data\n", - "data = tj.RVData.guess_from_table(sub_tbl, t_ref=data_tbl.meta['t_ref'])" + "data = tj.RVData.guess_from_table(sub_tbl, t_ref=data_tbl.meta[\"t_ref\"])" ] }, { @@ -105,9 +101,11 @@ "outputs": [], "source": [ "prior = tj.JokerPrior.default(\n", - " P_min=2*u.day, P_max=1e3*u.day,\n", - " sigma_K0=30*u.km/u.s,\n", - " sigma_v=100*u.km/u.s)" + " P_min=2 * u.day,\n", + " P_max=1e3 * u.day,\n", + " sigma_K0=30 * u.km / u.s,\n", + " sigma_v=100 * u.km / u.s,\n", + ")" ] }, { @@ -123,8 +121,7 @@ "metadata": {}, "outputs": [], "source": [ - "prior_samples = prior.sample(size=250_000,\n", - " rng=rnd)" + "prior_samples = prior.sample(size=250_000, rng=rnd)" ] }, { @@ -134,8 +131,7 @@ "outputs": [], "source": [ "joker = tj.TheJoker(prior, rng=rnd)\n", - "joker_samples = joker.rejection_sample(data, prior_samples,\n", - " max_posterior_samples=256)\n", + "joker_samples = joker.rejection_sample(data, prior_samples, max_posterior_samples=256)\n", "joker_samples" ] }, @@ -173,9 +169,7 @@ "with prior.model:\n", " mcmc_init = joker.setup_mcmc(data, joker_samples)\n", "\n", - " trace = pmx.sample(tune=500, draws=500,\n", - " start=mcmc_init,\n", - " cores=1, chains=2)" + " trace = pm.sample(tune=500, draws=500, start=mcmc_init, cores=1, chains=2)" ] }, { @@ -207,7 +201,7 @@ "metadata": {}, "outputs": [], "source": [ - "mcmc_samples = joker.trace_to_samples(trace, data)\n", + "mcmc_samples = tj.JokerSamples.from_inference_data(prior, trace, data)\n", "mcmc_samples.wrap_K()\n", "mcmc_samples" ] @@ -226,15 +220,16 @@ "outputs": [], "source": [ "import pickle\n", - "with open('true-orbit.pkl', 'rb') as f:\n", + "\n", + "with open(\"true-orbit.pkl\", \"rb\") as f:\n", " truth = pickle.load(f)\n", "\n", "# make sure the angles are wrapped the same way\n", - "if np.median(mcmc_samples['omega']) < 0:\n", - " truth['omega'] = coord.Angle(truth['omega']).wrap_at(np.pi*u.radian)\n", + "if np.median(mcmc_samples[\"omega\"]) < 0:\n", + " truth[\"omega\"] = coord.Angle(truth[\"omega\"]).wrap_at(np.pi * u.radian)\n", "\n", - "if np.median(mcmc_samples['M0']) < 0:\n", - " truth['M0'] = coord.Angle(truth['M0']).wrap_at(np.pi*u.radian)" + "if np.median(mcmc_samples[\"M0\"]) < 0:\n", + " truth[\"M0\"] = coord.Angle(truth[\"M0\"]).wrap_at(np.pi * u.radian)" ] }, { @@ -277,7 +272,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.3" + "version": "3.11.8" }, "toc": { "base_numbering": 1, diff --git a/docs/examples/5-Calibration-offsets.ipynb b/docs/examples/5-Calibration-offsets.ipynb index 245286b2..7e5d4183 100644 --- a/docs/examples/5-Calibration-offsets.ipynb +++ b/docs/examples/5-Calibration-offsets.ipynb @@ -35,18 +35,15 @@ "source": [ "import astropy.table as at\n", "import astropy.units as u\n", - "from astropy.visualization.units import quantity_support\n", - "import matplotlib.pyplot as plt\n", "import numpy as np\n", - "%matplotlib inline\n", "import corner\n", - "import pymc3 as pm\n", - "import pymc3_ext as pmx\n", - "import exoplanet as xo\n", - "import exoplanet.units as xu\n", + "import pymc as pm\n", + "import thejoker.units as xu\n", "import arviz as az\n", "\n", - "import thejoker as tj" + "import thejoker as tj\n", + "\n", + "%matplotlib inline" ] }, { @@ -73,9 +70,9 @@ "outputs": [], "source": [ "data = []\n", - "for filename in ['data-survey1.ecsv', 'data-survey2.ecsv']:\n", + "for filename in [\"data-survey1.ecsv\", \"data-survey2.ecsv\"]:\n", " tbl = at.QTable.read(filename)\n", - " _data = tj.RVData.guess_from_table(tbl, t_ref=tbl.meta['t_ref'])\n", + " _data = tj.RVData.guess_from_table(tbl, t_ref=tbl.meta[\"t_ref\"])\n", " data.append(_data)" ] }, @@ -92,7 +89,7 @@ "metadata": {}, "outputs": [], "source": [ - "for d, color in zip(data, ['tab:blue', 'tab:red']):\n", + "for d, color in zip(data, [\"tab:blue\", \"tab:red\"]):\n", " _ = d.plot(color=color)" ] }, @@ -110,14 +107,15 @@ "outputs": [], "source": [ "with pm.Model() as model:\n", - " dv0_1 = xu.with_unit(pm.Normal('dv0_1', 0, 10),\n", - " u.km/u.s)\n", + " dv0_1 = xu.with_unit(pm.Normal(\"dv0_1\", 0, 10), u.km / u.s)\n", "\n", " prior = tj.JokerPrior.default(\n", - " P_min=2*u.day, P_max=256*u.day,\n", - " sigma_K0=30*u.km/u.s,\n", - " sigma_v=100*u.km/u.s,\n", - " v0_offsets=[dv0_1])" + " P_min=2 * u.day,\n", + " P_max=256 * u.day,\n", + " sigma_K0=30 * u.km / u.s,\n", + " sigma_v=100 * u.km / u.s,\n", + " v0_offsets=[dv0_1],\n", + " )" ] }, { @@ -133,8 +131,7 @@ "metadata": {}, "outputs": [], "source": [ - "prior_samples = prior.sample(size=1_000_000,\n", - " rng=rnd)" + "prior_samples = prior.sample(size=1_000_000, rng=rnd)" ] }, { @@ -144,8 +141,7 @@ "outputs": [], "source": [ "joker = tj.TheJoker(prior, rng=rnd)\n", - "joker_samples = joker.rejection_sample(data, prior_samples,\n", - " max_posterior_samples=128)\n", + "joker_samples = joker.rejection_sample(data, prior_samples, max_posterior_samples=128)\n", "joker_samples" ] }, @@ -180,8 +176,7 @@ "metadata": {}, "outputs": [], "source": [ - "_ = tj.plot_rv_curves(joker_samples, data=data,\n", - " apply_mean_v0_offset=False)" + "_ = tj.plot_rv_curves(joker_samples, data=data, apply_mean_v0_offset=False)" ] }, { @@ -200,10 +195,7 @@ "with prior.model:\n", " mcmc_init = joker.setup_mcmc(data, joker_samples)\n", "\n", - " trace = pmx.sample(\n", - " tune=500, draws=500,\n", - " start=mcmc_init,\n", - " cores=1, chains=2)" + " trace = pm.sample(tune=500, draws=500, start=mcmc_init, cores=1, chains=2)" ] }, { @@ -230,8 +222,8 @@ "metadata": {}, "outputs": [], "source": [ - "mcmc_samples = joker.trace_to_samples(trace, data)\n", - "mcmc_samples.wrap_K()" + "mcmc_samples = tj.JokerSamples.from_inference_data(prior, trace, data)\n", + "mcmc_samples = mcmc_samples.wrap_K()" ] }, { @@ -242,7 +234,7 @@ "source": [ "df = mcmc_samples.tbl.to_pandas()\n", "colnames = mcmc_samples.par_names\n", - "colnames.pop(colnames.index('s'))\n", + "colnames.pop(colnames.index(\"s\"))\n", "_ = corner.corner(df[colnames])" ] } @@ -263,7 +255,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.3" + "version": "3.11.8" }, "toc": { "base_numbering": 1, diff --git a/docs/examples/Strader-circular-only.ipynb b/docs/examples/Strader-circular-only.ipynb index 9747d77c..c214d10b 100644 --- a/docs/examples/Strader-circular-only.ipynb +++ b/docs/examples/Strader-circular-only.ipynb @@ -26,22 +26,20 @@ "metadata": {}, "outputs": [], "source": [ - "from astropy.io import ascii\n", - "from astropy.time import Time\n", + "import arviz as az\n", "import astropy.units as u\n", + "import corner\n", "import matplotlib.pyplot as plt\n", - "%matplotlib inline\n", "import numpy as np\n", + "import pymc as pm\n", + "import pymc_ext as pmx\n", + "import pytensor.tensor as pt\n", + "import thejoker as tj\n", + "import thejoker.units as xu\n", + "from astropy.io import ascii\n", + "from astropy.time import Time\n", "\n", - "import pymc3 as pm\n", - "import pymc3_ext as pmx\n", - "import exoplanet.units as xu\n", - "import exoplanet as xo\n", - "from pymc3_ext.distributions import Angle\n", - "import corner\n", - "import arviz as az\n", - "\n", - "import thejoker as tj" + "%matplotlib inline" ] }, { @@ -114,9 +112,10 @@ " 2458247.5435496 160.5 14.2\n", " 2458278.5472619 197.1 15.9\n", " 2458278.5613912 183.7 15.7\"\"\",\n", - " names=['BJD', 'rv', 'rv_err'])\n", - "tbl['rv'].unit = u.km/u.s\n", - "tbl['rv_err'].unit = u.km/u.s" + " names=[\"BJD\", \"rv\", \"rv_err\"],\n", + ")\n", + "tbl[\"rv\"].unit = u.km / u.s\n", + "tbl[\"rv_err\"].unit = u.km / u.s" ] }, { @@ -133,9 +132,10 @@ "outputs": [], "source": [ "data = tj.RVData(\n", - " t=Time(tbl['BJD'], format='jd', scale='tcb'),\n", - " rv=u.Quantity(tbl['rv']),\n", - " rv_err=u.Quantity(tbl['rv_err']))" + " t=Time(tbl[\"BJD\"], format=\"jd\", scale=\"tcb\"),\n", + " rv=u.Quantity(tbl[\"rv\"]),\n", + " rv_err=u.Quantity(tbl[\"rv_err\"]),\n", + ")" ] }, { @@ -169,14 +169,16 @@ "source": [ "with pm.Model() as model:\n", " # Allow extra error to account for under-estimated error bars\n", - " e = xu.with_unit(pm.Constant('e', 0),\n", - " u.one)\n", + " e = xu.with_unit(pm.ConstantData(\"e\", 0), u.one)\n", + " omega = xu.with_unit(pm.ConstantData(\"omega\", 0), u.rad)\n", "\n", " prior = tj.JokerPrior.default(\n", - " P_min=0.1*u.day, P_max=100*u.day, # Range of periods to consider\n", - " sigma_K0=50*u.km/u.s, P0=1*u.year, # scale of the prior on semiamplitude, K\n", - " sigma_v=50*u.km/u.s, # std dev of the prior on the systemic velocity, v0\n", - " pars={'e': e}\n", + " P_min=0.1 * u.day,\n", + " P_max=100 * u.day, # Range of periods to consider\n", + " sigma_K0=50 * u.km / u.s,\n", + " P0=1 * u.year, # scale of the prior on semiamplitude, K\n", + " sigma_v=50 * u.km / u.s, # std dev of the prior on the systemic velocity, v0\n", + " pars={\"e\": e, \"omega\": omega},\n", " )" ] }, @@ -195,9 +197,9 @@ "source": [ "# Run rejection sampling with The Joker:\n", "joker = tj.TheJoker(prior, rng=rnd)\n", - "samples = joker.rejection_sample(data,\n", - " prior_samples=100_000,\n", - " max_posterior_samples=256)\n", + "samples = joker.rejection_sample(\n", + " data, prior_samples=1_000_000, max_posterior_samples=100\n", + ")\n", "samples" ] }, @@ -230,7 +232,7 @@ "metadata": {}, "outputs": [], "source": [ - "tj.plot_phase_fold(samples, data=data);" + "_ = tj.plot_phase_fold(samples[0], data=data)" ] }, { @@ -246,35 +248,40 @@ "metadata": {}, "outputs": [], "source": [ - "import aesara_theano_fallback.tensor as tt\n", - "with pm.Model():\n", - "\n", - " # To sample with pymc3, we have to set any constant variables\n", - " # as \"Deterministic\" objects. We can ignore eccentricity and\n", - " # the argument of pericenter by setting them both to 0:\n", - " e = xu.with_unit(pm.Deterministic('e', tt.constant(0)),\n", - " u.one)\n", - " omega = xu.with_unit(pm.Deterministic('omega', tt.constant(0)),\n", - " u.radian)\n", + "with pm.Model() as model:\n", + " # To sample with pymc, we have to set any constant variables as \"Deterministic\"\n", + " # objects. We can ignore eccentricity and the argument of pericenter by setting\n", + " # them both to 0:\n", + " e = xu.with_unit(pm.Deterministic(\"e\", pt.constant(0)), u.one)\n", + " omega = xu.with_unit(pm.Deterministic(\"omega\", pt.constant(0)), u.radian)\n", "\n", " # We use the same prior parameters as before:\n", " prior_mcmc = tj.JokerPrior.default(\n", - " P_min=0.1*u.day, P_max=10*u.day,\n", - " sigma_K0=50*u.km/u.s, P0=1*u.year,\n", - " sigma_v=50*u.km/u.s,\n", - " pars={'e': e, 'omega': omega}\n", + " P_min=0.1 * u.day,\n", + " P_max=10 * u.day,\n", + " sigma_K0=50 * u.km / u.s,\n", + " P0=1 * u.year,\n", + " sigma_v=50 * u.km / u.s,\n", + " pars={\"e\": e, \"omega\": omega},\n", " )\n", "\n", - " # Now we use the sample returned from The Joker to set up\n", - " # our initialization for standard MCMC:\n", + " # Now we use the sample returned from The Joker to set up our initialization for\n", + " # standard MCMC:\n", " joker_mcmc = tj.TheJoker(prior_mcmc, rng=rnd)\n", " mcmc_init = joker_mcmc.setup_mcmc(data, samples)\n", "\n", - " trace = pmx.sample(\n", - " tune=500, draws=1000,\n", - " start=mcmc_init,\n", + " opt_init = pmx.optimize(mcmc_init)\n", + "\n", + " trace = pm.sample(\n", + " tune=1000,\n", + " draws=1000,\n", + " start=opt_init,\n", " random_seed=seed,\n", - " cores=1, chains=2)" + " cores=1,\n", + " chains=2,\n", + " init=\"adapt_full\",\n", + " target_accept=0.9,\n", + " )" ] }, { @@ -290,7 +297,11 @@ "metadata": {}, "outputs": [], "source": [ - "az.summary(trace, var_names=prior_mcmc.par_names)" + "par_names = prior_mcmc.par_names.copy()\n", + "par_names.pop(par_names.index(\"e\"))\n", + "par_names.pop(par_names.index(\"omega\"))\n", + "# par_names.pop(par_names.index(\"s\"))\n", + "az.summary(trace, var_names=par_names)" ] }, { @@ -306,8 +317,8 @@ "metadata": {}, "outputs": [], "source": [ - "mcmc_samples = joker_mcmc.trace_to_samples(trace, data=data)\n", - "mcmc_samples.wrap_K()" + "mcmc_samples = tj.JokerSamples.from_inference_data(prior, trace, data=data)\n", + "mcmc_samples = mcmc_samples.wrap_K()" ] }, { @@ -324,11 +335,13 @@ "outputs": [], "source": [ "df = mcmc_samples.tbl.to_pandas()\n", - "df = df.drop(columns=['e', 's', 'omega'])\n", + "df = df.drop(columns=[\"e\", \"s\", \"omega\", \"ln_posterior\", \"ln_prior\", \"ln_likelihood\"])\n", "\n", "true_P = 1.097195\n", "true_T0 = 2457780.8373\n", - "true_M0 = (2*np.pi * (true_T0 - data.t_ref.tcb.jd) / true_P) % (2*np.pi) - (2*np.pi)\n", + "true_M0 = (2 * np.pi * (true_T0 - data.t_ref.tcb.jd) / true_P) % (2 * np.pi) - (\n", + " 2 * np.pi\n", + ")\n", "_ = corner.corner(df, truths=[true_P, true_M0, 210, 32.0])" ] }, @@ -351,9 +364,9 @@ "_ = tj.plot_phase_fold(mcmc_samples.mean(), data, ax=axes[1], residual=True)\n", "\n", "for ax in axes:\n", - " ax.set_ylabel(f'RV [{data.rv.unit:latex_inline}]')\n", + " ax.set_ylabel(f\"RV [{data.rv.unit:latex_inline}]\")\n", "\n", - "axes[1].axhline(0, zorder=-10, color='tab:green', alpha=0.5)\n", + "axes[1].axhline(0, zorder=-10, color=\"tab:green\", alpha=0.5)\n", "axes[1].set_ylim(-50, 50)" ] }, @@ -381,7 +394,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.3" + "version": "3.11.8" }, "toc": { "base_numbering": 1, diff --git a/docs/examples/Thompson-black-hole.ipynb b/docs/examples/Thompson-black-hole.ipynb index 8f4f3417..0a1d91b2 100644 --- a/docs/examples/Thompson-black-hole.ipynb +++ b/docs/examples/Thompson-black-hole.ipynb @@ -26,22 +26,18 @@ "metadata": {}, "outputs": [], "source": [ - "from astropy.io import ascii\n", - "from astropy.time import Time\n", + "import arviz as az\n", "import astropy.units as u\n", + "import corner\n", "import matplotlib.pyplot as plt\n", - "%matplotlib inline\n", "import numpy as np\n", - "\n", - "import pymc3 as pm\n", - "import pymc3_ext as pmx\n", - "import exoplanet.units as xu\n", - "import exoplanet as xo\n", - "import corner\n", - "import arviz as az\n", - "\n", + "import pymc as pm\n", "import thejoker as tj\n", - "from twobody.transforms import get_m2_min" + "import thejoker.units as xu\n", + "from astropy.io import ascii\n", + "from astropy.time import Time\n", + "\n", + "%matplotlib inline" ] }, { @@ -82,9 +78,10 @@ " 8123.79627 -25.810 0.115\n", " 8136.59960 15.691 0.146\n", " 8143.78352 34.281 0.087\"\"\",\n", - " names=['HJD', 'rv', 'rv_err'])\n", - "tres_tbl['rv'].unit = u.km/u.s\n", - "tres_tbl['rv_err'].unit = u.km/u.s" + " names=[\"HJD\", \"rv\", \"rv_err\"],\n", + ")\n", + "tres_tbl[\"rv\"].unit = u.km / u.s\n", + "tres_tbl[\"rv_err\"].unit = u.km / u.s" ] }, { @@ -97,9 +94,10 @@ " \"\"\"6204.95544 -37.417 0.011\n", " 6229.92499 34.846 0.010\n", " 6233.87715 42.567 0.010\"\"\",\n", - " names=['HJD', 'rv', 'rv_err'])\n", - "apogee_tbl['rv'].unit = u.km/u.s\n", - "apogee_tbl['rv_err'].unit = u.km/u.s" + " names=[\"HJD\", \"rv\", \"rv_err\"],\n", + ")\n", + "apogee_tbl[\"rv\"].unit = u.km / u.s\n", + "apogee_tbl[\"rv_err\"].unit = u.km / u.s" ] }, { @@ -109,14 +107,16 @@ "outputs": [], "source": [ "tres_data = tj.RVData(\n", - " t=Time(tres_tbl['HJD'] + 2450000, format='jd', scale='tcb'),\n", - " rv=u.Quantity(tres_tbl['rv']),\n", - " rv_err=u.Quantity(tres_tbl['rv_err']))\n", + " t=Time(tres_tbl[\"HJD\"] + 2450000, format=\"jd\", scale=\"tcb\"),\n", + " rv=u.Quantity(tres_tbl[\"rv\"]),\n", + " rv_err=u.Quantity(tres_tbl[\"rv_err\"]),\n", + ")\n", "\n", "apogee_data = tj.RVData(\n", - " t=Time(apogee_tbl['HJD'] + 2450000, format='jd', scale='tcb'),\n", - " rv=u.Quantity(apogee_tbl['rv']),\n", - " rv_err=u.Quantity(apogee_tbl['rv_err']))" + " t=Time(apogee_tbl[\"HJD\"] + 2450000, format=\"jd\", scale=\"tcb\"),\n", + " rv=u.Quantity(apogee_tbl[\"rv\"]),\n", + " rv_err=u.Quantity(apogee_tbl[\"rv_err\"]),\n", + ")" ] }, { @@ -132,7 +132,7 @@ "metadata": {}, "outputs": [], "source": [ - "for d, name in zip([tres_data, apogee_data], ['TRES', 'APOGEE']):\n", + "for d, name in zip([tres_data, apogee_data], [\"TRES\", \"APOGEE\"]):\n", " d.plot(color=None, label=name)\n", "plt.legend(fontsize=18)" ] @@ -176,14 +176,15 @@ "source": [ "with pm.Model() as model:\n", " # Allow extra error to account for under-estimated error bars\n", - " s = xu.with_unit(pm.Lognormal('s', -2, 1),\n", - " u.km/u.s)\n", + " s = xu.with_unit(pm.Lognormal(\"s\", -2, 1), u.km / u.s)\n", "\n", " prior = tj.JokerPrior.default(\n", - " P_min=16*u.day, P_max=128*u.day, # Range of periods to consider\n", - " sigma_K0=30*u.km/u.s, P0=1*u.year, # scale of the prior on semiamplitude, K\n", - " sigma_v=25*u.km/u.s, # std dev of the prior on the systemic velocity, v0\n", - " s=s\n", + " P_min=16 * u.day,\n", + " P_max=128 * u.day, # Range of periods to consider\n", + " sigma_K0=30 * u.km / u.s,\n", + " P0=1 * u.year, # scale of the prior on semiamplitude, K\n", + " sigma_v=25 * u.km / u.s, # std dev of the prior on the systemic velocity, v0\n", + " s=s,\n", " )" ] }, @@ -201,8 +202,7 @@ "outputs": [], "source": [ "# Generate a large number of prior samples:\n", - "prior_samples = prior.sample(size=1_000_000,\n", - " rng=rnd)" + "prior_samples = prior.sample(size=1_000_000, rng=rnd)" ] }, { @@ -213,8 +213,7 @@ "source": [ "# Run rejection sampling with The Joker:\n", "joker = tj.TheJoker(prior, rng=rnd)\n", - "samples = joker.rejection_sample(tres_data, prior_samples,\n", - " max_posterior_samples=256)\n", + "samples = joker.rejection_sample(tres_data, prior_samples, max_posterior_samples=256)\n", "samples" ] }, @@ -252,7 +251,7 @@ "metadata": {}, "outputs": [], "source": [ - "samples.tbl['P', 'e', 'K']" + "samples.tbl[\"P\", \"e\", \"K\"]" ] }, { @@ -270,7 +269,7 @@ "metadata": {}, "outputs": [], "source": [ - "_ = tres_data.plot(phase_fold=samples[0]['P'])" + "_ = tres_data.plot(phase_fold=samples[0][\"P\"])" ] }, { @@ -313,8 +312,8 @@ "metadata": {}, "outputs": [], "source": [ - "tres_data.plot(color=None, phase_fold=np.mean(samples['P']))\n", - "apogee_data.plot(color=None, phase_fold=np.mean(samples['P']))" + "tres_data.plot(color=None, phase_fold_period=np.mean(samples[\"P\"]))\n", + "apogee_data.plot(color=None, phase_fold_period=np.mean(samples[\"P\"]))" ] }, { @@ -335,25 +334,24 @@ "with pm.Model() as model:\n", " # The parameter that represents the constant velocity offset between\n", " # APOGEE and TRES:\n", - " dv0_1 = xu.with_unit(pm.Normal('dv0_1', 0, 5.),\n", - " u.km/u.s)\n", + " dv0_1 = xu.with_unit(pm.Normal(\"dv0_1\", 0, 5.0), u.km / u.s)\n", "\n", " # The same extra uncertainty parameter as previously defined\n", - " s = xu.with_unit(pm.Lognormal('s', -2, 1),\n", - " u.km/u.s)\n", + " s = xu.with_unit(pm.Lognormal(\"s\", -2, 1), u.km / u.s)\n", "\n", " # We can restrict the prior on prior now, using the above\n", " prior_joint = tj.JokerPrior.default(\n", " # P_min=16*u.day, P_max=128*u.day,\n", - " P_min=75*u.day, P_max=90*u.day,\n", - " sigma_K0=30*u.km/u.s, P0=1*u.year,\n", - " sigma_v=25*u.km/u.s,\n", + " P_min=75 * u.day,\n", + " P_max=90 * u.day,\n", + " sigma_K0=30 * u.km / u.s,\n", + " P0=1 * u.year,\n", + " sigma_v=25 * u.km / u.s,\n", " v0_offsets=[dv0_1],\n", - " s=s\n", + " s=s,\n", " )\n", "\n", - "prior_samples_joint = prior_joint.sample(size=10_000_000,\n", - " rng=rnd)" + "prior_samples_joint = prior_joint.sample(size=1_000_000, rng=rnd)" ] }, { @@ -364,9 +362,9 @@ "source": [ "# Run rejection sampling with The Joker:\n", "joker_joint = tj.TheJoker(prior_joint, rng=rnd)\n", - "samples_joint = joker_joint.rejection_sample(data,\n", - " prior_samples_joint,\n", - " max_posterior_samples=256)\n", + "samples_joint = joker_joint.rejection_sample(\n", + " data, prior_samples_joint, max_posterior_samples=256\n", + ")\n", "samples_joint" ] }, @@ -374,7 +372,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Here we again only get one sample back from *The Joker*, because these ata are so constraining:" + "Here we again only get one sample back from *The Joker*, because these data are so constraining:" ] }, { @@ -399,40 +397,37 @@ "metadata": {}, "outputs": [], "source": [ - "from pymc3_ext.distributions import Angle\n", + "from pymc_ext.distributions import angle\n", "\n", "with pm.Model():\n", - "\n", " # See note above: when running MCMC, we will sample in the parameters\n", " # (M0 - omega, omega) instead of (M0, omega)\n", - " M0_m_omega = xu.with_unit(Angle('M0_m_omega'), u.radian)\n", - " omega = xu.with_unit(Angle('omega'), u.radian)\n", - " # M0 = xu.with_unit(Angle('M0'), u.radian)\n", - " M0 = xu.with_unit(pm.Deterministic('M0', M0_m_omega + omega),\n", - " u.radian)\n", + " M0_m_omega = xu.with_unit(angle(\"M0_m_omega\"), u.radian)\n", + " omega = xu.with_unit(angle(\"omega\"), u.radian)\n", + " # M0 = xu.with_unit(angle('M0'), u.radian)\n", + " M0 = xu.with_unit(pm.Deterministic(\"M0\", M0_m_omega + omega), u.radian)\n", "\n", " # The same offset and extra uncertainty parameters as above:\n", - " dv0_1 = xu.with_unit(pm.Normal('dv0_1', 0, 5.), u.km/u.s)\n", - " s = xu.with_unit(pm.Lognormal('s', -2, 0.5),\n", - " u.km/u.s)\n", + " dv0_1 = xu.with_unit(pm.Normal(\"dv0_1\", 0, 5.0), u.km / u.s)\n", + " s = xu.with_unit(pm.Lognormal(\"s\", -2, 0.5), u.km / u.s)\n", "\n", " prior_mcmc = tj.JokerPrior.default(\n", - " P_min=16*u.day, P_max=128*u.day,\n", - " sigma_K0=30*u.km/u.s, P0=1*u.year,\n", - " sigma_v=25*u.km/u.s,\n", + " P_min=16 * u.day,\n", + " P_max=128 * u.day,\n", + " sigma_K0=30 * u.km / u.s,\n", + " P0=1 * u.year,\n", + " sigma_v=25 * u.km / u.s,\n", " v0_offsets=[dv0_1],\n", " s=s,\n", - " pars={'M0': M0, 'omega': omega}\n", + " pars={\"M0\": M0, \"omega\": omega},\n", " )\n", "\n", " joker_mcmc = tj.TheJoker(prior_mcmc, rng=rnd)\n", " mcmc_init = joker_mcmc.setup_mcmc(data, samples_joint)\n", "\n", - " trace = pmx.sample(\n", - " tune=500, draws=1000,\n", - " start=mcmc_init,\n", - " random_seed=seed,\n", - " cores=1, chains=2)" + " trace = pm.sample(\n", + " tune=1000, draws=500, start=mcmc_init, random_seed=seed, cores=1, chains=2\n", + " )" ] }, { @@ -464,8 +459,8 @@ "metadata": {}, "outputs": [], "source": [ - "mcmc_samples = joker_mcmc.trace_to_samples(trace, data=data)\n", - "mcmc_samples.wrap_K()" + "mcmc_samples = tj.JokerSamples.from_inference_data(prior_joint, trace, data=data)\n", + "mcmc_samples = mcmc_samples.wrap_K()" ] }, { @@ -500,13 +495,13 @@ "source": [ "fig, axes = plt.subplots(2, 1, figsize=(6, 8), sharex=True)\n", "\n", - "_ = tj.plot_phase_fold(mcmc_samples.median(), data, ax=axes[0], add_labels=False)\n", - "_ = tj.plot_phase_fold(mcmc_samples.median(), data, ax=axes[1], residual=True)\n", + "_ = tj.plot_phase_fold(mcmc_samples.median_period(), data, ax=axes[0], add_labels=False)\n", + "_ = tj.plot_phase_fold(mcmc_samples.median_period(), data, ax=axes[1], residual=True)\n", "\n", "for ax in axes:\n", - " ax.set_ylabel(f'RV [{apogee_data.rv.unit:latex_inline}]')\n", + " ax.set_ylabel(f\"RV [{apogee_data.rv.unit:latex_inline}]\")\n", "\n", - "axes[1].axhline(0, zorder=-10, color='tab:green', alpha=0.5)\n", + "axes[1].axhline(0, zorder=-10, color=\"tab:green\", alpha=0.5)\n", "axes[1].set_ylim(-1, 1)" ] }, @@ -523,17 +518,16 @@ "metadata": {}, "outputs": [], "source": [ - "mfs = u.Quantity([mcmc_samples.get_orbit(i).m_f\n", - " for i in np.random.choice(len(mcmc_samples), 1024)])\n", - "plt.hist(mfs.to_value(u.Msun), bins=32);\n", - "plt.xlabel(rf'$f(M)$ [{u.Msun:latex_inline}]');\n", - "\n", + "mfs = u.Quantity(\n", + " [mcmc_samples.get_orbit(i).m_f for i in rnd.choice(len(mcmc_samples), 1024)]\n", + ")\n", + "plt.hist(mfs.to_value(u.Msun), bins=32)\n", + "plt.xlabel(rf\"$f(M)$ [{u.Msun:latex_inline}]\")\n", "# Values from Thompson et al., showing 1-sigma region\n", - "plt.axvline(0.766, zorder=100, color='tab:orange')\n", - "plt.axvspan(0.766 - 0.00637,\n", - " 0.766 + 0.00637,\n", - " zorder=10, color='tab:orange',\n", - " alpha=0.4, lw=0)" + "plt.axvline(0.766, zorder=100, color=\"tab:orange\")\n", + "plt.axvspan(\n", + " 0.766 - 0.00637, 0.766 + 0.00637, zorder=10, color=\"tab:orange\", alpha=0.4, lw=0\n", + ")" ] }, { @@ -560,7 +554,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.3" + "version": "3.11.8" }, "toc": { "base_numbering": 1, diff --git a/docs/examples/notebook_setup.py b/docs/examples/notebook_setup.py index e81a2a48..5f9c0c91 100644 --- a/docs/examples/notebook_setup.py +++ b/docs/examples/notebook_setup.py @@ -15,6 +15,3 @@ plt.rcParams["figure.dpi"] = 200 plt.rcParams["font.size"] = 16 plt.rcParams["font.family"] = "serif" -plt.rcParams["font.serif"] = ["Liberation Serif"] -plt.rcParams["font.cursive"] = ["Liberation Serif"] -plt.rcParams["mathtext.fontset"] = "custom" From ef5560dc5b6ac4f4cc91c33efece71e48ae17b91 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Tue, 5 Mar 2024 21:11:18 -0500 Subject: [PATCH 27/50] corner py --- pyproject.toml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index f98a534d..74e7c178 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,7 +58,9 @@ docs = [ ] tutorials = [ "thejoker[docs]", - "jupyter-client" + "jupyter-client", + "corner", + "arviz" ] [tool.setuptools.packages.find] From f7548f45eb8bf1f82796586d8b99fa633f0ac718 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Tue, 5 Mar 2024 21:35:27 -0500 Subject: [PATCH 28/50] what why --- thejoker/tests/test_sampler.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/thejoker/tests/test_sampler.py b/thejoker/tests/test_sampler.py index 2f70f228..e16418ff 100644 --- a/thejoker/tests/test_sampler.py +++ b/thejoker/tests/test_sampler.py @@ -213,9 +213,9 @@ def test_iterative_rejection_sample(tmpdir, prior): @pytest.mark.parametrize("prior", priors) def test_continue_mcmc(prior): - data, orbit = make_data(n_times=10) + data, orbit = make_data(n_times=8) - prior_samples = prior.sample(size=16384, return_logprobs=True) + prior_samples = prior.sample(size=16384) joker = TheJoker(prior) joker_samples = joker.rejection_sample(data, prior_samples) From 54543d9b721d8455077f63440070d9ed24bc622e Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Tue, 5 Mar 2024 22:09:52 -0500 Subject: [PATCH 29/50] skip mcmc test on macos ci --- thejoker/tests/test_sampler.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/thejoker/tests/test_sampler.py b/thejoker/tests/test_sampler.py index e16418ff..66e614c0 100644 --- a/thejoker/tests/test_sampler.py +++ b/thejoker/tests/test_sampler.py @@ -1,4 +1,5 @@ import os +import platform import astropy.units as u import numpy as np @@ -212,6 +213,7 @@ def test_iterative_rejection_sample(tmpdir, prior): @pytest.mark.parametrize("prior", priors) +@pytest.mark.skipif(platform.system() == "Darwin", reason="Test fails on MacOS CI") def test_continue_mcmc(prior): data, orbit = make_data(n_times=8) From e9c4574dd5f4c372465940db884c47f3221eb0a6 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Tue, 5 Mar 2024 22:20:43 -0500 Subject: [PATCH 30/50] oops update docs config and run notebooks --- docs/conf.py | 18 ++++++++++++++++++ docs/run_notebooks.py | 15 ++++++++++++--- 2 files changed, 30 insertions(+), 3 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index 0955fe0f..f1cccf9f 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -14,6 +14,7 @@ "sphinx_copybutton", "sphinx_automodapi.automodapi", "sphinx_automodapi.smart_resolver", + "rtds_action", ] source_suffix = [".rst", ".md"] @@ -60,3 +61,20 @@ ] always_document_param_types = True + +# The name of your GitHub repository +rtds_action_github_repo = "adrn/thejoker" + +# The path where the artifact should be extracted +# Note: this is relative to the conf.py file! +rtds_action_path = "examples" + +# The "prefix" used in the `upload-artifact` step of the action +rtds_action_artifact_prefix = "notebooks-for-" + +# A GitHub personal access token is required, more info below +rtds_action_github_token = os.environ["GITHUB_TOKEN"] + +# Whether or not to raise an error on Read the Docs if the +# artifact containing the notebooks can't be downloaded (optional) +rtds_action_error_if_missing = False diff --git a/docs/run_notebooks.py b/docs/run_notebooks.py index 7a174b3f..d0d32ff4 100644 --- a/docs/run_notebooks.py +++ b/docs/run_notebooks.py @@ -19,19 +19,23 @@ def process_notebook(filename, kernel_name=None): notebook = nbformat.read(f, as_version=4) ep = ExecutePreprocessor(timeout=-1, kernel_name=kernel_name) - ep.log.setLevel(logging.DEBUG) + ep.log.setLevel(logging.INFO) ep.log.addHandler(logging.StreamHandler()) + success = True try: ep.preprocess(notebook, {"metadata": {"path": "examples/"}}) except CellExecutionError as e: - msg = f"error while running: {filename}\n\n" + msg = f"Error while running: {filename}\n\n" msg += e.traceback print(msg) + success = False finally: with open(os.path.join(filename), mode="w") as f: nbformat.write(notebook, f) + return success + if __name__ == "__main__": if len(sys.argv) == 2: @@ -42,5 +46,10 @@ def process_notebook(filename, kernel_name=None): nbsphinx_kernel_name = os.environ.get("NBSPHINX_KERNEL", "python3") + fail = False for filename in sorted(glob.glob(pattern)): - process_notebook(filename, kernel_name=nbsphinx_kernel_name) + success = process_notebook(filename, kernel_name=nbsphinx_kernel_name) + fail = fail or not success + + if fail: + sys.exit(1) From 23410b9d2be3b18fd787ce9e9fc08102530bb8cd Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Tue, 5 Mar 2024 22:38:16 -0500 Subject: [PATCH 31/50] REVERT THIS - testing if multipool is the issue --- thejoker/tests/test_sampler.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/thejoker/tests/test_sampler.py b/thejoker/tests/test_sampler.py index 66e614c0..c89c7708 100644 --- a/thejoker/tests/test_sampler.py +++ b/thejoker/tests/test_sampler.py @@ -1,12 +1,11 @@ import os -import platform import astropy.units as u import numpy as np import pymc as pm import pytest from astropy.time import Time -from schwimmbad import MultiPool, SerialPool +from schwimmbad import SerialPool from twobody import KeplerOrbit from thejoker.data import RVData @@ -103,10 +102,10 @@ def test_marginal_ln_likelihood(tmpdir, case): assert len(ll) == len(prior_samples) # NOTE: this makes it so I can't parallelize tests, I think - with MultiPool(processes=2) as pool: - joker = TheJoker(prior, pool=pool) - ll = joker.marginal_ln_likelihood(data, filename) - assert len(ll) == len(prior_samples) + # with MultiPool(processes=2) as pool: + # joker = TheJoker(prior, pool=pool) + # ll = joker.marginal_ln_likelihood(data, filename) + # assert len(ll) == len(prior_samples) priors = [ @@ -213,7 +212,6 @@ def test_iterative_rejection_sample(tmpdir, prior): @pytest.mark.parametrize("prior", priors) -@pytest.mark.skipif(platform.system() == "Darwin", reason="Test fails on MacOS CI") def test_continue_mcmc(prior): data, orbit = make_data(n_times=8) From c9073802b54d27e1f28e0a63b14ef198e89f7f9d Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Tue, 5 Mar 2024 22:40:41 -0500 Subject: [PATCH 32/50] fix make-data notebook --- docs/examples/data-survey1.ecsv | 18 +- docs/examples/data-survey2.ecsv | 14 +- docs/examples/data-triple.ecsv | 460 ++++++++++++++-------------- docs/examples/data.ecsv | 426 +++++++++++++------------- docs/examples/make-data.ipynb | 192 ++++++------ docs/examples/true-orbit-triple.pkl | Bin 1659 -> 1545 bytes docs/examples/true-orbit.pkl | Bin 1659 -> 1545 bytes 7 files changed, 560 insertions(+), 550 deletions(-) diff --git a/docs/examples/data-survey1.ecsv b/docs/examples/data-survey1.ecsv index b9d87895..dfe8f1c8 100644 --- a/docs/examples/data-survey1.ecsv +++ b/docs/examples/data-survey1.ecsv @@ -1,11 +1,11 @@ -# %ECSV 0.9 +# %ECSV 1.0 # --- # datatype: # - {name: bjd, datatype: float64} # - {name: rv, unit: km / s, datatype: float64} # - {name: rv_err, unit: km / s, datatype: float64} # meta: !!omap -# - t_ref: !astropy.time.Time {format: mjd, in_subfmt: '*', jd1: 2458515.0, jd2: 0.15522072320891311, out_subfmt: '*', precision: 3, scale: tcb} +# - t_ref: !astropy.time.Time {format: mjd, in_subfmt: '*', jd1: 2458515.0, jd2: 0.155220723230741, out_subfmt: '*', precision: 3, scale: tcb} # - __serialized_columns__: # rv: # __class__: astropy.units.quantity.Quantity @@ -18,12 +18,12 @@ # schema: astropy-2.0 bjd rv rv_err 2458515.1552207232 43.47671387955206 2.4227872187885664 -2458526.8309797375 38.000886500699686 2.0527770245831274 -2458539.856422522 32.47042267761668 1.3438940211186683 -2458543.489005615 31.693867158629324 0.14023090332462226 -2458558.941049686 37.15189528120144 0.4735940895006947 +2458526.8309797375 38.000886500717556 2.0527770245831274 +2458539.856422522 32.4704226776271 1.3438940211186683 +2458543.489005615 31.693867158632205 0.14023090332462226 +2458558.941049686 37.15189528117272 0.4735940895006947 2458579.227795638 46.99827383411713 0.3133784294944628 2458584.2517544227 46.30732167983992 2.9299449995725007 -2458588.7661287608 44.99952569375152 2.3559405959648614 -2458602.0750950654 39.15529512761318 2.3485689715740943 -2458628.016668719 32.203589649183016 1.8095757137276618 +2458588.7661287608 44.99952569376658 2.3559405959648614 +2458602.0750950654 39.155295127616746 2.3485689715740943 +2458628.016668719 32.20358964917636 1.8095757137276618 diff --git a/docs/examples/data-survey2.ecsv b/docs/examples/data-survey2.ecsv index 11cbf2f2..931d2f5b 100644 --- a/docs/examples/data-survey2.ecsv +++ b/docs/examples/data-survey2.ecsv @@ -1,11 +1,11 @@ -# %ECSV 0.9 +# %ECSV 1.0 # --- # datatype: # - {name: bjd, datatype: float64} # - {name: rv, unit: km / s, datatype: float64} # - {name: rv_err, unit: km / s, datatype: float64} # meta: !!omap -# - t_ref: !astropy.time.Time {format: mjd, in_subfmt: '*', jd1: 2458515.0, jd2: 0.15522072320891311, out_subfmt: '*', precision: 3, scale: tcb} +# - t_ref: !astropy.time.Time {format: mjd, in_subfmt: '*', jd1: 2458515.0, jd2: 0.155220723230741, out_subfmt: '*', precision: 3, scale: tcb} # - __serialized_columns__: # rv: # __class__: astropy.units.quantity.Quantity @@ -17,10 +17,10 @@ # value: !astropy.table.SerializedColumn {name: rv_err} # schema: astropy-2.0 bjd rv rv_err -2458638.66232473 43.67199782467397 1.3302801152625319 -2458644.8434247165 48.33668952819079 1.1645032359248524 +2458638.66232473 43.671997824638325 1.3302801152625319 +2458644.8434247165 48.33668952816298 1.1645032359248524 2458669.3095975104 48.69619418101378 0.6171738434348514 -2458699.6537869107 36.38900512229526 0.7471665495221848 -2458705.4573607757 36.95849243399448 0.5252801133029275 -2458714.947581842 42.61001209445934 0.5739898939015465 +2458699.6537869107 36.389005122298784 0.7471665495221848 +2458705.4573607757 36.95849243399019 0.5252801133029275 +2458714.947581842 42.61001209443589 0.5739898939015465 2458716.2755749053 43.690038798890654 0.17440907292151522 diff --git a/docs/examples/data-triple.ecsv b/docs/examples/data-triple.ecsv index a3af9061..7dbc6a85 100644 --- a/docs/examples/data-triple.ecsv +++ b/docs/examples/data-triple.ecsv @@ -1,4 +1,4 @@ -# %ECSV 0.9 +# %ECSV 1.0 # --- # datatype: # - {name: bjd, datatype: float64} @@ -17,259 +17,259 @@ # schema: astropy-2.0 bjd rv rv_err 2458499.9378344314 39.27095201189804 2.008993781154447 -2458500.9692732077 38.905226908148286 0.1656822971288193 -2458501.7584133083 38.58976817820053 0.8719830933655139 -2458502.3954306515 38.3114650549298 0.24761236240699092 -2458504.3471037457 37.31887115135462 0.41825281286097593 -2458506.482957222 35.97402027320257 0.506793124220416 -2458506.8822173015 35.69160868925746 0.7611942873357234 -2458507.0216006897 35.59072661860586 0.6015601071517347 -2458507.198626581 35.46090285889963 0.35209422894817444 -2458507.377715795 35.32764653259349 0.11499884562518728 -2458510.128967821 33.06034467265575 0.15484845891705262 -2458510.837018213 32.424660485262976 0.8629516490335437 +2458500.9692732077 38.905226908145536 0.1656822971288193 +2458501.7584133083 38.58976817820663 0.8719830933655139 +2458502.3954306515 38.31146505491988 0.24761236240699092 +2458504.3471037457 37.31887115133818 0.41825281286097593 +2458506.482957222 35.97402027319752 0.506793124220416 +2458506.8822173015 35.691608689262694 0.7611942873357234 +2458507.0216006897 35.59072661859527 0.6015601071517347 +2458507.198626581 35.46090285887812 0.35209422894817444 +2458507.377715795 35.32764653257167 0.11499884562518728 +2458510.128967821 33.06034467262993 0.15484845891705262 +2458510.837018213 32.42466048527617 0.8629516490335437 2458511.520167288 31.801897508448615 0.8110108322490306 2458513.549676771 29.986619236276628 1.0055607774143827 -2458517.3322867523 27.730452849302385 0.13199744964204602 -2458520.7528358875 28.356606585437845 0.2369709084395939 +2458517.3322867523 27.730452849297215 0.13199744964204602 +2458520.7528358875 28.356606585429397 0.2369709084395939 2458521.5853484576 28.905389893808568 0.22506374732093037 -2458522.8070115624 29.908838492129554 1.6670426945676025 -2458525.316756085 32.38359972707677 1.3794384334836447 -2458526.4185919776 33.521605851345825 1.3411034225878844 -2458527.576654805 34.694856672563645 2.8698173476737803 -2458528.111970773 35.22199559575967 2.364704841833626 -2458529.8390300074 36.83353391585537 0.9249374796599684 -2458530.604024408 37.49811380270383 0.5797835892283674 -2458531.771079307 38.45001802514051 1.6845389508832564 -2458531.9641530444 38.600231741031195 0.6675655129694912 -2458532.6917056805 39.14787066443605 1.4171901284667194 -2458534.0926596513 40.122103514701195 0.1863187081001272 -2458534.1499005174 40.15971817042143 0.5881775840787171 -2458537.076265389 41.868260456636456 0.5313839173251673 -2458537.191489296 41.927368396503674 0.53330536067345 -2458537.8113832935 42.23542071367911 0.6015267565751019 +2458522.8070115624 29.908838492116473 1.6670426945676025 +2458525.316756085 32.383599727099316 1.3794384334836447 +2458526.4185919776 33.52160585136826 1.3411034225878844 +2458527.576654805 34.69485667256363 2.8698173476737803 +2458528.111970773 35.221995595788 2.364704841833626 +2458529.8390300074 36.83353391584241 0.9249374796599684 +2458530.604024408 37.49811380269765 0.5797835892283674 +2458531.771079307 38.45001802512342 1.6845389508832564 +2458531.9641530444 38.60023174103681 0.6675655129694912 +2458532.6917056805 39.14787066442538 1.4171901284667194 +2458534.0926596513 40.12210351471556 0.1863187081001272 +2458534.149900518 40.15971817043575 0.5881775840787171 +2458537.076265389 41.8682604566477 0.5313839173251673 +2458537.191489296 41.92736839651852 0.53330536067345 +2458537.8113832935 42.23542071367206 0.6015267565751019 2458537.9657686786 42.309574416528875 0.833653309215884 -2458538.381849783 42.504433875411834 0.809332893418062 -2458540.743097261 43.47906010425952 2.684593000883861 -2458540.779631013 43.49246641231207 0.21346298137036643 -2458540.9921093797 43.56946644358569 0.1643903127786134 +2458538.381849783 42.504433875421874 0.809332893418062 +2458540.743097261 43.479060104254174 2.684593000883861 +2458540.779631013 43.492466412306726 0.21346298137036643 +2458540.9921093797 43.56946644358828 0.1643903127786134 2458541.5790782217 43.77367420369404 1.242025039313169 -2458542.337274912 44.01935711600308 0.23185696592260163 -2458544.2187311375 44.54424679872594 0.2634840517494057 -2458544.2979028192 44.56375854819504 0.1858451223697014 -2458546.274845702 44.9855696995223 0.520345394246201 -2458549.374411684 45.39837994750941 0.41137440309835466 -2458549.814878059 45.432428152190724 2.9023568719989417 -2458552.4403351406 45.50430851610058 1.0298554929276607 -2458552.8901611604 45.49337237379266 0.7538916413235267 -2458553.2193882978 45.480890484811624 0.7345379013621794 -2458554.0071165985 45.43533796137429 0.2676117067532668 -2458558.3401405653 44.752328167258334 2.1417909484327398 -2458562.891887693 43.08684507683734 0.1494630022235265 +2458542.337274912 44.01935711601211 0.23185696592260163 +2458544.2187311375 44.54424679873317 0.2634840517494057 +2458544.2979028192 44.56375854820217 0.1858451223697014 +2458546.274845702 44.98556969952759 0.520345394246201 +2458549.374411684 45.39837994751127 0.41137440309835466 +2458549.814878059 45.432428152189715 2.9023568719989417 +2458552.4403351406 45.50430851610034 1.0298554929276607 +2458552.8901611604 45.49337237379289 0.7538916413235267 +2458553.2193882978 45.48089048481036 0.7345379013621794 +2458554.0071165985 45.43533796137379 0.2676117067532668 +2458558.3401405653 44.75232816725106 2.1417909484327398 +2458562.891887693 43.08684507684092 0.1494630022235265 2458563.541054481 42.755027316771574 1.2686219946097248 -2458564.9075494427 41.97255918085013 0.15346859084547762 -2458564.928980646 41.95937025950825 0.5592235852758531 -2458565.095073066 41.85619404717419 0.21115268862687603 -2458565.688085095 41.47396967389899 0.6772268687009975 -2458566.0026242025 41.26253370014487 0.19696542931955294 -2458566.6836572993 40.78457413992845 0.567941231847918 +2458564.9075494427 41.9725591808546 0.15346859084547762 +2458564.928980646 41.959370259512724 0.5592235852758531 +2458565.095073066 41.85619404716051 0.21115268862687603 +2458565.688085095 41.4739696739038 0.6772268687009975 +2458566.0026242025 41.26253370013991 0.19696542931955294 +2458566.6836572993 40.78457413993894 0.567941231847918 2458566.9534362047 40.58784464709728 0.11134317071026828 -2458567.2519966373 40.36544076499406 0.23450551289578472 -2458568.8843326373 39.07308678933066 1.2679214911098027 -2458569.1489504385 38.853752776187555 0.601913981921152 -2458572.329534809 36.21173926251653 0.5529717086305411 +2458567.2519966373 40.36544076496666 0.23450551289578472 +2458568.8843326373 39.07308678933667 1.2679214911098027 +2458569.1489504385 38.85375277616936 0.601913981921152 +2458572.329534809 36.21173926249404 0.5529717086305411 2458572.570175387 36.0284082814722 0.9010168031111891 2458574.529938415 34.79182836715995 0.6135955943459428 -2458576.3368066563 34.26809465871612 0.24808496921970213 -2458577.629907819 34.370195949588265 2.162504002612899 -2458579.7197366427 35.38493314992718 0.29494605843239186 -2458581.5050042775 36.90217225559072 0.10649917303969393 -2458581.645047653 37.03853902271546 0.19242158077910812 -2458583.464109813 38.93834279115301 0.41077765230150565 -2458584.8743430697 40.479319576135694 0.10342867478137877 -2458585.028750741 40.64749746519492 0.6585010412587005 -2458586.284440072 41.993416335807176 2.691003276398369 -2458588.510278856 44.222327918424305 0.6911252459055384 -2458588.8714121594 44.56018747388936 0.20445643434195737 -2458590.686624103 46.14884315912025 0.12481417193686359 -2458591.2123410483 46.57461744629048 0.646928264384951 -2458593.324573916 48.13454310731777 1.0670210343371629 +2458576.3368066563 34.268094658714325 0.24808496921970213 +2458577.629907819 34.37019594958477 2.162504002612899 +2458579.7197366427 35.384933149916904 0.29494605843239186 +2458581.5050042775 36.90217225559773 0.10649917303969393 +2458581.645047653 37.03853902270831 0.19242158077910812 +2458583.464109813 38.93834279116882 0.41077765230150565 +2458584.8743430697 40.47931957611981 0.10342867478137877 +2458585.028750741 40.647497465210755 0.6585010412587005 +2458586.284440072 41.99341633583779 2.691003276398369 +2458588.510278856 44.222327918431176 0.6911252459055384 +2458588.8714121594 44.56018747388262 0.20445643434195737 +2458590.686624103 46.14884315910825 0.12481417193686359 +2458591.2123410488 46.57461744631364 0.646928264384951 +2458593.324573916 48.13454310733766 1.0670210343371629 2458594.55418422 48.93694581970476 0.9169355360938037 -2458595.2096541026 49.334857301786776 0.7099480541576314 -2458597.9311756245 50.78203504309509 0.35633010229856604 -2458599.21105146 51.35708934356261 0.3824200103253299 -2458599.3629127964 51.421133866999334 1.9660505419962635 -2458599.3691206276 51.42373340015031 0.11817040965459912 -2458601.3267903854 52.17307408836693 0.2806982788181881 -2458601.820179267 52.340439245700516 0.2696919904679525 -2458602.378518043 52.51974234668013 2.965305715722433 -2458602.418392197 52.53214159202897 0.5657940772185184 -2458603.1252326258 52.74305977310536 0.1869923501466016 -2458603.6775187505 52.89627218476699 1.8704068765384905 -2458604.7535473574 53.16599684270297 0.3635965516504767 -2458605.4055640083 53.31111057998035 0.37806094277072727 -2458605.5064679473 53.33233887721271 0.32043128702765633 -2458606.804172129 53.576051853114386 0.1680286493949399 +2458595.2096541026 49.33485730180399 0.7099480541576314 +2458597.9311756245 50.782035043091646 0.35633010229856604 +2458599.21105146 51.35708934357496 0.3824200103253299 +2458599.362912797 51.42113386701152 1.9660505419962635 +2458599.3691206276 51.42373340015946 0.11817040965459912 +2458601.3267903854 52.17307408837451 0.2806982788181881 +2458601.820179267 52.34043924569571 0.2696919904679525 +2458602.378518043 52.51974234668921 2.965305715722433 +2458602.418392197 52.53214159203575 0.5657940772185184 +2458603.1252326258 52.74305977311161 0.1869923501466016 +2458603.6775187505 52.89627218476309 1.8704068765384905 +2458604.7535473574 53.16599684269958 0.3635965516504767 +2458605.4055640083 53.31111057998499 0.37806094277072727 +2458605.5064679473 53.33233887721423 0.32043128702765633 +2458606.804172129 53.57605185311195 0.1680286493949399 2458610.9708635164 53.985206282534634 0.4676711206528206 -2458611.442485392 53.99424778316889 0.6665269648369389 -2458612.772249151 53.97628827260972 0.4778097074368812 -2458615.2129865023 53.76552059290973 0.37322387857426087 -2458616.1065599713 53.626105654868915 0.22107594552171897 -2458616.1234581396 53.62312762929799 1.941127954493647 -2458616.717265221 53.510191325899406 1.9703697920300995 -2458618.124941107 53.175457446871086 0.6983562943302118 -2458618.1636914215 53.16485527862642 1.1621134834401563 -2458618.8948865323 52.95025277461468 2.3252508438342883 -2458621.505436317 51.943865578790735 0.1740553827604509 -2458622.7788545527 51.30623939361758 0.15524876274805874 -2458623.504578622 50.89754943211297 0.19838662433170084 -2458623.8274050197 50.70508491852116 0.7710030355925556 -2458624.726844928 50.13447912126128 0.12882916182707782 -2458624.747714525 50.12064657846979 1.2610812653977455 -2458625.0811663736 49.8960605879743 0.11881032415917035 -2458625.727949624 49.44177073980637 0.10926440925322338 -2458625.750768838 49.42530745013968 2.814774096334648 -2458626.2666229447 49.045618658193035 0.3387750144223217 -2458626.8865343016 48.571589907532285 0.6903676490827431 -2458628.7452242393 47.06665742466761 0.5301252929504008 -2458630.565770755 45.58188595791667 0.9349881368331905 -2458632.386068454 44.30407143027982 2.5897471729215384 +2458611.442485392 53.99424778316905 0.6665269648369389 +2458612.772249151 53.976288272610276 0.4778097074368812 +2458615.2129865023 53.76552059290574 0.37322387857426087 +2458616.1065599713 53.62610565486507 0.22107594552171897 +2458616.1234581396 53.62312762929413 1.941127954493647 +2458616.7172652204 53.51019132590237 1.9703697920300995 +2458618.124941107 53.17545744686515 0.6983562943302118 +2458618.1636914215 53.164855278618425 1.1621134834401563 +2458618.8948865323 52.95025277461696 2.3252508438342883 +2458621.505436317 51.94386557878739 0.1740553827604509 +2458622.778854552 51.30623939362937 0.15524876274805874 +2458623.504578622 50.89754943210869 0.19838662433170084 +2458623.8274050197 50.70508491852998 0.7710030355925556 +2458624.726844928 50.13447912127092 0.12882916182707782 +2458624.747714525 50.12064657847945 1.2610812653977455 +2458625.0811663736 49.896060587964335 0.11881032415917035 +2458625.727949624 49.44177073981686 0.10926440925322338 +2458625.750768838 49.42530745015015 2.814774096334648 +2458626.2666229447 49.04561865817122 0.3387750144223217 +2458626.8865343016 48.571589907543625 0.6903676490827431 +2458628.7452242393 47.06665742467966 0.5301252929504008 +2458630.565770755 45.58188595792236 0.9349881368331905 +2458632.386068454 44.30407143026684 2.5897471729215384 2458632.5898929387 44.18597899531268 0.6849572956422009 -2458632.885712266 44.02634039796506 0.30478015513583656 -2458633.3107718504 43.82356390131066 0.11500762846362844 -2458636.204892672 43.49047824803964 0.9274511988886456 -2458636.5858795308 43.59827891155638 0.42734431889235736 -2458638.3316617925 44.52043088138531 2.7105359603686727 -2458639.529264593 45.49108345741045 0.754754453749501 -2458640.26416489 46.184199735816435 0.8934213550153699 +2458632.885712266 44.026340397968795 0.30478015513583656 +2458633.310771851 43.82356390129793 0.11500762846362844 +2458636.204892672 43.490478248046536 0.9274511988886456 +2458636.5858795308 43.59827891155396 0.42734431889235736 +2458638.3316617925 44.520430881405986 2.7105359603686727 +2458639.529264593 45.491083457417005 0.754754453749501 +2458640.26416489 46.18419973584501 0.8934213550153699 2458640.5170123368 46.435510308659644 0.9163919100922346 -2458642.1516529405 48.15979077787583 0.6523657807975388 -2458642.9835628923 49.06834622883131 0.381948429960075 -2458644.656136076 50.87447791634247 0.49149827839621346 -2458645.860685533 52.11786921699642 1.0575560634685441 -2458646.765329834 53.00606430928372 0.22817179377296543 -2458647.0520731797 53.27850384103365 0.18784970866476805 -2458647.311964013 53.52150228405162 2.233010518546652 -2458647.7607401083 53.93216248722787 1.0573472609219101 -2458648.8550824937 54.885216115366504 2.0274456541043606 -2458650.24250251 55.99458035900717 0.3556114120365529 -2458653.142561982 57.97133225639712 0.7308032542442472 -2458654.09818496 58.528359311427174 0.8062206052435784 -2458655.553645495 59.29433748673622 2.197499480694631 -2458655.8205519086 59.42448276922298 2.070848876973854 -2458657.3739412334 60.12170482369095 0.24569986602104002 -2458658.632426726 60.61460047558608 2.836896735596055 -2458659.090205151 60.77862967680762 0.3616984237835981 -2458661.6693102145 61.557964310980935 0.8213046317254442 -2458661.970445023 61.633422588266995 3.089209874901623 -2458664.806589622 62.19031116830659 0.40810137954600356 -2458665.6271729697 62.3001041648131 0.9179403453049487 -2458666.7130005476 62.40992316530998 3.0370889252899156 -2458667.0945009952 62.43884519066549 0.8788990418287487 -2458668.414699187 62.49957042609752 0.9692333756333517 -2458669.439618293 62.50373380061642 1.5213917499704739 -2458669.706668256 62.49850094843086 1.520076470108819 -2458671.4065454635 62.40185632557086 2.642550444155089 -2458672.388480312 62.29416197758076 0.9405234129234002 -2458672.603266781 62.265325315562556 0.23801002767908566 -2458673.117043802 62.188446362586255 0.12352936959714454 -2458676.503348102 61.38189904581722 0.13169739445687073 -2458676.5128372447 61.37885303552596 2.7762049656096166 -2458676.551107868 61.36652066849223 1.7177511235914709 -2458677.4965149644 61.03723519447482 0.7763630302735172 -2458678.730654603 60.53311681490936 2.2385006738165316 -2458678.8308989312 60.488322460899624 2.0192901243554884 -2458679.4450754654 60.200839736317285 0.3462297256643072 -2458679.6154656536 60.1170566798854 1.63748735794655 -2458681.369863755 59.14914970109586 0.19887485401319488 -2458681.8360715983 58.858897514108506 2.9905916267298047 -2458682.2526453803 58.58763254910595 0.2979743145364124 +2458642.1516529405 48.15979077789956 0.6523657807975388 +2458642.9835628923 49.06834622883927 0.381948429960075 +2458644.656136076 50.8744779163271 0.49149827839621346 +2458645.860685533 52.11786921698911 1.0575560634685441 +2458646.765329834 53.00606430926976 0.22817179377296543 +2458647.0520731797 53.27850384104737 0.18784970866476805 +2458647.311964013 53.52150228407862 2.233010518546652 +2458647.7607401083 53.932162487214725 1.0573472609219101 +2458648.8550824937 54.88521611535428 2.0274456541043606 +2458650.24250251 55.99458035902931 0.3556114120365529 +2458653.142561982 57.97133225641478 0.7308032542442472 +2458654.09818496 58.5283593114394 0.8062206052435784 +2458655.5536454944 59.29433748673622 2.197499480694631 +2458655.8205519086 59.42448276921596 2.070848876973854 +2458657.3739412334 60.12170482370309 0.24569986602104002 +2458658.632426726 60.61460047558075 2.836896735596055 +2458659.090205151 60.77862967681524 0.3616984237835981 +2458661.669310214 61.55796431097721 0.8213046317254442 +2458661.970445023 61.633422588268786 3.089209874901623 +2458664.806589622 62.19031116830443 0.40810137954600356 +2458665.6271729697 62.30010416481221 0.9179403453049487 +2458666.7130005476 62.40992316530879 3.0370889252899156 +2458667.0945009952 62.43884519066649 0.8788990418287487 +2458668.414699187 62.49957042609802 0.9692333756333517 +2458669.439618293 62.5037338006162 1.5213917499704739 +2458669.706668256 62.49850094843122 1.520076470108819 +2458671.4065454635 62.401856325569554 2.642550444155089 +2458672.388480312 62.29416197757793 0.9405234129234002 +2458672.603266781 62.26532531556356 0.23801002767908566 +2458673.117043802 62.18844636258275 0.12352936959714454 +2458676.503348102 61.381899045814876 0.13169739445687073 +2458676.5128372447 61.37885303552361 2.7762049656096166 +2458676.551107868 61.36652066849459 1.7177511235914709 +2458677.4965149644 61.03723519447211 0.7763630302735172 +2458678.730654603 60.53311681491909 2.2385006738165316 +2458678.8308989312 60.488322460906176 2.0192901243554884 +2458679.4450754654 60.20083973631017 0.3462297256643072 +2458679.6154656536 60.117056679892606 1.63748735794655 +2458681.369863755 59.14914970108262 0.19887485401319488 +2458681.8360715983 58.85889751411317 2.9905916267298047 +2458682.2526453803 58.587632549086564 0.2979743145364124 2458682.566988757 58.3754964990636 0.26265516981018255 -2458683.389570199 57.790356379750264 0.1877042235803976 -2458684.0216642665 57.31190740204639 0.68079594428688 -2458684.6012278446 56.85229414456144 2.2490926077881914 -2458685.404821025 56.184678413378805 0.6782833719614333 -2458686.1590582514 55.530611951302355 0.11441318503569375 -2458687.236644577 54.56401285550905 0.1202290837878777 -2458688.3125864435 53.587231196940515 0.5989857671825178 -2458688.7479626173 53.19727469850335 2.1622996448879914 -2458692.156780839 50.70045698262314 0.7499317643580609 -2458692.766937284 50.44332867927932 1.317165408773449 +2458683.389570199 57.79035637973418 0.1877042235803976 +2458684.0216642665 57.31190740203514 0.68079594428688 +2458684.6012278446 56.85229414456734 2.2490926077881914 +2458685.404821025 56.18467841336024 0.6782833719614333 +2458686.1590582514 55.530611951276704 0.11441318503569375 +2458687.236644577 54.56401285548261 0.1202290837878777 +2458688.3125864435 53.58723119691425 0.5989857671825178 +2458688.7479626173 53.19727469852279 2.1622996448879914 +2458692.156780839 50.70045698260901 0.7499317643580609 +2458692.766937284 50.44332867928449 1.317165408773449 2458694.578975429 50.18504924873144 1.7992982737596763 -2458698.192402917 51.92237393224195 0.15750622213241686 -2458700.739077453 54.26025074051744 0.43997779817577526 -2458702.640045006 56.13536285625857 0.10132625189275672 -2458702.750355968 56.242530848074594 0.4248660289420663 -2458703.763920462 57.20776000047498 1.939474694854101 +2458698.192402917 51.92237393226535 0.15750622213241686 +2458700.739077453 54.26025074050314 0.43997779817577526 +2458702.640045006 56.13536285624437 0.10132625189275672 +2458702.750355968 56.24253084806052 0.4248660289420663 +2458703.763920462 57.20776000046136 1.939474694854101 2458704.937250477 58.269067900188396 1.517756693761716 -2458706.522387948 59.58589983442129 0.3175037316942859 -2458707.07824788 60.01332272408253 1.1730716922155056 -2458707.2687107017 60.1555916801723 1.171679106892128 -2458707.4134367574 60.26226656124572 0.5459351873809183 -2458709.6955002095 61.782907550453785 0.14315469172410467 -2458709.8298823065 61.86318531759164 2.4505431089785272 -2458709.9109045127 61.911102421093005 0.7225544906979365 -2458712.4037619983 63.214161050382174 0.1851887531382049 -2458714.0357477074 63.8991038083093 1.0253505365488276 -2458714.1081865425 63.92663046788159 1.4340750097726003 +2458706.522387948 59.58589983442704 0.3175037316942859 +2458707.07824788 60.01332272409348 1.1730716922155056 +2458707.2687107017 60.1555916801993 1.171679106892128 +2458707.4134367574 60.262266561256396 0.5459351873809183 +2458709.6955002095 61.78290755044503 0.14315469172410467 +2458709.8298823065 61.863185317587316 2.4505431089785272 +2458709.9109045127 61.911102421088714 0.7225544906979365 +2458712.4037619983 63.21416105039218 0.1851887531382049 +2458714.0357477074 63.899103808314834 1.0253505365488276 +2458714.1081865425 63.926630467889844 1.4340750097726003 2458714.9326952337 64.22341616881728 0.14695182580066382 -2458715.1750054606 64.30496375881876 0.4339947154843559 -2458715.460098795 64.39767882895583 2.6575457196324073 -2458715.874209366 64.52621794575656 0.7604611975490089 -2458716.6701493477 64.75331501122406 1.5039738599352945 -2458717.8839150737 65.05072306508012 0.16872360448870846 +2458715.1750054606 64.30496375882839 0.4339947154843559 +2458715.460098795 64.3976788289605 2.6575457196324073 +2458715.874209366 64.52621794575438 0.7604611975490089 +2458716.6701493477 64.75331501122014 1.5039738599352945 +2458717.8839150737 65.0507230650785 0.16872360448870846 2458719.5715264576 65.37038484552342 0.251555401818196 -2458721.404752437 65.5994703866685 2.7841718888755898 -2458722.1256263372 65.65683890228772 0.1321843407272699 -2458722.914182628 65.69875041480196 2.1935827170715214 +2458721.404752437 65.59947038666985 2.7841718888755898 +2458722.1256263372 65.65683890228918 0.1321843407272699 +2458722.914182628 65.69875041480168 2.1935827170715214 2458723.9346860894 65.72081787794625 0.12454863070390058 -2458727.3881316422 65.52413328807874 0.7791749863733198 -2458727.7514479174 65.47853044347066 0.10488140473186187 -2458728.7121236105 65.3342292513253 2.2568308820215646 -2458729.6546154763 65.15837832871469 0.6339568803573913 -2458730.4394241096 64.98517066444848 0.6397449242511021 -2458730.8583014337 64.88246843541461 0.3826116012403758 +2458727.3881316422 65.52413328807614 0.7791749863733198 +2458727.7514479174 65.47853044347258 0.10488140473186187 +2458728.7121236105 65.33422925132774 2.2568308820215646 +2458729.6546154763 65.15837832871618 0.6339568803573913 +2458730.4394241096 64.98517066444504 0.6397449242511021 +2458730.8583014337 64.8824684354183 0.3826116012403758 2458730.9040930658 64.87080027769514 1.1115237695142635 -2458731.3492399245 64.75278225760663 1.0016359084104673 -2458731.4729332016 64.71849633603708 0.7989482887516248 -2458731.4767028 64.7174411772657 0.49235387654801294 -2458732.3724768804 64.44916633236603 0.19242435066240274 -2458732.4368653228 64.42851537410866 0.740061206957491 -2458733.0372002055 64.22690834006401 2.5619480682318465 -2458735.0447353334 63.42790019559681 0.28093673072808095 -2458735.1149851214 63.39628493466235 0.49058533955576455 +2458731.3492399245 64.75278225759864 1.0016359084104673 +2458731.4729332016 64.718496336033 0.7989482887516248 +2458731.4767028 64.71744117726163 0.49235387654801294 +2458732.3724768804 64.44916633235908 0.19242435066240274 +2458732.4368653228 64.42851537410397 0.740061206957491 +2458733.0372002055 64.22690834005893 2.5619480682318465 +2458735.0447353334 63.42790019559027 0.28093673072808095 +2458735.1149851214 63.396284934649195 0.49058533955576455 2458736.9458670467 62.47858714381885 0.12199658149974707 -2458737.321895205 62.266823170869486 0.13426204604656036 -2458744.1668420266 56.88022990908795 0.1991438233507365 -2458745.6398492367 55.382570009150236 0.8815480413816492 -2458746.5032714573 54.483907215009694 0.2562568428523935 -2458747.2674526228 53.693273319639076 2.6590371437428058 -2458748.8496229746 52.144735609280524 0.4210195419775625 -2458749.366416382 51.68773471530906 0.8636344673643035 -2458751.1829068772 50.40580243438146 0.7051134936322967 -2458752.3958529127 49.92465633101922 0.5565755855433756 -2458753.8998918934 49.824280997074204 0.3976748337679327 -2458754.6547187455 49.9805998026777 0.4668878931605293 -2458758.343388597 52.21976070791676 3.1190703851192163 +2458737.321895205 62.266823170852795 0.13426204604656036 +2458744.1668420266 56.88022990906639 0.1991438233507365 +2458745.6398492367 55.382570009157796 0.8815480413816492 +2458746.5032714573 54.48390721500969 0.2562568428523935 +2458747.2674526228 53.693273319609254 2.6590371437428058 +2458748.8496229746 52.144735609287125 0.4210195419775625 +2458749.366416382 51.68773471529042 0.8636344673643035 +2458751.1829068772 50.405802434365896 0.7051134936322967 +2458752.3958529127 49.92465633101371 0.5565755855433756 +2458753.8998918934 49.82428099707334 0.3976748337679327 +2458754.6547187455 49.980599802675556 0.4668878931605293 +2458758.343388597 52.2197607079346 3.1190703851192163 2458759.917944817 53.55556269202489 0.1374402713875284 -2458761.831231671 55.18897000213629 2.8133214222625758 -2458768.464569113 59.56778428961978 0.2399620581758453 -2458768.5819420405 59.62305180134105 1.5742454045673095 +2458761.831231671 55.188970002124194 2.8133214222625758 +2458768.464569113 59.56778428962323 0.2399620581758453 +2458768.5819420405 59.62305180133763 1.5742454045673095 2458770.5460998467 60.43999007877256 0.222025120934236 -2458773.8488958 61.39278916733532 0.38650757698972943 +2458773.8488958 61.39278916733374 0.38650757698972943 2458775.5640330412 61.70228135725388 1.1716209888640634 -2458775.840023363 61.74120234540741 2.4363732030264633 +2458775.840023363 61.741202345405426 2.4363732030264633 2458775.953186657 61.756314526903985 0.4512904637746437 -2458777.14573308 61.88623730525252 0.8212354027300535 -2458778.786889266 61.980288848295636 0.3320340189700579 +2458777.14573308 61.88623730525506 0.8212354027300535 +2458778.786889266 61.980288848295224 0.3320340189700579 2458781.52885256 61.928556538235156 0.20614510231981653 -2458784.2688611285 61.62309389630521 2.5734677832368176 -2458784.2977270763 61.618524129110895 2.561962764459666 -2458784.6890452486 61.55376133191355 2.9110683687157026 -2458785.3168836124 61.43883907316712 2.933140518166558 -2458787.002880427 61.06169698446683 2.2614800578616503 -2458789.7109172232 60.23564902008364 0.6935105563455366 -2458790.334076077 60.00451648155351 0.29092972750062335 +2458784.2688611285 61.62309389629946 2.5734677832368176 +2458784.2977270763 61.61852412910628 2.561962764459666 +2458784.6890452486 61.55376133191606 2.9110683687157026 +2458785.3168836124 61.43883907316145 2.933140518166558 +2458787.002880427 61.06169698446499 2.2614800578616503 +2458789.7109172232 60.23564902008887 0.6935105563455366 +2458790.334076077 60.00451648154234 0.29092972750062335 2458790.966275479 59.753344499423214 0.38669416300752857 -2458792.346680061 59.143740760549534 0.2747534829925239 -2458792.9836372067 58.83275639562224 0.2995300205518705 +2458792.346680061 59.14374076053919 0.2747534829925239 +2458792.9836372067 58.83275639561857 0.2995300205518705 diff --git a/docs/examples/data.ecsv b/docs/examples/data.ecsv index 0125995d..e9289bd7 100644 --- a/docs/examples/data.ecsv +++ b/docs/examples/data.ecsv @@ -1,11 +1,11 @@ -# %ECSV 0.9 +# %ECSV 1.0 # --- # datatype: # - {name: bjd, datatype: float64} # - {name: rv, unit: km / s, datatype: float64} # - {name: rv_err, unit: km / s, datatype: float64} # meta: !!omap -# - t_ref: !astropy.time.Time {format: mjd, in_subfmt: '*', jd1: 2458512.0, jd2: -0.20488644755096175, out_subfmt: '*', precision: 3, +# - t_ref: !astropy.time.Time {format: mjd, in_subfmt: '*', jd1: 2458512.0, jd2: -0.20488644757278962, out_subfmt: '*', precision: 3, # scale: tcb} # - __serialized_columns__: # rv: @@ -19,259 +19,259 @@ # schema: astropy-2.0 bjd rv rv_err 2458511.7951135524 37.63538635612137 0.22341820389458492 -2458512.1973919393 37.794388800166274 0.48399198155187484 -2458512.5399584738 37.9124751984127 0.1671605217360056 -2458512.8695262037 38.01019743708635 0.11470700018223313 +2458512.1973919393 37.79438880018236 0.48399198155187484 +2458512.5399584738 37.91247519841736 0.1671605217360056 +2458512.8695262037 38.01019743708833 0.11470700018223313 2458513.6255345265 38.17142772241421 0.3371802798016645 2458513.6787027447 38.1793017494875 1.338474665902923 -2458513.7290753396 38.186327076256646 2.1551502385179004 -2458514.6069552493 38.23842353393525 0.22387028602056924 -2458515.2999124355 38.180764019786515 2.608888611051979 -2458515.374372337 38.169159582948865 0.13147929534874903 -2458515.626131753 38.12199047194763 0.13105006526306565 -2458515.9115815205 38.05355539126889 1.1430648800333527 -2458516.1373007107 37.98807961580226 0.12638754213716372 +2458513.7290753396 38.18632707625566 2.1551502385179004 +2458514.6069552493 38.238423533935105 0.22387028602056924 +2458515.299912436 38.18076401977894 2.608888611051979 +2458515.374372337 38.169159582941745 0.13147929534874903 +2458515.6261317525 38.12199047194763 0.13105006526306565 +2458515.9115815205 38.05355539126499 1.1430648800333527 +2458516.1373007107 37.988079615790895 0.12638754213716372 2458516.6617032117 37.79690986997475 2.8574003324529063 -2458519.904824091 35.45113391641817 0.257818878799462 -2458520.095701421 35.25828342333601 0.15598603609819855 +2458519.904824091 35.451133916410924 0.257818878799462 +2458520.095701421 35.25828342329879 0.15598603609819855 2458520.6915558754 34.62656963313482 1.4212867171701629 2458520.8352296217 34.46818064751871 0.4302058307838804 -2458521.6052562613 33.58674837244351 2.2393041165074066 -2458521.88827662 33.25171204066023 0.5478829820558546 -2458521.9608255355 33.16509837166233 0.10279581688342261 -2458522.3325251793 32.71761710910887 0.1757983055626769 -2458522.979423407 31.929742318746257 0.14095428238633567 -2458525.212076863 29.282912494390168 0.18707865561107262 -2458525.4250065605 29.04764793248382 0.2545963976014491 +2458521.6052562613 33.58674837243496 2.2393041165074066 +2458521.88827662 33.251712040651555 0.5478829820558546 +2458521.9608255355 33.16509837164493 0.10279581688342261 +2458522.3325251793 32.71761710906482 0.1757983055626769 +2458522.979423407 31.929742318719594 0.14095428238633567 +2458525.212076863 29.282912494333385 0.18707865561107262 +2458525.4250065605 29.047647932443994 0.2545963976014491 2458526.726917548 27.716276983566384 0.1297381350371654 -2458527.9426433467 26.67184008614556 2.0815560120336163 -2458528.0039909477 26.624804823998183 1.0231497648292653 +2458527.9426433467 26.67184008613434 2.0815560120336163 +2458528.0039909477 26.62480482398155 1.0231497648292653 2458528.8550834605 26.031260898832237 0.46341063581672076 -2458529.1186102657 25.870124423702947 2.2287942207115714 -2458529.205223191 25.819524297942948 0.8269925924539391 -2458529.8809215585 25.464940800137423 1.053374187966504 -2458530.090803717 25.36927180653545 1.1721099226570535 -2458530.134139603 25.350369165926114 0.4562427566440821 -2458530.134691676 25.350130233090983 0.27526206257462377 +2458529.1186102657 25.870124423681446 2.2287942207115714 +2458529.205223191 25.819524297917738 0.8269925924539391 +2458529.8809215585 25.464940800133988 1.053374187966504 +2458530.090803717 25.36927180652266 1.1721099226570535 +2458530.134139603 25.350369165907217 0.4562427566440821 +2458530.134691676 25.35013023307524 0.27526206257462377 2458530.644092005 25.149679158674587 2.1386409008999814 -2458530.9584806263 25.045749974977237 0.35779108981157376 -2458531.0544676706 25.01699431463352 0.4520550475014811 -2458531.1849017097 24.98013262014708 0.25966306391147265 -2458532.133611955 24.78725294013038 0.676810195445213 -2458532.3559681317 24.76071144460364 0.6415262738208175 -2458532.7750877053 24.7293018511864 0.4442226291335937 -2458533.7523140265 24.74684825619526 0.3387800439645821 -2458534.179651081 24.792524602860396 0.6495530357097828 -2458534.4472189583 24.83233186692637 0.37418546430662336 -2458534.4616757273 24.834723108557945 1.4745125856500165 -2458534.57294269 24.853942179915343 0.16987063691333779 -2458535.1062234063 24.96565728354451 1.2041835223899529 -2458535.364269518 25.031018770915146 0.20559001389326279 -2458537.9858829523 26.066212592808842 2.500559542743992 -2458538.331549126 26.247317169517597 0.3779963174623941 -2458538.668001884 26.4322941138862 0.3059639246186404 -2458538.7130355756 26.45768390333857 2.246283118785033 +2458530.9584806263 25.04574997497054 0.35779108981157376 +2458531.0544676706 25.01699431462501 0.4520550475014811 +2458531.1849017097 24.98013262013517 0.25966306391147265 +2458532.133611955 24.787252940125466 0.676810195445213 +2458532.3559681317 24.760711444599863 0.6415262738208175 +2458532.7750877053 24.729301851186733 0.4442226291335937 +2458533.7523140265 24.746848256194674 0.3387800439645821 +2458534.179651081 24.792524602866195 0.6495530357097828 +2458534.4472189583 24.83233186693116 0.37418546430662336 +2458534.4616757273 24.834723108562784 1.4745125856500165 +2458534.5729426895 24.853942179916647 0.16987063691333779 +2458535.1062234063 24.965657283553213 1.2041835223899529 +2458535.364269518 25.03101877092487 0.20559001389326279 +2458537.9858829523 26.06621259281998 2.500559542743992 +2458538.331549126 26.247317169541052 0.3779963174623941 +2458538.668001884 26.432294113886204 0.3059639246186404 +2458538.713035575 26.457683903338573 2.246283118785033 2458538.82726244 26.522738038596252 0.11777799029988209 -2458539.3811393795 26.85105444859142 2.927199921278786 -2458539.415960697 26.872385581204686 2.4718976014731795 -2458539.4911145433 26.9186937314235 0.9929856568136601 -2458540.716614934 27.722574573345764 0.7151567282886854 -2458540.723064148 27.727032230793686 0.6116870301936291 +2458539.3811393795 26.851054448613667 2.927199921278786 +2458539.415960697 26.872385581222552 2.4718976014731795 +2458539.4911145433 26.918693731436996 0.9929856568136601 +2458540.716614934 27.722574573340736 0.7151567282886854 +2458540.723064148 27.72703223079368 0.6116870301936291 2458540.8206339017 27.794741797905086 0.13339042359994682 2458540.8579934235 27.820800847243408 0.13682743854283835 -2458541.0772961485 27.975224608224295 1.0863447718461388 -2458541.121880845 28.006918244987197 2.7349633825095236 -2458541.880405759 28.560682744215686 0.5826046016920249 -2458542.0969222724 28.723478125698442 0.10035237920729409 -2458542.6187597695 29.123613881997606 1.7509062458140376 -2458542.7214284437 29.203552250190114 1.76112684962545 -2458542.772975497 29.243829567883832 0.3197185504113186 -2458544.366087825 30.52871243664606 1.0464728059481498 -2458544.902084453 30.974626244867988 0.1630705266305978 -2458545.243083203 31.26070878491114 0.17573447451801208 -2458545.301802088 31.310122846250522 1.167209479848952 -2458545.3339542365 31.33719678323974 0.5819569446018467 -2458545.4436923917 31.42968646757187 0.7645192138070683 +2458541.0772961485 27.975224608244947 1.0863447718461388 +2458541.121880845 28.006918245013097 2.7349633825095236 +2458541.880405759 28.560682744221122 0.5826046016920249 +2458542.0969222724 28.723478125725965 0.10035237920729409 +2458542.6187597695 29.123613882003255 1.7509062458140376 +2458542.7214284437 29.203552250190118 1.76112684962545 +2458542.772975497 29.243829567883836 0.3197185504113186 +2458544.366087825 30.52871243667614 1.0464728059481498 +2458544.902084453 30.974626244874063 0.1630705266305978 +2458545.243083203 31.260708784953987 0.17573447451801208 +2458545.301802088 31.31012284628728 1.167209479848952 +2458545.3339542365 31.337196783270386 0.5819569446018467 +2458545.4436923917 31.429686467602544 0.7645192138070683 2458545.6898913775 31.63760219294733 0.6103216468817618 -2458545.856581011 31.778634673531748 2.574550616215446 -2458546.1974915667 32.06750128805181 1.5241870572606315 -2458546.770303743 32.553135634552504 2.2319530840165 -2458546.9335481036 32.69135410175073 0.27379175845475445 -2458547.1001992957 32.83227809610074 2.9472428888211595 +2458545.856581011 31.778634673537915 2.574550616215446 +2458546.1974915667 32.06750128808882 1.5241870572606315 +2458546.7703037425 32.553135634552504 2.2319530840165 +2458546.9335481036 32.691354101756886 0.27379175845475445 +2458547.1001992957 32.83227809612532 2.9472428888211595 2458548.739866051 34.19709294966671 1.6910896872089258 -2458548.87038594 34.303057925709766 0.1410422831020223 -2458549.094597208 34.48379719752555 0.1351015739531875 -2458549.6615879387 34.932539951662115 0.19376777695228434 -2458550.0870243954 35.26008955520981 0.3195670255134891 -2458550.28416869 35.408816354037995 0.21338142093369603 -2458550.9575470984 35.90005764537306 0.11960799710854207 -2458551.3440377475 36.1687713434041 0.4650714656398719 -2458551.457620365 36.24570153386527 0.24760977756060598 +2458548.87038594 34.30305792571566 0.1410422831020223 +2458549.094597208 34.48379719755474 0.1351015739531875 +2458549.6615879387 34.932539951667785 0.19376777695228434 +2458550.0870243954 35.26008955523744 0.3195670255134891 +2458550.28416869 35.4088163540707 0.21338142093369603 +2458550.9575470984 35.90005764538337 0.11960799710854207 +2458551.3440377475 36.16877134343386 0.4650714656398719 +2458551.457620365 36.24570153388487 0.24760977756060598 2458551.668151446 36.38569140890009 0.49431626756089414 -2458551.8872192274 36.52760101157714 1.4970324442130165 -2458552.119522199 36.67367105141511 0.30384983729086273 -2458553.343705049 37.35824481743809 0.1149355776040928 -2458554.157271688 37.72078292302159 1.7856936732699211 -2458554.478539183 37.84053538034509 0.5438482807033748 -2458554.982222368 37.99915301027728 0.23239313128588251 -2458555.9511024794 38.195282416927185 2.468231957604805 -2458556.9639741285 38.230761504718 1.2670465735636296 -2458557.364910584 38.19308723122357 0.4290974625704362 -2458558.007952889 38.06836286246205 0.33165427674446596 +2458551.8872192274 36.52760101158179 1.4970324442130165 +2458552.119522199 36.67367105143762 0.30384983729086273 +2458553.343705049 37.358244817459685 0.1149355776040928 +2458554.157271688 37.720782923038804 1.7856936732699211 +2458554.478539183 37.84053538035275 0.5438482807033748 +2458554.982222368 37.99915301028133 0.23239313128588251 +2458555.9511024794 38.19528241692898 2.468231957604805 +2458556.9639741285 38.23076150471718 1.2670465735636296 +2458557.364910584 38.19308723121782 0.4290974625704362 +2458558.007952889 38.06836286245645 0.33165427674446596 2458558.6671129093 37.85618846670452 0.12781378122908615 2458558.670031597 37.85505695606132 0.2682121554224838 -2458559.9843785786 37.174003891584 1.1850703888319583 -2458561.21007865 36.24299633740116 0.12985309701614323 -2458561.5982932933 35.894015915425264 0.3670429816577177 -2458561.8922199663 35.61419897013813 0.7715379885696934 -2458562.44906581 35.05051189330296 1.2880561209959978 -2458563.1304501933 34.3091446370546 1.3008696034572902 -2458563.4649259015 33.92814944018391 0.27473362450584543 -2458563.5589009146 33.81941134802707 0.31179101270895127 -2458564.753893603 32.391450194749936 0.17219581821825913 -2458565.1345074004 31.927002014438074 1.2414572290814219 -2458565.460461512 31.529043638324556 0.2501522105202238 +2458559.9843785786 37.17400389156988 1.1850703888319583 +2458561.21007865 36.2429963373633 0.12985309701614323 +2458561.5982932933 35.894015915411735 0.3670429816577177 +2458561.8922199663 35.61419897013104 0.7715379885696934 +2458562.44906581 35.05051189327244 1.2880561209959978 +2458563.1304501933 34.309144637013716 1.3008696034572902 +2458563.4649259015 33.92814944015034 0.27473362450584543 +2458563.5589009146 33.819411348010185 0.31179101270895127 +2458564.753893603 32.39145019475879 0.17219581821825913 +2458565.1345074004 31.927002014393633 1.2414572290814219 +2458565.460461512 31.529043638289064 0.2501522105202238 2458565.8306855797 31.078960791318927 0.5255753785563327 -2458566.0213195374 30.84872284672518 2.2366022664001677 -2458566.1508215494 30.693115784011074 0.10224532121251691 -2458566.2781605255 30.540842970680185 1.9277033006421684 +2458566.0213195374 30.848722846698887 2.2366022664001677 +2458566.1508215494 30.69311578395874 0.10224532121251691 +2458566.2781605255 30.540842970619444 1.9277033006421684 2458566.8458072143 29.87316037462867 1.2875771420944915 -2458567.113176119 29.566376428668043 0.9283859768023894 -2458567.3598506446 29.288562675589226 0.4079583845132085 -2458567.6027899208 29.02037460709473 0.8630334769827732 +2458567.113176119 29.566376428626686 0.9283859768023894 +2458567.3598506446 29.288562675540515 0.4079583845132085 +2458567.6027899208 29.020374607086776 0.8630334769827732 2458569.686046281 27.000033601768962 2.647861896916208 2458569.730383666 26.963307719247123 0.1389710023048171 2458569.84148702 26.872534875042053 0.12637405395923684 -2458570.058896495 26.700159546248017 0.32414964704313437 -2458570.1874520415 26.60153914332664 0.14275286554365627 -2458570.2244934235 26.573582929154988 2.3803178384024384 +2458570.058896495 26.70015954622541 0.32414964704313437 +2458570.1874520415 26.60153914329357 0.14275286554365627 +2458570.2244934235 26.57358292912216 2.3803178384024384 2458570.7423330545 26.20455749959222 1.5172855479287055 2458570.7705419404 26.18563413148256 0.17283331504031468 -2458571.0172255393 26.0253897675659 0.3634464008050491 -2458571.022812415 26.02186970477842 0.41984302025602 -2458571.1485215756 25.943947130574863 2.7016828276438933 -2458571.373926399 25.81038398699036 0.31176663081514633 -2458573.208608696 25.016614280425003 0.5251855657916831 -2458573.555412109 24.924257313363764 0.6215030487050776 +2458571.0172255393 26.025389767547555 0.3634464008050491 +2458571.022812415 26.021869704760096 0.41984302025602 +2458571.1485215756 25.94394713055266 2.7016828276438933 +2458571.373926399 25.810383986969455 0.31176663081514633 +2458573.208608696 25.01661428041224 0.5251855657916831 +2458573.555412109 24.92425731336026 0.6215030487050776 2458573.6363293286 24.90527165051254 1.4883494469612204 2458573.75827769 24.87847375723161 2.0047811551285584 -2458574.607295453 24.751160240830533 1.0142891001670458 -2458574.7447718913 24.740067211333844 0.3507468888371042 -2458574.9594708732 24.72790675073891 0.22058145880207508 -2458575.5202305717 24.725169226133694 0.39546979403522764 +2458574.607295453 24.751160240829876 1.0142891001670458 +2458574.7447718913 24.74006721133384 0.3507468888371042 +2458574.959470873 24.7279067507383 0.22058145880207508 +2458575.5202305717 24.725169226134156 0.39546979403522764 2458575.76684139 24.736899980128484 0.7307294263789043 -2458576.3157991534 24.79032360604769 1.413784344608477 -2458576.664600299 24.84319593580483 2.7494480646949486 -2458577.321499425 24.980809077760167 2.042224766757052 -2458577.44164857 25.011160241436485 1.5193707483503014 +2458576.3157991534 24.7903236060534 1.413784344608477 +2458576.664600299 24.843195935806083 2.7494480646949486 +2458577.321499425 24.980809077770907 2.042224766757052 +2458577.44164857 25.011160241444028 1.5193707483503014 2458577.798057939 25.110264845770526 2.243828377115662 -2458578.4440777106 25.323156836095507 1.10681618645793 -2458580.1419169893 26.067842975010283 2.61960217421659 -2458580.320438804 26.160264704751942 0.5892813631176571 -2458580.5231411536 26.26819201112164 0.13696624203245664 +2458578.4440777106 25.32315683610603 1.10681618645793 +2458580.1419169893 26.067842975028857 2.61960217421659 +2458580.320438804 26.16026470477485 0.5892813631176571 +2458580.5231411536 26.26819201113343 0.13696624203245664 2458580.6790191904 26.353299575806943 1.242461214154429 -2458581.458338464 26.804992399562195 2.646773708190422 -2458582.517002403 27.482411928181474 0.1113167214457624 +2458581.458338464 26.804992399579852 2.646773708190422 +2458582.517002403 27.48241192819613 0.1113167214457624 2458582.8236835045 27.69100122902624 0.8573125724736741 -2458583.3819039958 28.083514698197 0.28651856055448477 -2458584.619173907 29.005648535844447 0.40389986354060536 -2458584.9887659927 29.293143484541055 2.383132770434437 -2458585.2525704843 29.50127279252495 0.31738345242813737 -2458585.485611641 29.687006817798178 0.8683992717746128 -2458585.550582178 29.739085634380853 0.4316165838873574 -2458588.953128109 32.57853331879285 2.4561121534316137 -2458589.2688752552 32.84565731418336 0.3572845524654675 -2458589.5520900493 33.08453313437572 0.14095138991390269 +2458583.3819039958 28.083514698223087 0.28651856055448477 +2458584.619173907 29.005648535850067 0.40389986354060536 +2458584.9887659927 29.29314348455819 2.383132770434437 +2458585.2525704843 29.501272792559593 0.31738345242813737 +2458585.485611641 29.68700681782147 0.8683992717746128 +2458585.550582178 29.73908563439253 0.4316165838873574 +2458588.953128109 32.57853331880517 2.4561121534316137 +2458589.2688752552 32.84565731422025 0.3572845524654675 +2458589.5520900493 33.08453313438797 0.14095138991390269 2458589.7556874095 33.25567032655153 0.42510683925578957 2458590.6493888632 33.998206275493786 0.4052085434888088 -2458590.8992207763 34.202393862251064 2.1284734291803384 -2458591.1628828607 34.4158358428011 0.31413876399455787 -2458591.4564221804 34.65066614248212 0.27606202275888125 +2458590.8992207763 34.202393862257 2.1284734291803384 +2458591.1628828607 34.4158358428304 0.31413876399455787 +2458591.4564221804 34.65066614250525 0.27606202275888125 2458592.793054921 35.67197004299393 0.40789391022273064 -2458594.248453634 36.658853867632736 1.0877137223796849 -2458594.330772908 36.7096552320603 0.43000140392164243 +2458594.248453634 36.658853867659836 1.0877137223796849 +2458594.330772908 36.70965523208708 0.43000140392164243 2458594.661509989 36.9075044953399 0.6068475922305185 -2458595.170485553 37.19104655852463 1.1623218898607743 -2458595.2120056194 37.212987126945805 0.10969758209932548 -2458596.9540171996 37.946368064868295 0.15839404518426234 +2458595.170485553 37.191046558547804 1.1623218898607743 +2458595.2120056194 37.21298712696876 0.10969758209932548 +2458596.9540171996 37.94636806487273 0.15839404518426234 2458597.7140886877 38.13444098784455 1.2236950808204705 2458598.694059555 38.23818889469797 2.2948970335496153 2458598.8110234602 38.239390362436026 1.9977481736288853 2458598.831362406 38.239347076425304 1.150281471741598 -2458599.2498248946 38.221614101044516 2.97017466608847 -2458599.9043666334 38.12773045227851 1.9026471642181997 -2458601.267807699 37.66271794671086 1.0980673942133412 -2458603.4640042316 36.15447109339607 0.9123942597198087 -2458603.6215343312 36.013198260710745 1.9041337012527948 -2458603.887690115 35.765627056985196 2.0904233485588994 -2458604.204925004 35.456697188815184 0.2576659158284269 +2458599.2498248946 38.22161410104097 2.97017466608847 +2458599.9043666334 38.127730452277014 1.9026471642181997 +2458601.267807699 37.66271794669005 1.0980673942133412 +2458603.4640042316 36.15447109337033 0.9123942597198087 +2458603.6215343312 36.01319826070412 1.9041337012527948 +2458603.887690115 35.76562705697829 2.0904233485588994 +2458604.204925004 35.456697188771685 0.2576659158284269 2458604.783394053 34.85814955453509 1.9099794864950546 -2458605.394019192 34.18411742345329 0.7626235354774431 -2458605.559942743 33.99470667888917 0.10006439401856808 -2458605.864967669 33.64069057607246 0.1975373058549149 -2458606.5938742952 32.771268518428585 1.4770569656766863 -2458607.364738551 31.832464024907722 0.3260511570413087 -2458607.443547457 31.736210083805314 1.8390411207098787 -2458607.458327727 31.71816181666353 0.8016942788558589 -2458607.753146916 31.35867851173813 3.0198293755096155 -2458607.834049526 31.260296128720395 1.398882859282722 -2458608.302156811 30.69492059571507 0.13790075427094659 -2458608.531287327 30.421490705202377 1.8058506748302139 -2458608.84721951 30.049269598277764 0.990531369267568 +2458605.394019192 34.18411742341203 0.7626235354774431 +2458605.559942743 33.99470667887244 0.10006439401856808 +2458605.864967669 33.64069057606394 0.1975373058549149 +2458606.5938742952 32.77126851841098 1.4770569656766863 +2458607.364738551 31.832464024863295 0.3260511570413087 +2458607.443547457 31.73621008376978 1.8390411207098787 +2458607.458327727 31.71816181662798 0.8016942788558589 +2458607.753146916 31.358678511746984 3.0198293755096155 +2458607.8340495257 31.260296128711545 1.398882859282722 +2458608.302156811 30.694920595662715 0.13790075427094659 +2458608.531287327 30.421490705185075 1.8058506748302139 +2458608.84721951 30.04926959827776 0.990531369267568 2458609.702732182 29.078288741807796 0.21388687352231717 -2458610.100676085 28.64961764146608 0.5959456022338969 +2458610.100676085 28.649617641427646 0.5959456022338969 2458610.684298204 28.05268461990862 0.17399136091603062 -2458611.0237629865 27.724617560479846 0.3627180055815208 +2458611.0237629865 27.724617560459237 0.3627180055815208 2458611.7057417915 27.112027200089713 1.3097328444868137 -2458612.588065033 26.41848036870156 0.4786847536851103 -2458612.8156960513 26.258532382946537 0.9821711576214502 -2458612.8512979904 26.234234270946395 2.2336633370302765 -2458613.5528910076 25.79541633768526 1.1279442467003273 -2458613.7178363446 25.703369054103824 0.1374311086127715 -2458614.0152574293 25.548120925897713 0.41563316018358154 -2458614.3557728645 25.38729511068919 0.2979883260701478 -2458614.4936125544 25.327310167876284 0.5220090042826792 -2458615.460065639 24.988529266761528 2.2144799838588387 -2458616.0274894005 24.85492419961752 0.3268514922119821 -2458616.215406039 24.821032722156012 1.294627047607252 -2458617.40092397 24.721323463178404 0.5039819638915918 -2458618.124243317 24.752459662092168 1.8054186346463887 -2458619.582590718 25.00808484453209 0.22570988438055403 -2458619.88120014 25.08983186396636 0.6097943888061523 -2458620.3720763857 25.244331973272676 0.8865793524688304 -2458620.4839499546 25.282933446898774 0.3673242403538802 -2458620.6168157184 25.330369749007563 0.7002513842636655 -2458621.1425666166 25.53454732364628 0.27868138162595496 -2458621.2097744853 25.56249226624152 3.0278799709088293 -2458621.324310234 25.611053674390313 1.0737200000321394 -2458622.227921901 26.033887732374076 0.8794073265724713 -2458624.25329212 27.20768557849734 2.522442566828421 -2458624.3296302403 27.257266547154213 0.45815190845502335 -2458625.029118305 27.72728987447476 2.9641762763192836 -2458625.325566464 27.934560475217253 0.1364648947937836 -2458625.539268654 28.08676009203145 3.036108524153798 -2458625.5620748666 28.1031362130567 1.2009865117917156 -2458627.668461853 29.71104598955161 0.7411563288591058 -2458627.7730806074 29.79506029034281 0.3425546024473731 +2458612.588065033 26.418480368696326 0.4786847536851103 +2458612.815696051 26.258532382951525 0.9821711576214502 +2458612.8512979904 26.23423427094639 2.2336633370302765 +2458613.5528910076 25.7954163376728 1.1279442467003273 +2458613.7178363446 25.70336905410382 0.1374311086127715 +2458614.0152574293 25.548120925886817 0.41563316018358154 +2458614.3557728645 25.387295110672977 0.2979883260701478 +2458614.4936125544 25.327310167867022 0.5220090042826792 +2458615.460065639 24.98852926675345 2.2144799838588387 +2458616.0274894005 24.85492419961188 0.3268514922119821 +2458616.215406039 24.821032722148725 1.294627047607252 +2458617.40092397 24.721323463178273 0.5039819638915918 +2458618.124243317 24.7524596620954 1.8054186346463887 +2458619.582590718 25.008084844535844 0.22570988438055403 +2458619.88120014 25.089831863968463 0.6097943888061523 +2458620.3720763857 25.244331973287498 0.8865793524688304 +2458620.4839499546 25.282933446908977 0.3673242403538802 +2458620.6168157184 25.33036974901021 0.7002513842636655 +2458621.1425666166 25.534547323661293 0.27868138162595496 +2458621.2097744853 25.56249226625981 3.0278799709088293 +2458621.324310234 25.611053674409046 1.0737200000321394 +2458622.227921901 26.033887732396124 0.8794073265724713 +2458624.25329212 27.20768557852559 2.522442566828421 +2458624.3296302403 27.25726654718267 0.45815190845502335 +2458625.029118305 27.727289874489855 2.9641762763192836 +2458625.325566464 27.934560475248123 0.1364648947937836 +2458625.5392686534 28.086760092041906 3.036108524153798 +2458625.5620748666 28.103136213067156 1.2009865117917156 +2458627.668461853 29.71104598955745 0.7411563288591058 +2458627.7730806074 29.795060290342818 0.3425546024473731 2458627.798571989 29.815579632735133 0.20166476827281782 -2458627.9599760915 29.945929352130356 2.9773974520320747 -2458628.2126703407 30.151416015285264 0.22533700062538536 -2458628.480195208 30.370692926345793 2.0900733413764665 -2458628.5599593506 30.436390124152563 0.12177864688520286 -2458629.073666818 30.862591208910104 1.591829215992026 -2458629.082890462 30.870287886130374 2.0154718662643663 -2458629.538034037 31.251682937794833 0.23300682372018364 -2458631.4067446967 32.833008060857814 2.920315439008531 -2458631.439220519 32.86044439880763 0.3611990849813782 -2458631.9908884 33.32473249694268 0.17279606014459417 -2458632.2657493623 33.55439972968167 0.42582900028334236 -2458632.357793869 33.63099698917818 2.8740112294535667 -2458632.5502393283 33.79057060477975 0.1479915242733024 -2458633.265320477 34.37520779709099 1.0951117382578566 -2458633.3711873037 34.46044294859951 0.16000665547396625 +2458627.9599760915 29.945929352142137 2.9773974520320747 +2458628.212670341 30.15141601532092 0.22533700062538536 +2458628.480195208 30.370692926369745 2.0900733413764665 +2458628.5599593506 30.436390124164575 0.12177864688520286 +2458629.073666818 30.86259120893438 1.591829215992026 +2458629.082890462 30.870287886160703 2.0154718662643663 +2458629.538034037 31.25168293781321 0.23300682372018364 +2458631.4067446967 32.83300806088855 2.920315439008531 +2458631.439220519 32.86044439883223 0.3611990849813782 +2458631.9908884 33.324732496960955 0.17279606014459417 +2458632.2657493623 33.55439972971805 0.42582900028334236 +2458632.357793869 33.630996989208434 2.8740112294535667 +2458632.5502393283 33.79057060479176 0.1479915242733024 +2458633.265320477 34.37520779713209 1.0951117382578566 +2458633.371187304 34.46044294862873 0.16000665547396625 2458634.777263851 35.5482433458863 1.066462842693139 2458635.876221915 36.32121488397856 0.17034520038873252 -2458635.928955045 36.35614081211861 2.603902593414404 -2458636.4556365996 36.692457645567515 2.3472898783794176 +2458635.928955045 36.356140812128224 2.603902593414404 +2458636.4556365996 36.69245764558545 2.3472898783794176 2458636.7794928127 36.88713792983214 2.996540209142504 -2458637.50402638 37.2851642167502 2.077893722672471 +2458637.50402638 37.28516421676137 2.077893722672471 diff --git a/docs/examples/make-data.ipynb b/docs/examples/make-data.ipynb index 25e2ba52..412a0f3e 100644 --- a/docs/examples/make-data.ipynb +++ b/docs/examples/make-data.ipynb @@ -15,14 +15,15 @@ "source": [ "import pickle\n", "\n", - "from astropy.time import Time\n", "import astropy.units as u\n", - "from astropy.table import QTable\n", + "import matplotlib.pyplot as plt\n", "import numpy as np\n", - "%matplotlib inline\n", - "\n", + "from astropy.table import QTable\n", + "from astropy.time import Time\n", + "from thejoker import RVData\n", "from twobody import KeplerOrbit, PolynomialRVTrend\n", - "from thejoker import RVData" + "\n", + "%matplotlib inline" ] }, { @@ -49,25 +50,32 @@ "source": [ "n_data = 256\n", "\n", - "t0 = Time('2019-1-1') + rnd.uniform(0., 40) * u.day\n", + "t0 = Time(\"2019-1-1\") + rnd.uniform(0.0, 40) * u.day\n", "\n", "truth1 = dict()\n", - "truth1['P'] = rnd.uniform(40, 80) * u.day\n", - "truth1['M0'] = rnd.uniform(0., 2*np.pi) * u.radian\n", - "truth1['omega'] = rnd.uniform(0., 2*np.pi) * u.radian\n", - "truth1['e'] = 0.1 * u.one\n", - "truth1['K'] = rnd.uniform(5, 15) * u.km/u.s\n", - "truth1['v0'] = rnd.uniform(-50, 50) * u.km/u.s\n", + "truth1[\"P\"] = rnd.uniform(40, 80) * u.day\n", + "truth1[\"M0\"] = rnd.uniform(0.0, 2 * np.pi) * u.radian\n", + "truth1[\"omega\"] = rnd.uniform(0.0, 2 * np.pi) * u.radian\n", + "truth1[\"e\"] = 0.1 * u.one\n", + "truth1[\"K\"] = rnd.uniform(5, 15) * u.km / u.s\n", + "truth1[\"v0\"] = rnd.uniform(-50, 50) * u.km / u.s\n", "\n", - "orbit = KeplerOrbit(P=truth1['P'], e=truth1['e'], omega=truth1['omega'],\n", - " M0=truth1['M0'], t0=t0, K=truth1['K'],\n", - " i=90*u.deg, Omega=0*u.deg, # these don't matter\n", - " barycenter=PolynomialRVTrend([truth1['v0']]))\n", + "orbit = KeplerOrbit(\n", + " P=truth1[\"P\"],\n", + " e=truth1[\"e\"],\n", + " omega=truth1[\"omega\"],\n", + " M0=truth1[\"M0\"],\n", + " t0=t0,\n", + " K=truth1[\"K\"],\n", + " i=90 * u.deg,\n", + " Omega=0 * u.deg, # these don't matter\n", + " barycenter=PolynomialRVTrend([truth1[\"v0\"]]),\n", + ")\n", "\n", - "t = t0 + truth1['P'] * np.concatenate(([0], np.sort(rnd.uniform(0, 3., n_data))))\n", + "t = t0 + truth1[\"P\"] * np.concatenate(([0], np.sort(rnd.uniform(0, 3.0, n_data))))\n", "\n", "rv = orbit.radial_velocity(t)\n", - "err = 10 ** rnd.uniform(-1, 0.5, size=len(rv)) * u.km/u.s\n", + "err = 10 ** rnd.uniform(-1, 0.5, size=len(rv)) * u.km / u.s\n", "data = RVData(t, rv, rv_err=err)" ] }, @@ -78,11 +86,11 @@ "outputs": [], "source": [ "tbl = QTable()\n", - "tbl['bjd'] = data.t.tcb.jd\n", - "tbl['rv'] = data.rv\n", - "tbl['rv_err'] = data.rv_err\n", - "tbl.meta['t_ref'] = data.t_ref\n", - "tbl.write('data.ecsv', overwrite=True)" + "tbl[\"bjd\"] = data.t.tcb.jd\n", + "tbl[\"rv\"] = data.rv\n", + "tbl[\"rv_err\"] = data.rv_err\n", + "tbl.meta[\"t_ref\"] = data.t_ref\n", + "tbl.write(\"data.ecsv\", overwrite=True)" ] }, { @@ -91,7 +99,7 @@ "metadata": {}, "outputs": [], "source": [ - "with open('true-orbit.pkl', 'wb') as f:\n", + "with open(\"true-orbit.pkl\", \"wb\") as f:\n", " pickle.dump(truth1, f)" ] }, @@ -110,46 +118,60 @@ "source": [ "n_data = 256\n", "\n", - "t0 = Time('2019-1-1') + rnd.uniform(0., 40) * u.day\n", + "t0 = Time(\"2019-1-1\") + rnd.uniform(0.0, 40) * u.day\n", "\n", "truth1 = dict()\n", - "truth1['P'] = rnd.uniform(40, 80) * u.day\n", - "truth1['M0'] = rnd.uniform(0., 2*np.pi) * u.radian\n", - "truth1['omega'] = rnd.uniform(0., 2*np.pi) * u.radian\n", - "truth1['e'] = 0.25 * u.one\n", - "truth1['K'] = rnd.uniform(5, 15) * u.km/u.s\n", - "truth1['v0'] = rnd.uniform(-50, 50) * u.km/u.s\n", + "truth1[\"P\"] = rnd.uniform(40, 80) * u.day\n", + "truth1[\"M0\"] = rnd.uniform(0.0, 2 * np.pi) * u.radian\n", + "truth1[\"omega\"] = rnd.uniform(0.0, 2 * np.pi) * u.radian\n", + "truth1[\"e\"] = 0.25 * u.one\n", + "truth1[\"K\"] = rnd.uniform(5, 15) * u.km / u.s\n", + "truth1[\"v0\"] = rnd.uniform(-50, 50) * u.km / u.s\n", "\n", - "orbit = KeplerOrbit(P=truth1['P'], e=truth1['e'], omega=truth1['omega'],\n", - " M0=truth1['M0'], t0=t0, K=truth1['K'],\n", - " i=90*u.deg, Omega=0*u.deg, # these don't matter\n", - " barycenter=PolynomialRVTrend([truth1['v0']]))\n", + "orbit = KeplerOrbit(\n", + " P=truth1[\"P\"],\n", + " e=truth1[\"e\"],\n", + " omega=truth1[\"omega\"],\n", + " M0=truth1[\"M0\"],\n", + " t0=t0,\n", + " K=truth1[\"K\"],\n", + " i=90 * u.deg,\n", + " Omega=0 * u.deg, # these don't matter\n", + " barycenter=PolynomialRVTrend([truth1[\"v0\"]]),\n", + ")\n", "\n", - "with open('true-orbit-triple.pkl', 'wb') as f:\n", + "with open(\"true-orbit-triple.pkl\", \"wb\") as f:\n", " pickle.dump(truth1, f)\n", "\n", "truth2 = dict()\n", - "truth2['P'] = 10 * rnd.uniform(40, 80) * u.day\n", - "truth2['M0'] = rnd.uniform(0., 2*np.pi) * u.radian\n", - "truth2['omega'] = rnd.uniform(0., 2*np.pi) * u.radian\n", - "truth2['e'] = 0.1\n", - "truth2['K'] = 13 * u.km/u.s\n", + "truth2[\"P\"] = 10 * rnd.uniform(40, 80) * u.day\n", + "truth2[\"M0\"] = rnd.uniform(0.0, 2 * np.pi) * u.radian\n", + "truth2[\"omega\"] = rnd.uniform(0.0, 2 * np.pi) * u.radian\n", + "truth2[\"e\"] = 0.1\n", + "truth2[\"K\"] = 13 * u.km / u.s\n", "\n", - "orbit2 = KeplerOrbit(P=truth2['P'], e=truth2['e'], omega=truth2['omega'],\n", - " M0=truth2['M0'], t0=t0, K=truth2['K'],\n", - " i=90*u.deg, Omega=0*u.deg)\n", + "orbit2 = KeplerOrbit(\n", + " P=truth2[\"P\"],\n", + " e=truth2[\"e\"],\n", + " omega=truth2[\"omega\"],\n", + " M0=truth2[\"M0\"],\n", + " t0=t0,\n", + " K=truth2[\"K\"],\n", + " i=90 * u.deg,\n", + " Omega=0 * u.deg,\n", + ")\n", "\n", - "t = t0 + truth1['P'] * np.concatenate(([0], np.sort(rnd.uniform(0, 5., n_data))))\n", + "t = t0 + truth1[\"P\"] * np.concatenate(([0], np.sort(rnd.uniform(0, 5.0, n_data))))\n", "\n", "rv = orbit.radial_velocity(t) + orbit2.radial_velocity(t)\n", - "err = 10**rnd.uniform(-1, 0.5, size=len(rv)) * u.km/u.s\n", + "err = 10 ** rnd.uniform(-1, 0.5, size=len(rv)) * u.km / u.s\n", "data = RVData(t, rv, rv_err=err)\n", "\n", "tbl = QTable()\n", - "tbl['bjd'] = data.t.tcb.jd\n", - "tbl['rv'] = data.rv\n", - "tbl['rv_err'] = data.rv_err\n", - "tbl.write('data-triple.ecsv', overwrite=True)" + "tbl[\"bjd\"] = data.t.tcb.jd\n", + "tbl[\"rv\"] = data.rv\n", + "tbl[\"rv_err\"] = data.rv_err\n", + "tbl.write(\"data-triple.ecsv\", overwrite=True)" ] }, { @@ -176,44 +198,39 @@ "source": [ "n_data = 16\n", "\n", - "t0 = Time('2019-1-1') + rnd.uniform(0., 40) * u.day\n", + "t0 = Time(\"2019-1-1\") + rnd.uniform(0.0, 40) * u.day\n", "\n", "truth1 = dict()\n", - "truth1['P'] = rnd.uniform(40, 80) * u.day\n", - "truth1['M0'] = rnd.uniform(0., 2*np.pi) * u.radian\n", - "truth1['omega'] = rnd.uniform(0., 2*np.pi) * u.radian\n", - "truth1['e'] = 0.13 * u.one\n", - "truth1['K'] = rnd.uniform(5, 15) * u.km/u.s\n", - "truth1['v0'] = rnd.uniform(-50, 50) * u.km/u.s\n", + "truth1[\"P\"] = rnd.uniform(40, 80) * u.day\n", + "truth1[\"M0\"] = rnd.uniform(0.0, 2 * np.pi) * u.radian\n", + "truth1[\"omega\"] = rnd.uniform(0.0, 2 * np.pi) * u.radian\n", + "truth1[\"e\"] = 0.13 * u.one\n", + "truth1[\"K\"] = rnd.uniform(5, 15) * u.km / u.s\n", + "truth1[\"v0\"] = rnd.uniform(-50, 50) * u.km / u.s\n", "\n", - "orbit = KeplerOrbit(P=truth1['P'], e=truth1['e'], omega=truth1['omega'],\n", - " M0=truth1['M0'], t0=t0, K=truth1['K'],\n", - " i=90*u.deg, Omega=0*u.deg, # these don't matter\n", - " barycenter=PolynomialRVTrend([truth1['v0']]))\n", + "orbit = KeplerOrbit(\n", + " P=truth1[\"P\"],\n", + " e=truth1[\"e\"],\n", + " omega=truth1[\"omega\"],\n", + " M0=truth1[\"M0\"],\n", + " t0=t0,\n", + " K=truth1[\"K\"],\n", + " i=90 * u.deg,\n", + " Omega=0 * u.deg, # these don't matter\n", + " barycenter=PolynomialRVTrend([truth1[\"v0\"]]),\n", + ")\n", "\n", - "t = t0 + truth1['P'] * np.concatenate(([0], np.sort(rnd.uniform(0, 3., n_data))))\n", + "t = t0 + truth1[\"P\"] * np.concatenate(([0], np.sort(rnd.uniform(0, 3.0, n_data))))\n", "\n", "rv = orbit.radial_velocity(t)\n", - "err = 10**rnd.uniform(-1, 0.5, size=len(rv)) * u.km/u.s\n", + "err = 10 ** rnd.uniform(-1, 0.5, size=len(rv)) * u.km / u.s\n", "\n", "data1 = RVData(t[:10], rv[:10], rv_err=err[:10])\n", "\n", - "rv[10:] += 4.8 * u.km/u.s\n", + "rv[10:] += 4.8 * u.km / u.s\n", "data2 = RVData(t[10:], rv[10:], rv_err=err[10:])" ] }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import matplotlib as mpl\n", - "import matplotlib.pyplot as plt\n", - "%matplotlib inline\n", - "import numpy as np" - ] - }, { "cell_type": "code", "execution_count": null, @@ -233,26 +250,19 @@ "source": [ "for i, data in enumerate([data1, data2]):\n", " tbl = QTable()\n", - " tbl['bjd'] = data.t.tcb.jd\n", - " tbl['rv'] = data.rv\n", - " tbl['rv_err'] = data.rv_err\n", - " tbl.meta['t_ref'] = data1.t0\n", - " tbl.write(f'data-survey{i+1}.ecsv', overwrite=True)" + " tbl[\"bjd\"] = data.t.tcb.jd\n", + " tbl[\"rv\"] = data.rv\n", + " tbl[\"rv_err\"] = data.rv_err\n", + " tbl.meta[\"t_ref\"] = data1.t_ref\n", + " tbl.write(f\"data-survey{i+1}.ecsv\", overwrite=True)" ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] } ], "metadata": { "kernelspec": { - "display_name": "Python [conda env:root] *", + "display_name": "thejoker2024", "language": "python", - "name": "conda-root-py" + "name": "python3" }, "language_info": { "codemirror_mode": { @@ -264,7 +274,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.3" + "version": "3.11.8" }, "nbsphinx": { "orphan": true diff --git a/docs/examples/true-orbit-triple.pkl b/docs/examples/true-orbit-triple.pkl index 3955c145815d10f964ccef7300d5f068823621a7..2cd734d1fc94a815f25b75ee754baafab873c0b9 100644 GIT binary patch delta 473 zcmey()5*iyz%un8>qgcWjFbN}s!Y~rGOl+@%rK9gqLJa8;gXtRSvw`ex`!hpzqCj} zBSUjaX;NZ_ZR`|p#>5PJcLxgwFvxJokk@RSQVY}s6#Oay5mbtul3@r}?U3PBnpB#U zVGK07hcRVJ4|7Ul<&@M6Q=lq;u#Fi3i6HBEfcCKiDP!Vf2^MpCdw2T}P=|o*1Ue7spo-F@;Y?#T(Pa>279O6Pzb2y`D`c4kg~ZfZ#?Fi^5{ zL7|dimHz?^ CA)?Fx delta 564 zcmeC=`OU-Hz%sRfZ6oUoMxh?o_>9Ejj453lGb`20t-0S6O@7O$!joa!#nHmW-OI(H zbY-$2lX1OsVupF_6paj*4A;~Q%i1X!);$~<`K3h)8X1~XK&mGPKI;qO=4dHRO3bi} zo#M@ynBm~=V8H+eAO*K(u5vkJ&Ac(gF+*OnaY`-FFrb335)cJSu~RY(;W`~Nyi1cx zlQN8fR`xKaOzD}tnu%LA0PKd0z(lA6*i#ZK6*N*H_IqU8Y@2iBC)ecjOkslY840Nw zra&X(AzBh9n=!|MJ!Rls`f$mz1)`IuFq<+aO+LkJF7M!O{{iYukSlUG__vD4l5|b@hBp4kg2eO0;PRN)DF?s^T=!uiJvuGKDJeZys$FObV-7Bae(C+{b zfeQ!QKP$hpUIjG1XYwx=IgU9HBj!$4V3n{38S!GH<%PIf%R->n4O2lb(eP&U`VR!D msG$Zk6zE=!jLBfNle1af!%@AWkqgcWjFbN}s!Y~rGOl+@%rK9gqLJa8;gXtRSvw`ex`!hpzqCj} zBSUjaX;NZ_ZR`|p#>5PJcLxgwFvxJokk@RSQVY}s6#Oay5mbtul3@r}?U3PBnpB#U zVGK07hcRVJ4|7Ul<&@M6Q=lq;u#Fi3i6HBEfcCKiDP!Vf2^MpCdw2T}P=|o*1Ue7spo-F@;Y?#T(Pa>279O6Pzb2y`D`c4kg~ZfZ#?Fi^5{ zL7|dimHz?^ CA)?Fx delta 564 zcmeC=`OU-Hz%sRfZ6oUoMxh?o_>9Ejj453lGb`20t-0S6O@7O$!joa!#nHmW-OI(H zbY-$2lX1OsVupF_6paj*4A;~Q%i1X!);$~<`K3h)8X1~XK&mGPKI;qO=4dHRO3bi} zo#M@ynBm~=V8H+eAO*K(u5vkJ&Ac(gF+*OnaY`-FFrb335)cJSu~RY(;W`~Nyi1cx zlQN8fR`xKaOzD}tnu%LA0PKd0z(lA6*i#ZK6*N*H_IqU8Y@2iBC)ecjOkslY840Nw zra&X(AzBh9n=!|MJ!Rls`f$mz1)`IuFq<+aO+LkJF7M!O{{iYukSlUG__vD4l5|b@hBp4kg2eO0;PRN)DF?s^T=!uiJvuGKDJeZys$FObV-7Bae(C+{b zfeQ!QKP$hpUIjG1XYwx=IgU9HBj!$4V3n{38S!GH<%PIf%R->n4O2lb(eP&U`VR!D msG$Zk6zE=!jLBfNle1af!%@AWk Date: Tue, 5 Mar 2024 22:50:25 -0500 Subject: [PATCH 33/50] skip multipool test on mac --- thejoker/tests/test_sampler.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/thejoker/tests/test_sampler.py b/thejoker/tests/test_sampler.py index c89c7708..2f0988fb 100644 --- a/thejoker/tests/test_sampler.py +++ b/thejoker/tests/test_sampler.py @@ -1,11 +1,12 @@ import os +import platform import astropy.units as u import numpy as np import pymc as pm import pytest from astropy.time import Time -from schwimmbad import SerialPool +from schwimmbad import MultiPool, SerialPool from twobody import KeplerOrbit from thejoker.data import RVData @@ -102,10 +103,11 @@ def test_marginal_ln_likelihood(tmpdir, case): assert len(ll) == len(prior_samples) # NOTE: this makes it so I can't parallelize tests, I think - # with MultiPool(processes=2) as pool: - # joker = TheJoker(prior, pool=pool) - # ll = joker.marginal_ln_likelihood(data, filename) - # assert len(ll) == len(prior_samples) + if platform.system() != "Darwin": # this test fails on CI + with MultiPool(processes=2) as pool: + joker = TheJoker(prior, pool=pool) + ll = joker.marginal_ln_likelihood(data, filename) + assert len(ll) == len(prior_samples) priors = [ From 0829c56611988871bba37e9d77cd8dfa6a1c5750 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Tue, 5 Mar 2024 22:58:51 -0500 Subject: [PATCH 34/50] fix notebook 1 --- docs/examples/1-Getting-started.ipynb | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/docs/examples/1-Getting-started.ipynb b/docs/examples/1-Getting-started.ipynb index 4e1c23c5..23883f56 100644 --- a/docs/examples/1-Getting-started.ipynb +++ b/docs/examples/1-Getting-started.ipynb @@ -37,12 +37,12 @@ "outputs": [], "source": [ "import astropy.table as at\n", - "from astropy.time import Time\n", "import astropy.units as u\n", - "from astropy.visualization.units import quantity_support\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "import thejoker as tj\n", + "from astropy.time import Time\n", + "from astropy.visualization.units import quantity_support\n", "\n", "%matplotlib inline" ] @@ -394,7 +394,7 @@ " ax.scatter(joker_samples[\"P\"], joker_samples[\"e\"], s=20, lw=0, alpha=0.5)\n", "\n", "ax.set_xscale(\"log\")\n", - "ax.set_xlim(prior.pars[\"P\"].distribution.a, prior.pars[\"P\"].distribution.b)\n", + "ax.set_xlim(1, 1e3)\n", "ax.set_ylim(0, 1)\n", "\n", "ax.set_xlabel(\"$P$ [day]\")\n", @@ -438,7 +438,7 @@ " )\n", "\n", "ax.set_xscale(\"log\")\n", - "ax.set_xlim(prior.pars[\"P\"].distribution.a, prior.pars[\"P\"].distribution.b)\n", + "ax.set_xlim(1, 1e3)\n", "ax.set_ylim(0, 1)\n", "\n", "ax.set_xlabel(\"$P$ [day]\")\n", @@ -455,9 +455,9 @@ ], "metadata": { "kernelspec": { - "display_name": "Python [conda env:root] *", + "display_name": "thejoker2024", "language": "python", - "name": "conda-root-py" + "name": "python3" }, "language_info": { "codemirror_mode": { @@ -469,7 +469,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.3" + "version": "3.11.8" }, "toc": { "base_numbering": 1, From 07ee421607b111ccc63bfeff2e858aa1a5c2256c Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Tue, 5 Mar 2024 22:59:17 -0500 Subject: [PATCH 35/50] exit early if nb fails --- docs/run_notebooks.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/docs/run_notebooks.py b/docs/run_notebooks.py index d0d32ff4..cb82e644 100644 --- a/docs/run_notebooks.py +++ b/docs/run_notebooks.py @@ -46,10 +46,7 @@ def process_notebook(filename, kernel_name=None): nbsphinx_kernel_name = os.environ.get("NBSPHINX_KERNEL", "python3") - fail = False for filename in sorted(glob.glob(pattern)): success = process_notebook(filename, kernel_name=nbsphinx_kernel_name) - fail = fail or not success - - if fail: - sys.exit(1) + if not success: + sys.exit(1) From 76873d23d18664719fd2e15ade4776cd725a9b69 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Tue, 5 Mar 2024 23:15:52 -0500 Subject: [PATCH 36/50] note about exoplanet --- thejoker/_keplerian_orbit.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/thejoker/_keplerian_orbit.py b/thejoker/_keplerian_orbit.py index 1b584334..18fc980a 100644 --- a/thejoker/_keplerian_orbit.py +++ b/thejoker/_keplerian_orbit.py @@ -1,3 +1,7 @@ +""" +NOTE: This comes from the exoplanet-devs/exoplanet project! Modified to work with pymc +instead of pymc3. +""" __all__ = [ "KeplerianOrbit", "get_true_anomaly", From 11f42dd1a2cdbb42b4dc34c628c734c80d17dc21 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Tue, 5 Mar 2024 23:20:40 -0500 Subject: [PATCH 37/50] docs tweaks --- docs/index.rst | 42 +++++++++++++++++++++++++----------------- docs/install.rst | 3 ++- 2 files changed, 27 insertions(+), 18 deletions(-) diff --git a/docs/index.rst b/docs/index.rst index 2300afa3..183dd563 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -6,16 +6,22 @@ Introduction ============ |thejoker| [#f1]_ is a custom Monte Carlo sampler for the `two-body problem -`_ that generates posterior -samplings in Keplerian orbital parameters given radial velocity observations of -stars. It is designed to deliver converged posterior samplings even when the -radial velocity measurements are sparse or very noisy. It is therefore useful -for constraining the orbital properties of binary star or star-planet systems. -Though it fundamentally assumes that any system has two massive bodies (and only -the primary is observed), |thejoker| can also be used for hierarchical systems -in which the velocity perturbations from a third or other bodies are much longer -than the dominant companion. See the paper [#f2]_ for more details about the -method and applications. +`_ that generates posterior samplings in +Keplerian orbital parameters given radial velocity observations of stars. It is designed +to deliver converged posterior samplings even when the radial velocity measurements are +sparse or very noisy. It is therefore useful for constraining the orbital properties of +binary star or star-planet systems. Though it fundamentally assumes that any system has +two massive bodies (and only the primary is observed), |thejoker| can also be used for +hierarchical systems in which the velocity perturbations from a third or other bodies +are much longer than the dominant companion. See the paper [#f2]_ for more details about +the method and applications. + +.. note:: + + New with v1.3, ``thejoker`` is no longer based on ``pymc3`` and is now instead built + on `pymc `_. If you used ``thejoker`` previously + with ``pymc3``, you will have to upgrade and modify some code to use the new version + of ``pymc`` instead. .. toctree:: :maxdepth: 1 @@ -70,8 +76,8 @@ object: >>> rv = [38.77, 39.70, 37.45, 38.31, 38.31] * u.km/u.s >>> err = [0.184, 0.261, 0.112, 0.155, 0.223] * u.km/u.s >>> data = tj.RVData(t=t, rv=rv, rv_err=err) - >>> ax = data.plot() # doctest: +SKIP - >>> ax.set_xlim(-10, 200) # doctest: +SKIP + >>> ax = data.plot() # doctest: +SKIP + >>> ax.set_xlim(-10, 200) # doctest: +SKIP .. plot:: :align: center @@ -85,7 +91,7 @@ object: err = [0.184, 0.261, 0.112, 0.155, 0.223] * u.km/u.s data = RVData(t=t, rv=rv, rv_err=err) - ax = data.plot() # doctest: +SKIP + ax = data.plot() # doctest: +SKIP ax.set_xlim(-10, 200) We next need to specify the prior distributions for the parameters of @@ -96,16 +102,18 @@ along with parameters that specify the prior over the linear parameters in *The Joker* (the velocity semi-amplitude, ``K``, and the systemic velocity, ``v0``): >>> import numpy as np - >>> prior = tj.JokerPrior.default(P_min=2*u.day, P_max=256*u.day, - ... sigma_K0=30*u.km/u.s, - ... sigma_v=100*u.km/u.s) + >>> prior = tj.JokerPrior.default( + ... P_min=2*u.day, P_max=256*u.day, + ... sigma_K0=30*u.km/u.s, + ... sigma_v=100*u.km/u.s + ... ) With the data and prior created, we can now instantiate the sampler object and run the rejection sampler: >>> joker = tj.TheJoker(prior) >>> prior_samples = prior.sample(size=100_000) - >>> samples = joker.rejection_sample(data, prior_samples) # doctest: +SKIP + >>> samples = joker.rejection_sample(data, prior_samples) # doctest: +SKIP Of the 100_000 prior samples we generated, only a handful pass the rejection sampling step of |thejoker|. Let's visualize the surviving samples in the diff --git a/docs/install.rst b/docs/install.rst index 6c389711..a67ee761 100644 --- a/docs/install.rst +++ b/docs/install.rst @@ -31,4 +31,5 @@ To install from source (i.e. from the cloned Git repository), also use pip: Dependencies ============ -See the `pyproject.toml `_ file for the most up-to-date list of dependencies. +See the `pyproject.toml `_ +file for the most up-to-date list of dependencies. From a098fcfddcf226b8b30487308a0eaf1f6d411269 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Tue, 5 Mar 2024 23:24:34 -0500 Subject: [PATCH 38/50] import os --- docs/conf.py | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/conf.py b/docs/conf.py index f1cccf9f..8a11fab2 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -1,4 +1,5 @@ import importlib.metadata +import os project = "thejoker" copyright = "2024, Adrian Price-Whelan" From 20d35578cf67b72e21c87f895d8ff7e28eb5ea55 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Tue, 5 Mar 2024 23:42:55 -0500 Subject: [PATCH 39/50] fix various sphinx things --- docs/conf.py | 7 ++++--- docs/examples/notebook_setup.py | 2 ++ docs/index.rst | 8 +++++--- pyproject.toml | 3 +++ 4 files changed, 14 insertions(+), 6 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index 8a11fab2..8863f4f1 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -11,11 +11,12 @@ "sphinx.ext.intersphinx", "sphinx.ext.mathjax", "sphinx.ext.napoleon", - "sphinx_autodoc_typehints", + # "sphinx_autodoc_typehints", "sphinx_copybutton", "sphinx_automodapi.automodapi", "sphinx_automodapi.smart_resolver", "rtds_action", + "matplotlib.sphinxext.plot_directive", ] source_suffix = [".rst", ".md"] @@ -37,7 +38,7 @@ intersphinx_mapping = { "python": ("https://docs.python.org/3", None), "pandas": ("https://pandas.pydata.org/pandas-docs/stable/", None), - "pymc": ("https://docs.pymc.io/", None), + "pymc": ("https://www.pymc.io/projects/docs/en/stable", None), "h5py": ("http://docs.h5py.org/en/latest/", None), "twobody": ("https://twobody.readthedocs.io/en/latest/", None), "schwimmbad": ("https://schwimmbad.readthedocs.io/en/latest/", None), @@ -74,7 +75,7 @@ rtds_action_artifact_prefix = "notebooks-for-" # A GitHub personal access token is required, more info below -rtds_action_github_token = os.environ["GITHUB_TOKEN"] +rtds_action_github_token = os.environ.get("GITHUB_TOKEN", "") # Whether or not to raise an error on Read the Docs if the # artifact containing the notebooks can't be downloaded (optional) diff --git a/docs/examples/notebook_setup.py b/docs/examples/notebook_setup.py index 5f9c0c91..5372ee1a 100644 --- a/docs/examples/notebook_setup.py +++ b/docs/examples/notebook_setup.py @@ -1,3 +1,5 @@ +get_ipython().magic('config InlineBackend.figure_format = "retina"') # noqa + import warnings import matplotlib.pyplot as plt diff --git a/docs/index.rst b/docs/index.rst index 183dd563..5402cc7c 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -142,17 +142,17 @@ below to see how these were made): prior_samples = prior.sample(size=100_000) samples = joker.rejection_sample(data, prior_samples) - fig, ax = plt.subplots(1, 1, figsize=(6,6)) # doctest: +SKIP + fig, ax = plt.subplots(1, 1, figsize=(6, 4)) # doctest: +SKIP ax.scatter(samples['P'].value, samples['K'].to(u.km/u.s).value, marker='.', color='k', alpha=0.45) # doctest: +SKIP ax.set_xlabel("$P$ [day]") ax.set_ylabel("$K$ [km/s]") - ax.set_xlim(2, 256) + ax.set_xlim(0, 256) ax.set_ylim(0.75, 3.) ax.scatter(61.942, 1.3959, marker='o', color='#31a354', zorder=-100) - fig, ax = plt.subplots(1, 1, figsize=(8,5)) # doctest: +SKIP + fig, ax = plt.subplots(1, 1, figsize=(6, 4)) # doctest: +SKIP t_grid = np.linspace(-10, 210, 1024) plot_rv_curves(samples, t_grid, rv_unit=u.km/u.s, data=data, ax=ax, plot_kwargs=dict(color='#888888')) @@ -170,3 +170,5 @@ API .. [#f1] Short for Johannes Kepler. .. [#f2] ``_ + +.. |thejoker| replace:: *The Joker* \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 74e7c178..0157d085 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,9 @@ docs = [ "nbformat", "ipykernel", "matplotlib", + "sphinx_copybutton", + "rtds_action", + "pydata-sphinx-theme" ] tutorials = [ "thejoker[docs]", From ca6d17fa836df0c9261e6e5cdba33798316872d2 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Tue, 5 Mar 2024 23:56:57 -0500 Subject: [PATCH 40/50] ignore docs build dir --- docs/examples/notebook_setup.py | 6 +++++- docs/index.rst | 1 + pyproject.toml | 3 ++- 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/docs/examples/notebook_setup.py b/docs/examples/notebook_setup.py index 5372ee1a..009351d2 100644 --- a/docs/examples/notebook_setup.py +++ b/docs/examples/notebook_setup.py @@ -1,4 +1,8 @@ -get_ipython().magic('config InlineBackend.figure_format = "retina"') # noqa +from IPython.core.getipython import get_ipython + +ipy = get_ipython() +if ipy is not None: + ipy.magic('config InlineBackend.figure_format = "retina"') import warnings diff --git a/docs/index.rst b/docs/index.rst index 5402cc7c..88b4e159 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -126,6 +126,7 @@ below to see how these were made): :align: center :width: 512 + import matplotlib.pyplot as plt from thejoker import JokerPrior, TheJoker, RVData from thejoker.plot import plot_rv_curves import astropy.units as u diff --git a/pyproject.toml b/pyproject.toml index 0157d085..737d22eb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -82,7 +82,8 @@ testpaths = ["thejoker", "docs"] doctest_plus = "enabled" text_file_format = "rst" addopts = [ - "--doctest-rst", "-ra", "--showlocals", "--strict-markers", "--strict-config" + "--doctest-rst", "-ra", "--showlocals", "--strict-markers", "--strict-config", + "--ignore=docs/_build" ] xfail_strict = true filterwarnings = [ From 6322061a057296a3269a28c812bca3ba33887727 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Wed, 6 Mar 2024 08:25:10 -0500 Subject: [PATCH 41/50] trying to figure out what is failing on CI --- thejoker/tests/test_sampler.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/thejoker/tests/test_sampler.py b/thejoker/tests/test_sampler.py index 2f0988fb..a95930d3 100644 --- a/thejoker/tests/test_sampler.py +++ b/thejoker/tests/test_sampler.py @@ -92,8 +92,9 @@ def test_marginal_ln_likelihood(tmpdir, case): assert len(ll) == len(prior_samples) # save prior samples to a file and pass that instead - filename = str(tmpdir / "samples.hdf5") - prior_samples.write(filename, overwrite=True) + if platform.system() != "Darwin": # this test fails on CI?? + filename = str(tmpdir / "samples.hdf5") + prior_samples.write(filename, overwrite=True) ll = joker.marginal_ln_likelihood(data, filename) assert len(ll) == len(prior_samples) @@ -103,11 +104,10 @@ def test_marginal_ln_likelihood(tmpdir, case): assert len(ll) == len(prior_samples) # NOTE: this makes it so I can't parallelize tests, I think - if platform.system() != "Darwin": # this test fails on CI - with MultiPool(processes=2) as pool: - joker = TheJoker(prior, pool=pool) - ll = joker.marginal_ln_likelihood(data, filename) - assert len(ll) == len(prior_samples) + with MultiPool(processes=2) as pool: + joker = TheJoker(prior, pool=pool) + ll = joker.marginal_ln_likelihood(data, filename) + assert len(ll) == len(prior_samples) priors = [ From e5c07a004cfb4657cfbf335c22242552f6c8d87f Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Wed, 6 Mar 2024 08:32:06 -0500 Subject: [PATCH 42/50] add nbsphinx, duh --- docs/conf.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/conf.py b/docs/conf.py index 8863f4f1..6556df90 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -11,6 +11,7 @@ "sphinx.ext.intersphinx", "sphinx.ext.mathjax", "sphinx.ext.napoleon", + "nbsphinx", # "sphinx_autodoc_typehints", "sphinx_copybutton", "sphinx_automodapi.automodapi", @@ -64,6 +65,9 @@ always_document_param_types = True +# We execute the tutorial notebooks using GitHub Actions and upload to RTD: +nbsphinx_execute = "never" + # The name of your GitHub repository rtds_action_github_repo = "adrn/thejoker" From e5aeee4590f763da222da61402b71a161ddff811 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Wed, 6 Mar 2024 08:32:15 -0500 Subject: [PATCH 43/50] changelog entry --- CHANGES.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGES.rst b/CHANGES.rst index 2fa89f15..51ad025e 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -2,6 +2,8 @@ ---------------- - Fix bug with ``plot_phase_fold()`` that would cause it to fail if ``data=None``. +- Overhaul packaging layout to use more modern practices. +- Replace the ``pymc3`` backend to use the newer ``pymc`` instead. 1.2.2 (2022-03-03) From 940d4f460bd048dfcc94563fc1255c72d1192f41 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Wed, 6 Mar 2024 08:32:21 -0500 Subject: [PATCH 44/50] formatting --- docs/index.rst | 38 ++++++++++++++++++-------------------- docs/install.rst | 7 ++----- 2 files changed, 20 insertions(+), 25 deletions(-) diff --git a/docs/index.rst b/docs/index.rst index 88b4e159..02380364 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -58,17 +58,16 @@ Getting started Generating samples with |thejoker| requires three things: - #. **The data**, `~thejoker.RVData`: radial velocity measurements, + #. **The data**, :class:`thejoker.RVData`: radial velocity measurements, uncertainties, and observation times - #. **The priors**, `~thejoker.JokerPrior`: - the prior distributions over the parameters in *The Joker* - #. **The sampler**, `~thejoker.TheJoker`: the work horse that runs the + #. **The priors**, :class:`thejoker.JokerPrior`: the prior distributions over the + parameters in *The Joker* + #. **The sampler**, :class:`thejoker.TheJoker`: the work horse that runs the rejection sampler -Here, we'll work through a simple example to generate poserior samples for -orbital parameters given some sparse, simulated radial velocity data (shown -below). We'll first use these plain arrays to construct a `~thejoker.RVData` -object: +Here, we'll work through a simple example to generate posterior samples for orbital +parameters given some sparse, simulated radial velocity data (shown below). We'll first +use these plain arrays to construct a `~thejoker.RVData` object: >>> import astropy.units as u >>> import thejoker as tj @@ -94,12 +93,12 @@ object: ax = data.plot() # doctest: +SKIP ax.set_xlim(-10, 200) -We next need to specify the prior distributions for the parameters of -|thejoker|. The default prior, explained in the docstring of -`~thejoker.JokerPrior.default()`, assumes some reasonable defaults where -possible, but requires specifying the minimum and maximum period to sample over, -along with parameters that specify the prior over the linear parameters in *The -Joker* (the velocity semi-amplitude, ``K``, and the systemic velocity, ``v0``): +We next need to specify the prior distributions for the parameters of |thejoker|. The +default prior, explained in the docstring of :meth:`thejoker.JokerPrior.default()`, +assumes some reasonable defaults where possible, but requires specifying the minimum and +maximum period to sample over, along with parameters that specify the prior over the +linear parameters in *The Joker* (the velocity semi-amplitude, ``K``, and the systemic +velocity, ``v0``): >>> import numpy as np >>> prior = tj.JokerPrior.default( @@ -115,12 +114,11 @@ run the rejection sampler: >>> prior_samples = prior.sample(size=100_000) >>> samples = joker.rejection_sample(data, prior_samples) # doctest: +SKIP -Of the 100_000 prior samples we generated, only a handful pass the rejection -sampling step of |thejoker|. Let's visualize the surviving samples in the -subspace of the period :math:`P` and velocity semi-amplitude :math:`K`. We'll -also plot the true values as a green marker. As a separate plot, we'll also -visualize orbits computed from these posterior samples (check the source code -below to see how these were made): +Of the 100,000 prior samples we generated, only a handful pass the rejection sampling +step of |thejoker|. Let's visualize the surviving samples in the subspace of the period +:math:`P` and velocity semi-amplitude :math:`K`. We'll also plot the true values as a +green marker. As a separate plot, we'll also visualize orbits computed from these +posterior samples (check the source code below to see how these were made): .. plot:: :align: center diff --git a/docs/install.rst b/docs/install.rst index a67ee761..8ab55811 100644 --- a/docs/install.rst +++ b/docs/install.rst @@ -2,11 +2,8 @@ Installation ************ -We recommend installing |thejoker| into a new `Anaconda environment -`_. - -pip -=== +With ``pip`` +============ You can install the latest release of The Joker using ``pip``: From 7f42d95dab1a70d01c1874198d30b33ce6644204 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Wed, 6 Mar 2024 08:50:44 -0500 Subject: [PATCH 45/50] theme options --- docs/conf.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/docs/conf.py b/docs/conf.py index 6556df90..f8ba5b68 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -31,6 +31,21 @@ ] html_theme = "pydata_sphinx_theme" +html_theme_options = { + "github_url": "https://github.com/adrn/thejoker", + "use_edit_page_button": True, + "navigation_with_keys": False, +} + +html_context = { + "default_mode": "light", + "to_be_indexed": ["stable", "latest"], + "github_user": "adrn", + "github_repo": "thejoker", + "github_version": "main", + "doc_path": "docs", +} + html_static_path = ["_static"] html_css_files = [ "custom.css", From a8656b11e5862adfd605a66345bb8dadd7100073 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Wed, 6 Mar 2024 09:10:33 -0500 Subject: [PATCH 46/50] change theme and redo index page --- docs/api_docs.rst | 9 +++++++++ docs/conf.py | 42 +++++++++++++++++++++++++----------------- docs/index.rst | 44 ++++++++++---------------------------------- docs/tutorials.rst | 25 +++++++++++++++++++++++++ pyproject.toml | 2 +- 5 files changed, 70 insertions(+), 52 deletions(-) create mode 100644 docs/api_docs.rst create mode 100644 docs/tutorials.rst diff --git a/docs/api_docs.rst b/docs/api_docs.rst new file mode 100644 index 00000000..39ce8ac3 --- /dev/null +++ b/docs/api_docs.rst @@ -0,0 +1,9 @@ +***************** +API Documentation +***************** + +API +=== + +.. automodapi:: thejoker + :no-inheritance-diagram: \ No newline at end of file diff --git a/docs/conf.py b/docs/conf.py index f8ba5b68..36f34fc3 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -18,6 +18,7 @@ "sphinx_automodapi.smart_resolver", "rtds_action", "matplotlib.sphinxext.plot_directive", + "matplotlib.sphinxext.figmpl_directive", ] source_suffix = [".rst", ".md"] @@ -30,27 +31,30 @@ ".venv", ] -html_theme = "pydata_sphinx_theme" +# HTML theme +html_theme = "sphinx_book_theme" +html_copy_source = True +html_show_sourcelink = True +html_sourcelink_suffix = "" +html_title = "thejoker" +html_logo = "_static/thejoker.png" +html_favicon = "_static/icon.ico" +html_static_path = ["_static"] +html_css_files = ["custom.css"] html_theme_options = { - "github_url": "https://github.com/adrn/thejoker", + "path_to_docs": "docs", + "repository_url": "https://github.com/adrn/thejoker", + "repository_branch": "main", + "launch_buttons": { + "binderhub_url": "https://mybinder.org", + "notebook_interface": "classic", + }, "use_edit_page_button": True, - "navigation_with_keys": False, -} - -html_context = { - "default_mode": "light", - "to_be_indexed": ["stable", "latest"], - "github_user": "adrn", - "github_repo": "thejoker", - "github_version": "main", - "doc_path": "docs", + "use_issues_button": True, + "use_repository_button": True, + "use_download_button": True, } -html_static_path = ["_static"] -html_css_files = [ - "custom.css", -] - intersphinx_mapping = { "python": ("https://docs.python.org/3", None), "pandas": ("https://pandas.pydata.org/pandas-docs/stable/", None), @@ -78,6 +82,10 @@ ("py:class", "_io.BytesIO"), ] +plot_srcset = ["2.0x"] # for retina displays +plot_rcparams = {"font.size": 16, "font.family": "serif", "figure.figsize": (6, 4)} +plot_apply_rcparams = True + always_document_param_types = True # We execute the tutorial notebooks using GitHub Actions and upload to RTD: diff --git a/docs/index.rst b/docs/index.rst index 02380364..cd36dd38 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -27,32 +27,10 @@ the method and applications. :maxdepth: 1 install + tutorials + api_docs changes - -Tutorials -========= - -.. toctree:: - :maxdepth: 1 - - examples/1-Getting-started.ipynb - examples/2-Customize-prior.ipynb - examples/3-Polynomial-velocity-trend.ipynb - examples/4-Continue-sampling-mcmc.ipynb - examples/5-Calibration-offsets.ipynb - - -Science demonstrations -====================== - -.. toctree:: - :maxdepth: 1 - - examples/Thompson-black-hole.ipynb - examples/Strader-circular-only.ipynb - - Getting started =============== @@ -82,6 +60,7 @@ use these plain arrays to construct a `~thejoker.RVData` object: :align: center :width: 512 + import matplotlib.pyplot as plt from thejoker.data import RVData import astropy.units as u @@ -90,7 +69,8 @@ use these plain arrays to construct a `~thejoker.RVData` object: err = [0.184, 0.261, 0.112, 0.155, 0.223] * u.km/u.s data = RVData(t=t, rv=rv, rv_err=err) - ax = data.plot() # doctest: +SKIP + fig, ax = plt.subplots(figsize=(6, 4)) + ax = data.plot(ax=ax) ax.set_xlim(-10, 200) We next need to specify the prior distributions for the parameters of |thejoker|. The @@ -112,7 +92,7 @@ run the rejection sampler: >>> joker = tj.TheJoker(prior) >>> prior_samples = prior.sample(size=100_000) - >>> samples = joker.rejection_sample(data, prior_samples) # doctest: +SKIP + >>> samples = joker.rejection_sample(data, prior_samples) Of the 100,000 prior samples we generated, only a handful pass the rejection sampling step of |thejoker|. Let's visualize the surviving samples in the subspace of the period @@ -128,6 +108,7 @@ posterior samples (check the source code below to see how these were made): from thejoker import JokerPrior, TheJoker, RVData from thejoker.plot import plot_rv_curves import astropy.units as u + import numpy as np t = [0., 49.452, 95.393, 127.587, 190.408] rv = [38.77, 39.70, 37.45, 38.31, 38.31] * u.km/u.s @@ -138,8 +119,10 @@ posterior samples (check the source code below to see how these were made): sigma_K0=30*u.km/u.s, sigma_v=100*u.km/u.s) joker = TheJoker(prior) - prior_samples = prior.sample(size=100_000) + rng = np.random.default_rng(seed=42) + prior_samples = prior.sample(size=100_000, rng=rng) samples = joker.rejection_sample(data, prior_samples) + samples = samples.wrap_K() fig, ax = plt.subplots(1, 1, figsize=(6, 4)) # doctest: +SKIP ax.scatter(samples['P'].value, samples['K'].to(u.km/u.s).value, @@ -158,13 +141,6 @@ posterior samples (check the source code below to see how these were made): ax.set_xlim(-5, 205) -API -=== - -.. automodapi:: thejoker - :no-inheritance-diagram: - - .. rubric:: Footnotes .. [#f1] Short for Johannes Kepler. diff --git a/docs/tutorials.rst b/docs/tutorials.rst new file mode 100644 index 00000000..f4d2529b --- /dev/null +++ b/docs/tutorials.rst @@ -0,0 +1,25 @@ +************************************ +Tutorials and Science Demonstrations +************************************ + +Tutorials +========= + +.. toctree:: + :maxdepth: 1 + + examples/1-Getting-started.ipynb + examples/2-Customize-prior.ipynb + examples/3-Polynomial-velocity-trend.ipynb + examples/4-Continue-sampling-mcmc.ipynb + examples/5-Calibration-offsets.ipynb + + +Science demonstrations +====================== + +.. toctree:: + :maxdepth: 1 + + examples/Thompson-black-hole.ipynb + examples/Strader-circular-only.ipynb \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 737d22eb..7550d851 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,7 +57,7 @@ docs = [ "matplotlib", "sphinx_copybutton", "rtds_action", - "pydata-sphinx-theme" + "sphinx_book_theme" ] tutorials = [ "thejoker[docs]", From 644f887f2b5aade96ed169ae0f3c81ebab914e89 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Wed, 6 Mar 2024 09:14:07 -0500 Subject: [PATCH 47/50] use rng when sampling from prior and add regression test --- thejoker/prior.py | 4 +--- thejoker/tests/test_prior.py | 14 ++++++++++++++ 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/thejoker/prior.py b/thejoker/prior.py index 01253e51..3d2eaa62 100644 --- a/thejoker/prior.py +++ b/thejoker/prior.py @@ -16,7 +16,6 @@ validate_poly_trend, validate_sigma_v, ) -from .utils import rng_context __all__ = ["JokerPrior"] @@ -354,8 +353,7 @@ def sample( par_names = list(sub_pars.keys()) par_list = [sub_pars[k] for k in par_names] - with rng_context(rng): - samples_values = pm.draw(par_list, draws=size) + samples_values = pm.draw(par_list, draws=size, random_seed=rng) raw_samples = { name: samples.astype(dtype) diff --git a/thejoker/tests/test_prior.py b/thejoker/tests/test_prior.py index 8dc62fdf..ff6efed2 100644 --- a/thejoker/tests/test_prior.py +++ b/thejoker/tests/test_prior.py @@ -210,3 +210,17 @@ def test_dtype(): samples = prior.sample(size=100, dtype=np.float32) for name in samples.par_names: assert samples[name].dtype == np.float32 + + +def test_rng_prior_sample(): + """Check that passing in a random number generator works as expected""" + + prior, _ = get_prior(0) + + rng = np.random.default_rng(seed=42) + prior_samples1 = prior.sample(size=100, rng=rng) + + rng = np.random.default_rng(seed=42) + prior_samples2 = prior.sample(size=100, rng=rng) + + assert np.allclose(prior_samples1["P"], prior_samples2["P"]) From d65fc945d314a26bcd6a84dbc393354e832e329d Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Wed, 6 Mar 2024 09:17:26 -0500 Subject: [PATCH 48/50] disable test on macos CI --- thejoker/tests/test_sampler.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/thejoker/tests/test_sampler.py b/thejoker/tests/test_sampler.py index a95930d3..1458195d 100644 --- a/thejoker/tests/test_sampler.py +++ b/thejoker/tests/test_sampler.py @@ -91,23 +91,23 @@ def test_marginal_ln_likelihood(tmpdir, case): ll = joker.marginal_ln_likelihood(data, prior_samples) assert len(ll) == len(prior_samples) - # save prior samples to a file and pass that instead - if platform.system() != "Darwin": # this test fails on CI?? + if platform.system() != "Darwin": # this part of the test fails on CI macos?? + # save prior samples to a file and pass that instead filename = str(tmpdir / "samples.hdf5") prior_samples.write(filename, overwrite=True) - ll = joker.marginal_ln_likelihood(data, filename) - assert len(ll) == len(prior_samples) + ll = joker.marginal_ln_likelihood(data, filename) + assert len(ll) == len(prior_samples) - # make sure batches work: - ll = joker.marginal_ln_likelihood(data, filename, n_batches=10) - assert len(ll) == len(prior_samples) + # make sure batches work: + ll = joker.marginal_ln_likelihood(data, filename, n_batches=10) + assert len(ll) == len(prior_samples) - # NOTE: this makes it so I can't parallelize tests, I think - with MultiPool(processes=2) as pool: - joker = TheJoker(prior, pool=pool) - ll = joker.marginal_ln_likelihood(data, filename) - assert len(ll) == len(prior_samples) + # NOTE: this makes it so I can't parallelize tests, I think + with MultiPool(processes=2) as pool: + joker = TheJoker(prior, pool=pool) + ll = joker.marginal_ln_likelihood(data, filename) + assert len(ll) == len(prior_samples) priors = [ From ebb2030423a3727849c62904de8a3e021697082d Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Wed, 6 Mar 2024 09:30:42 -0500 Subject: [PATCH 49/50] tight layout --- docs/index.rst | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/index.rst b/docs/index.rst index cd36dd38..f3ae9ef8 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -69,7 +69,7 @@ use these plain arrays to construct a `~thejoker.RVData` object: err = [0.184, 0.261, 0.112, 0.155, 0.223] * u.km/u.s data = RVData(t=t, rv=rv, rv_err=err) - fig, ax = plt.subplots(figsize=(6, 4)) + fig, ax = plt.subplots(figsize=(6, 4), layout="tight") ax = data.plot(ax=ax) ax.set_xlim(-10, 200) @@ -124,9 +124,9 @@ posterior samples (check the source code below to see how these were made): samples = joker.rejection_sample(data, prior_samples) samples = samples.wrap_K() - fig, ax = plt.subplots(1, 1, figsize=(6, 4)) # doctest: +SKIP + fig, ax = plt.subplots(1, 1, figsize=(6, 4), layout="tight") ax.scatter(samples['P'].value, samples['K'].to(u.km/u.s).value, - marker='.', color='k', alpha=0.45) # doctest: +SKIP + marker='.', color='k', alpha=0.45) ax.set_xlabel("$P$ [day]") ax.set_ylabel("$K$ [km/s]") ax.set_xlim(0, 256) @@ -134,7 +134,7 @@ posterior samples (check the source code below to see how these were made): ax.scatter(61.942, 1.3959, marker='o', color='#31a354', zorder=-100) - fig, ax = plt.subplots(1, 1, figsize=(6, 4)) # doctest: +SKIP + fig, ax = plt.subplots(1, 1, figsize=(6, 4), layout="tight") t_grid = np.linspace(-10, 210, 1024) plot_rv_curves(samples, t_grid, rv_unit=u.km/u.s, data=data, ax=ax, plot_kwargs=dict(color='#888888')) From 3516c310b30e447b0762f9877480132a30c2be56 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Wed, 6 Mar 2024 09:56:08 -0500 Subject: [PATCH 50/50] try pinning pytables --- pyproject.toml | 2 +- thejoker/tests/test_sampler.py | 28 +++++++++++++--------------- 2 files changed, 14 insertions(+), 16 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 7550d851..d10a153b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,7 @@ dependencies = [ "pymc>=5", "pymc_ext @ git+https://github.com/exoplanet-dev/pymc-ext", "exoplanet-core[pymc] @ git+https://github.com/exoplanet-dev/exoplanet-core", - "tables", + "tables<3.9.2", # https://github.com/PyTables/PyTables/issues/1093 "dill" ] diff --git a/thejoker/tests/test_sampler.py b/thejoker/tests/test_sampler.py index 1458195d..e16418ff 100644 --- a/thejoker/tests/test_sampler.py +++ b/thejoker/tests/test_sampler.py @@ -1,5 +1,4 @@ import os -import platform import astropy.units as u import numpy as np @@ -91,23 +90,22 @@ def test_marginal_ln_likelihood(tmpdir, case): ll = joker.marginal_ln_likelihood(data, prior_samples) assert len(ll) == len(prior_samples) - if platform.system() != "Darwin": # this part of the test fails on CI macos?? - # save prior samples to a file and pass that instead - filename = str(tmpdir / "samples.hdf5") - prior_samples.write(filename, overwrite=True) + # save prior samples to a file and pass that instead + filename = str(tmpdir / "samples.hdf5") + prior_samples.write(filename, overwrite=True) - ll = joker.marginal_ln_likelihood(data, filename) - assert len(ll) == len(prior_samples) + ll = joker.marginal_ln_likelihood(data, filename) + assert len(ll) == len(prior_samples) - # make sure batches work: - ll = joker.marginal_ln_likelihood(data, filename, n_batches=10) - assert len(ll) == len(prior_samples) + # make sure batches work: + ll = joker.marginal_ln_likelihood(data, filename, n_batches=10) + assert len(ll) == len(prior_samples) - # NOTE: this makes it so I can't parallelize tests, I think - with MultiPool(processes=2) as pool: - joker = TheJoker(prior, pool=pool) - ll = joker.marginal_ln_likelihood(data, filename) - assert len(ll) == len(prior_samples) + # NOTE: this makes it so I can't parallelize tests, I think + with MultiPool(processes=2) as pool: + joker = TheJoker(prior, pool=pool) + ll = joker.marginal_ln_likelihood(data, filename) + assert len(ll) == len(prior_samples) priors = [