├── .gitattributes ├── .github └── workflows │ └── python-app.yml ├── .gitignore ├── CHANGES.md ├── LICENSE ├── README.md ├── docs ├── Makefile ├── make.bat └── source │ ├── api.md │ ├── bank.md │ ├── conf.py │ ├── constants.md │ ├── index.md │ ├── installation.md │ ├── metric.md │ ├── noise.md │ ├── usage.md │ ├── utils.md │ └── waveforms.md ├── plot_style.mplstyle ├── pyproject.toml ├── requirements.txt ├── scripts ├── README.md ├── banks │ ├── 3pn-random-1-mm=0.8-eta_star=0.8-n_eff=10.npz │ ├── 3pn-random-1-mm=0.95-eta_star=0.9-n_eff=1000.npz │ ├── 3pn-random-100-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-101-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-102-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-103-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-104-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-105-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-106-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-107-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-108-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-109-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-110-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-111-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-112-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-113-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-114-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-115-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-116-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-117-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-118-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-120-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-121-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-122-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-123-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-124-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-125-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-126-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-127-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-128-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-129-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-130-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── 3pn-random-501-mm=0.95-eta_star=0.9-n_eff=10.npz │ ├── 3pn-random-502-mm=0.95-eta_star=0.9-n_eff=11.npz │ ├── 3pn-random-503-mm=0.95-eta_star=0.9-n_eff=13.npz │ ├── 3pn-random-504-mm=0.95-eta_star=0.9-n_eff=15.npz │ ├── 3pn-random-505-mm=0.95-eta_star=0.9-n_eff=17.npz │ ├── 3pn-random-506-mm=0.95-eta_star=0.9-n_eff=20.npz │ ├── 3pn-random-507-mm=0.95-eta_star=0.9-n_eff=23.npz │ ├── 3pn-random-508-mm=0.95-eta_star=0.9-n_eff=26.npz │ ├── 3pn-random-509-mm=0.95-eta_star=0.9-n_eff=30.npz │ ├── 3pn-random-510-mm=0.95-eta_star=0.9-n_eff=35.npz │ ├── 3pn-random-511-mm=0.95-eta_star=0.9-n_eff=40.npz │ ├── 3pn-random-512-mm=0.95-eta_star=0.9-n_eff=47.npz │ ├── 3pn-random-513-mm=0.95-eta_star=0.9-n_eff=54.npz │ ├── 3pn-random-514-mm=0.95-eta_star=0.9-n_eff=62.npz │ ├── 3pn-random-515-mm=0.95-eta_star=0.9-n_eff=71.npz │ ├── 3pn-random-516-mm=0.95-eta_star=0.9-n_eff=82.npz │ ├── 3pn-random-517-mm=0.95-eta_star=0.9-n_eff=95.npz │ ├── 3pn-random-518-mm=0.95-eta_star=0.9-n_eff=109.npz │ ├── 3pn-random-519-mm=0.95-eta_star=0.9-n_eff=126.npz │ ├── 3pn-random-520-mm=0.95-eta_star=0.9-n_eff=145.npz │ ├── 3pn-random-521-mm=0.95-eta_star=0.9-n_eff=167.npz │ ├── 3pn-random-522-mm=0.95-eta_star=0.9-n_eff=193.npz │ ├── 3pn-random-523-mm=0.95-eta_star=0.9-n_eff=222.npz │ ├── 3pn-random-524-mm=0.95-eta_star=0.9-n_eff=255.npz │ ├── 3pn-random-525-mm=0.95-eta_star=0.9-n_eff=294.npz │ ├── 3pn-random-526-mm=0.95-eta_star=0.9-n_eff=339.npz │ ├── 3pn-random-527-mm=0.95-eta_star=0.9-n_eff=390.npz │ ├── 3pn-random-528-mm=0.95-eta_star=0.9-n_eff=449.npz │ ├── 3pn-random-529-mm=0.95-eta_star=0.9-n_eff=517.npz │ ├── 3pn-random-530-mm=0.95-eta_star=0.9-n_eff=596.npz │ ├── 3pn-random-531-mm=0.95-eta_star=0.9-n_eff=686.npz │ ├── 3pn-random-532-mm=0.95-eta_star=0.9-n_eff=790.npz │ ├── 3pn-random-533-mm=0.95-eta_star=0.9-n_eff=910.npz │ ├── 3pn-random-534-mm=0.95-eta_star=0.9-n_eff=1048.npz │ ├── 3pn-random-535-mm=0.95-eta_star=0.9-n_eff=1206.npz │ ├── 3pn-random-536-mm=0.95-eta_star=0.9-n_eff=1389.npz │ ├── 3pn-random-537-mm=0.95-eta_star=0.9-n_eff=1599.npz │ ├── 3pn-random-538-mm=0.95-eta_star=0.9-n_eff=1842.npz │ ├── 3pn-random-539-mm=0.95-eta_star=0.9-n_eff=2120.npz │ ├── 3pn-random-540-mm=0.95-eta_star=0.9-n_eff=2442.npz │ ├── 3pn-random-541-mm=0.95-eta_star=0.9-n_eff=2811.npz │ ├── 3pn-random-542-mm=0.95-eta_star=0.9-n_eff=3237.npz │ ├── 3pn-random-543-mm=0.95-eta_star=0.9-n_eff=3727.npz │ ├── 3pn-random-544-mm=0.95-eta_star=0.9-n_eff=4291.npz │ ├── 3pn-random-545-mm=0.95-eta_star=0.9-n_eff=4941.npz │ ├── 3pn-random-546-mm=0.95-eta_star=0.9-n_eff=5689.npz │ ├── 3pn-random-547-mm=0.95-eta_star=0.9-n_eff=6551.npz │ ├── 3pn-random-548-mm=0.95-eta_star=0.9-n_eff=7543.npz │ ├── 3pn-random-549-mm=0.95-eta_star=0.9-n_eff=8685.npz │ ├── 3pn-random-550-mm=0.95-eta_star=0.9-n_eff=10000.npz │ ├── 3pn-stochastic-1000-mm=0.95-eta_star=0.9-n_eff=500.npz │ ├── README.md │ ├── tf2-random-1-mm=0.96-eta_star=0.9-n_eff=1000.npz │ └── tf2rs-random-1-mm=0.95-eta_star=0.993-n_eff=1300.npz ├── calc_bank_eff.py ├── figures │ ├── .gitignore │ ├── README.md │ ├── bank-effs.pdf │ ├── density.pdf │ ├── eta-dist.pdf │ ├── neff-scaling.pdf │ ├── scalar-curvature.pdf │ ├── scaling.pdf │ ├── taylorf2reducedspin_density.pdf │ └── threePN-eff-cdfs.pdf ├── genbank_2D_threePN.py ├── genbank_3D_taylorf2reducedspin.py ├── genbank_4D_taylorf2.py ├── job-2D-threePN.sh ├── job-3D-taylorfsreducedspin.sh ├── job-4D-taylorf2.sh ├── job-threePN-coverage.sh ├── job-threePN-est-p-variation.sh ├── job-threePN-scaling.sh ├── minimal_example.py ├── plot_bank_effs.py ├── plot_density.py ├── plot_eta_dist.py ├── plot_neff_scaling.py ├── plot_scalar_curvature.py ├── plot_scaling.py ├── threePN-banks-scaling │ ├── 3pn-random-15-mm=0.9-eta_star=0.975-n_eff=1000.npz │ ├── 3pn-random-16-mm=0.9-eta_star=0.95-n_eff=1000.npz │ ├── 3pn-random-17-mm=0.9-eta_star=0.925-n_eff=1000.npz │ ├── 3pn-random-18-mm=0.9-eta_star=0.9-n_eff=1000.npz │ ├── 3pn-random-19-mm=0.9-eta_star=0.875-n_eff=1000.npz │ ├── 3pn-random-20-mm=0.9-eta_star=0.85-n_eff=1000.npz │ ├── 3pn-random-5-mm=0.95-eta_star=0.9-n_eff=1000.npz │ ├── 3pn-random-6-mm=0.9-eta_star=0.9-n_eff=1000.npz │ ├── 3pn-random-7-mm=0.85-eta_star=0.9-n_eff=1000.npz │ ├── 3pn-random-8-mm=0.8-eta_star=0.9-n_eff=1000.npz │ ├── 3pn-random-9-mm=0.75-eta_star=0.9-n_eff=1000.npz │ ├── 3pn-stochastic-10-mm=0.95-eta_star=0.9-n_eff=1000.npz │ ├── 3pn-stochastic-11-mm=0.9-eta_star=0.9-n_eff=1000.npz │ ├── 3pn-stochastic-12-mm=0.85-eta_star=0.9-n_eff=1000.npz │ ├── 3pn-stochastic-13-mm=0.8-eta_star=0.9-n_eff=1000.npz │ ├── 3pn-stochastic-14-mm=0.75-eta_star=0.9-n_eff=1000.npz │ ├── 3pn-stochastic-21-mm=0.9-eta_star=0.975-n_eff=1000.npz │ ├── 3pn-stochastic-22-mm=0.9-eta_star=0.95-n_eff=1000.npz │ ├── 3pn-stochastic-23-mm=0.9-eta_star=0.925-n_eff=1000.npz │ ├── 3pn-stochastic-24-mm=0.9-eta_star=0.9-n_eff=1000.npz │ ├── 3pn-stochastic-25-mm=0.9-eta_star=0.875-n_eff=1000.npz │ └── 3pn-stochastic-26-mm=0.9-eta_star=0.85-n_eff=1000.npz ├── threePN-outputs-scaling │ ├── threePN-random-15-mm=0.90-eta_star=0.975.txt │ ├── threePN-random-16-mm=0.90-eta_star=0.950.txt │ ├── threePN-random-17-mm=0.90-eta_star=0.925.txt │ ├── threePN-random-18-mm=0.90-eta_star=0.900.txt │ ├── threePN-random-19-mm=0.90-eta_star=0.875.txt │ ├── threePN-random-20-mm=0.90-eta_star=0.850.txt │ ├── threePN-random-5-mm=0.95-eta_star=0.9.txt │ ├── threePN-random-6-mm=0.90-eta_star=0.9.txt │ ├── threePN-random-7-mm=0.85-eta_star=0.9.txt │ ├── threePN-random-8-mm=0.80-eta_star=0.9.txt │ ├── threePN-random-9-mm=0.75-eta_star=0.9.txt │ ├── threePN-stochastic-10-mm=0.95-eta_star=0.9.txt │ ├── threePN-stochastic-11-mm=0.90-eta_star=0.9.txt │ ├── threePN-stochastic-12-mm=0.85-eta_star=0.9.txt │ ├── threePN-stochastic-13-mm=0.80-eta_star=0.9.txt │ ├── threePN-stochastic-14-mm=0.75-eta_star=0.9.txt │ ├── threePN-stochastic-21-mm=0.90-eta_star=0.975.txt │ ├── threePN-stochastic-22-mm=0.90-eta_star=0.950.txt │ ├── threePN-stochastic-23-mm=0.90-eta_star=0.925.txt │ ├── threePN-stochastic-24-mm=0.90-eta_star=0.900.txt │ ├── threePN-stochastic-25-mm=0.90-eta_star=0.875.txt │ └── threePN-stochastic-26-mm=0.90-eta_star=0.850.txt ├── threePN-p.txt ├── threePN_est_p.py └── threePN_est_p_variation.py ├── setup.cfg ├── setup.py ├── src └── diffbank │ ├── __init__.py │ ├── bank.py │ ├── constants.py │ ├── metric.py │ ├── noise.py │ ├── noise_resources │ ├── O2_ASD.dat │ ├── O3a_Livingston_ASD.dat │ ├── __init__.py │ └── aLIGOZeroDetHighPower.dat │ ├── utils.py │ └── waveforms │ ├── __init__.py │ ├── taylorF2.py │ ├── taylorf2reducedspin.py │ ├── threePN_simple.py │ ├── twoPN_chirptimes.py │ └── twoPN_simple.py └── tests ├── README.md ├── test_bank.py ├── test_checkpointing.py ├── test_eff.py ├── test_est_ratio_max.py ├── test_genbank_2D_threePN.py ├── test_metric_speed.py └── test_taylorf2_metric.py /.gitattributes: -------------------------------------------------------------------------------- 1 | scripts/banks/*.npz filter=lfs diff=lfs merge=lfs -text 2 | -------------------------------------------------------------------------------- /.github/workflows/python-app.yml: -------------------------------------------------------------------------------- 1 | # This workflow will install Python dependencies and run tests with a single version of Python 2 | # For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions 3 | 4 | name: Python application 5 | 6 | on: 7 | push: 8 | branches: [ main ] 9 | pull_request: 10 | branches: [ main ] 11 | 12 | jobs: 13 | build: 14 | 15 | runs-on: ubuntu-latest 16 | 17 | steps: 18 | - uses: actions/checkout@v2 19 | - name: Set up Python 3.10 20 | uses: actions/setup-python@v2 21 | with: 22 | python-version: "3.10" 23 | - name: Install dependencies 24 | run: | 25 | python -m pip install --upgrade pip 26 | pip install pytest 27 | if [ -f requirements.txt ]; then pip install -r requirements.txt; fi 28 | - name: Install diffbank 29 | run: | 30 | pip install . 31 | - name: Test with pytest 32 | run: | 33 | cd tests/ 34 | pytest 35 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | .ipynb_checkpoints/ 3 | .python-version 4 | *__pycache__ 5 | *.egg-info 6 | tests/*.png 7 | tests/*.npz 8 | *slurm-* 9 | docs/build/* 10 | -------------------------------------------------------------------------------- /CHANGES.md: -------------------------------------------------------------------------------- 1 | # v0.0.1 2 | 3 | First code release to go with paper. 4 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) [2021] [Adam Coogan, Thomas Edwards] 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # diffbank 2 | 3 | [![docs](https://readthedocs.org/projects/diffbank/badge/?version=latest)](http://diffbank.readthedocs.io/?badge=latest) 4 | ![tests](https://github.com/adam-coogan/diffbank/actions/workflows/python-app.yml/badge.svg) 5 | [![arXiv](https://img.shields.io/badge/arXiv-2202.09380-b31b1b.svg)](https://arxiv.org/abs/2202.09380) 6 | 7 | 8 | Gravitational wave template bank generation made easy. For details, check out the 9 | [documentation](https://diffbank.readthedocs.io/en/latest/). 10 | 11 | ## Installation 12 | 13 | Clone the repo, `cd` to this directory, and install with `pip install .` (or 14 | `pip install -e .` for an editable install). You'll need [`git-lfs`](https://git-lfs.github.com/) 15 | if you want to download the banks in the repo. 16 | 17 | ## Citing 18 | 19 | If you use `diffbank` in your work, please cite [our paper](https://arxiv.org/abs/2202.09380): 20 | 21 | ```text 22 | @article{Coogan:2022qxs, 23 | author = "Coogan, Adam and Edwards, Thomas D. P. and Chia, Horng Sheng and 24 | George, Richard N. and Freese, Katherine and Messick, Cody and 25 | Setzer, Christian N. and Weniger, Christoph and Zimmerman, Aaron", 26 | title = "{Efficient Template Bank Generation with Differentiable Waveforms}", 27 | eprint = "2202.09380", 28 | archivePrefix = "arXiv", 29 | primaryClass = "astro-ph.IM", 30 | reportNumber = "UTTG 26-2021, NORDITA-2022-004", 31 | month = "2", 32 | year = "2022" 33 | } 34 | ``` 35 | -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- 1 | # Minimal makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line, and also 5 | # from the environment for the first two. 6 | SPHINXOPTS ?= 7 | SPHINXBUILD ?= sphinx-build 8 | SOURCEDIR = source 9 | BUILDDIR = build 10 | 11 | # Put it first so that "make" without argument is like "make help". 12 | help: 13 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 14 | 15 | .PHONY: help Makefile 16 | 17 | # Catch-all target: route all unknown targets to Sphinx using the new 18 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). 19 | %: Makefile 20 | @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 21 | -------------------------------------------------------------------------------- /docs/make.bat: -------------------------------------------------------------------------------- 1 | @ECHO OFF 2 | 3 | pushd %~dp0 4 | 5 | REM Command file for Sphinx documentation 6 | 7 | if "%SPHINXBUILD%" == "" ( 8 | set SPHINXBUILD=sphinx-build 9 | ) 10 | set SOURCEDIR=source 11 | set BUILDDIR=build 12 | 13 | if "%1" == "" goto help 14 | 15 | %SPHINXBUILD% >NUL 2>NUL 16 | if errorlevel 9009 ( 17 | echo. 18 | echo.The 'sphinx-build' command was not found. Make sure you have Sphinx 19 | echo.installed, then set the SPHINXBUILD environment variable to point 20 | echo.to the full path of the 'sphinx-build' executable. Alternatively you 21 | echo.may add the Sphinx directory to PATH. 22 | echo. 23 | echo.If you don't have Sphinx installed, grab it from 24 | echo.https://www.sphinx-doc.org/ 25 | exit /b 1 26 | ) 27 | 28 | %SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 29 | goto end 30 | 31 | :help 32 | %SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 33 | 34 | :end 35 | popd 36 | -------------------------------------------------------------------------------- /docs/source/api.md: -------------------------------------------------------------------------------- 1 | # API 2 | 3 | The {class}`diffbank.bank.Bank` class provides a simple interface to access all 4 | the functionality in `diffbank`. If you prefer a more functional interface, you 5 | can also use the functions in the {mod}`diffbank.metric` and {mod}`diffbank.utils` 6 | modules. 7 | 8 | A few waveform and detector noise models are included in the {mod}`diffbank.waveforms` 9 | and {mod}`diffbank.noise` modules. Finally, the {mod}`diffbank.constants` module 10 | defines a few physical constants. Note that `diffbank` uses SI units. 11 | 12 | ## Modules 13 | 14 | ```{eval-rst} 15 | .. toctree:: 16 | :maxdepth: 1 17 | 18 | bank 19 | constants 20 | metric 21 | noise 22 | utils 23 | waveforms 24 | ``` 25 | -------------------------------------------------------------------------------- /docs/source/bank.md: -------------------------------------------------------------------------------- 1 | # diffbank.bank.Bank 2 | 3 | ```{eval-rst} 4 | .. autoclass:: diffbank.bank.Bank 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | ``` 9 | -------------------------------------------------------------------------------- /docs/source/conf.py: -------------------------------------------------------------------------------- 1 | # Configuration file for the Sphinx documentation builder. 2 | # 3 | # This file only contains a selection of the most common options. For a full 4 | # list see the documentation: 5 | # https://www.sphinx-doc.org/en/master/usage/configuration.html 6 | 7 | # -- Path setup -------------------------------------------------------------- 8 | 9 | # If extensions (or modules to document with autodoc) are in another directory, 10 | # add these directories to sys.path here. If the directory is relative to the 11 | # documentation root, use os.path.abspath to make it absolute, like shown here. 12 | # 13 | import os 14 | import sys 15 | 16 | sys.path.insert(0, os.path.abspath(os.path.join("..", "..", "src"))) 17 | 18 | 19 | # -- Project information ----------------------------------------------------- 20 | 21 | project = "diffbank" 22 | copyright = "2022, Adam Coogan, Thomas Edwards" 23 | author = "Adam Coogan, Thomas Edwards" 24 | 25 | # The full version, including alpha/beta/rc tags 26 | release = "0.0.1" 27 | 28 | 29 | # -- General configuration --------------------------------------------------- 30 | 31 | # Add any Sphinx extension module names here, as strings. They can be 32 | # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom 33 | # ones. 34 | extensions = [ 35 | "myst_parser", 36 | "sphinx.ext.autodoc", 37 | "sphinx.ext.napoleon", 38 | "sphinx.ext.linkcode", 39 | "sphinx.ext.mathjax", 40 | "sphinx_autodoc_typehints", 41 | ] 42 | 43 | 44 | def linkcode_resolve(domain, info): 45 | if domain != "py": 46 | return None 47 | if not info["module"]: 48 | return None 49 | filename = info["module"].replace(".", "/") 50 | return f"https://github.com/adam-coogan/diffbank/blob/main/src/{filename}.py" 51 | 52 | 53 | # Add any paths that contain templates here, relative to this directory. 54 | templates_path = ["_templates"] 55 | 56 | # List of patterns, relative to source directory, that match files and 57 | # directories to ignore when looking for source files. 58 | # This pattern also affects html_static_path and html_extra_path. 59 | exclude_patterns = [] 60 | 61 | 62 | # -- Options for HTML output ------------------------------------------------- 63 | 64 | # The theme to use for HTML and HTML Help pages. See the documentation for 65 | # a list of builtin themes. 66 | # 67 | # html_theme = 'alabaster' 68 | # html_theme = "sphinx_rtd_theme" 69 | html_theme = "furo" 70 | 71 | # Add any paths that contain custom static files (such as style sheets) here, 72 | # relative to this directory. They are copied after the builtin static files, 73 | # so a file named "default.css" will overwrite the builtin "default.css". 74 | html_static_path = ["_static"] 75 | -------------------------------------------------------------------------------- /docs/source/constants.md: -------------------------------------------------------------------------------- 1 | # diffbank.constants 2 | 3 | ```{eval-rst} 4 | .. automodule:: diffbank.constants 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | ``` 9 | -------------------------------------------------------------------------------- /docs/source/index.md: -------------------------------------------------------------------------------- 1 | # diffbank 2 | 3 | `diffbank` makes creating template banks for gravitational wave searches simple 4 | and fast. In a nutshell, it combines differentiable waveforms based on the 5 | [jax](https://github.com/google/jax/) framework with a new variant on [random 6 | template bank generation](https://arxiv.org/abs/0809.5223). Read through the links 7 | below to get started. 8 | 9 | ```{eval-rst} 10 | .. toctree:: 11 | :maxdepth: 1 12 | 13 | installation 14 | usage 15 | api 16 | ``` 17 | 18 | ## Citing 19 | 20 | If you use `diffbank` in your work, please cite [our paper](https://arxiv.org/abs/2202.09380): 21 | 22 | ```text 23 | @article{Coogan:2022qxs, 24 | author = "Coogan, Adam and Edwards, Thomas D. P. and Chia, Horng Sheng and 25 | George, Richard N. and Freese, Katherine and Messick, Cody and 26 | Setzer, Christian N. and Weniger, Christoph and Zimmerman, Aaron", 27 | title = "{Efficient Template Bank Generation with Differentiable Waveforms}", 28 | eprint = "2202.09380", 29 | archivePrefix = "arXiv", 30 | primaryClass = "astro-ph.IM", 31 | reportNumber = "UTTG 26-2021, NORDITA-2022-004", 32 | month = "2", 33 | year = "2022" 34 | } 35 | ``` 36 | -------------------------------------------------------------------------------- /docs/source/installation.md: -------------------------------------------------------------------------------- 1 | # Installation 2 | 3 | `diffbank` requires python 3.6 or newer. `diffbank` is not yet on PyPI. For now, 4 | to install, clone the repository and run `pip`: 5 | 6 | ```bash 7 | git clone https://github.com/adam-coogan/diffbank 8 | cd diffbank 9 | pip install . 10 | ``` 11 | 12 | This will install the dependencies: `jax`, `numpy` and `tqdm`. The scripts, docs, 13 | tests and notebooks have additional dependencies. These can be installed by running 14 | 15 | ```bash 16 | pip install -r requirements.txt 17 | ``` 18 | 19 | from the root of the repository. You'll also need [`git-lfs`](https://git-lfs.github.com/) 20 | if you want to download the banks in the repository. 21 | -------------------------------------------------------------------------------- /docs/source/metric.md: -------------------------------------------------------------------------------- 1 | # diffbank.metric 2 | 3 | ```{eval-rst} 4 | .. automodule:: diffbank.metric 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | ``` 9 | -------------------------------------------------------------------------------- /docs/source/noise.md: -------------------------------------------------------------------------------- 1 | # diffbank.noise 2 | 3 | ```{eval-rst} 4 | .. automodule:: diffbank.noise 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | ``` 9 | -------------------------------------------------------------------------------- /docs/source/usage.md: -------------------------------------------------------------------------------- 1 | # Usage 2 | 3 | TL;DR: check out the [`minimal_example.py`](https://github.com/adam-coogan/diffbank/blob/main/scripts/minimal_example.py) 4 | script. 5 | 6 | There are two interfaces to `diffbank`: a set of bank generation functions and 7 | the convenient {class}`diffbank.bank.Bank` class wrapper. The first interface is 8 | more in keeping with `jax`'s functional approach, but we expect other users will 9 | find the class interface more convenient. This page will show you how to use both. 10 | 11 | Take a look at our [`genbank_*.py` scripts](https://github.com/adam-coogan/diffbank/tree/main/scripts) 12 | for more examples. 13 | 14 | ## Waveform model and setup 15 | 16 | Let's start by importing some `jax` modules and setting the random seed: 17 | 18 | ```python 19 | import jax.numpy as jnp 20 | from jax import random 21 | 22 | key = random.PRNGKey(526) 23 | ``` 24 | 25 | If you're unfamiliar with `jax` random number generation, take a look [here](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#random-numbers). 26 | It's a bit different than what you may be used to from e.g. `numpy`. 27 | 28 | We will use a simple 3.5PN waveform as the signal model, taking the component 29 | black hole masses to lie between 1 and 3 solar masses. 30 | 31 | ```python 32 | from diffbank.waveforms.threePN_simple import Psi, amp 33 | 34 | m_range = (2.0, 3.0) 35 | ``` 36 | 37 | Our bank generation scheme samples from the parameter space using the metric density. 38 | This is done using rejection sampling, which means the user must provide a base 39 | distribution, defined in terms of a sampler over the parameter space. Our base 40 | sampler draws uniformly from the `(m1, m2)` parameter space, with the restriction 41 | `m1 > m2`. This can be done using a convenience function from {mod}`diffbank.utils`: 42 | 43 | ```python 44 | from diffbank.utils import get_m1_m2_sampler 45 | 46 | sampler = get_m1_m2_sampler(m_range, m_range) 47 | ``` 48 | 49 | Using a nonuniform base density necessitates defining its density function as well. 50 | This density function need not be normalized. Here we will take it to be 1. 51 | 52 | Next we define a noise model. We will use an analytic version of the LIGO-I noise 53 | power spectral density: 54 | 55 | ```python 56 | from diffbank.noise import Sn_LIGOI as Sn 57 | ``` 58 | 59 | We also need a frequency grid to use for match calculations, in units of Hz: 60 | 61 | ```python 62 | fs = jnp.linspace(24.0, 512.0, 4880) 63 | ``` 64 | 65 | Lastly, bank generation requires setting the maximum mismatch `m_star`, target 66 | parameter space coverage fraction `eta` and number of effectualness points to use 67 | for convergence monitoring: 68 | 69 | ```python 70 | minimum_match = 0.95 71 | m_star = 1 - minimum_match 72 | eta = 0.9 73 | n_eff = 1000 74 | ``` 75 | 76 | ## `Bank` class interface 77 | 78 | Initializing a bank requires the definitions from above and a name: 79 | 80 | ```python 81 | from diffbank.bank import Bank 82 | 83 | bank = Bank(amp, Psi, fs, Sn, m_star, eta_star, sampler, name=f"3pn-bank") 84 | ``` 85 | 86 | Before we can generate the bank, we need to compute the maximum value of the ratio 87 | between the metric density and the base sampler density. For this waveform model 88 | we can easily find this point through numerical optimization, but it's more convenient 89 | to estimate it with [empirical supremum rejection sampling](https://bookdown.org/rdpeng/advstatcomp/rejection-sampling.html#empirical-supremum-rejection-sampling). 90 | The `est_ratio_max` method returns the estimated maximum ratio and the point at 91 | which it was attained, so we need only keep the first return value: 92 | 93 | ```python 94 | key, subkey = random.split(key) 95 | bank.ratio_max = bank.est_ratio_max(subkey)[0] 96 | ``` 97 | 98 | Now we can fill our bank! Let's use the random bank generation scheme from our paper: 99 | 100 | ```python 101 | key, subkey = random.split(key) 102 | bank.fill_bank(subkey, "random", n_eff) 103 | ``` 104 | 105 | This will print a [`tqdm`](https://github.com/tqdm/tqdm) progress bar to monitor 106 | bank generation and take a few minutes to run. Afterwards, you can check the templates' 107 | positions in `bank.templates`. 108 | 109 | We could generate a stochastic bank by instead passing `"stochastic"` as the second 110 | argument to `fill_bank`, which would take much longer to generate. 111 | 112 | To estimate the coverage of our bank at 1000 points sampled from the metric density, 113 | we can run 114 | 115 | ```python 116 | key, subkey = random.split(key) 117 | bank.calc_bank_effectualness(subkey, 1000) 118 | ``` 119 | 120 | This will populate the attributes `effectualness`, `effectualness_points`, `eta_est` 121 | and `eta_est_err` of `bank`. The first two are the effectualness and sampled points. 122 | The second are the resulting Monte Carlo estimate (and associated error) of the 123 | banks' coverage. 124 | 125 | We can save the bank to the current working directory with 126 | 127 | ```python 128 | bank.save() 129 | ``` 130 | 131 | and reload it with 132 | 133 | ```python 134 | Bank.load("3pn-bank.npz", amp, Psi, Sn, sampler) 135 | ``` 136 | 137 | Note the variables you must provide when loading a bank. This is because we do not 138 | use e.g. `pickle` to save function attributes. 139 | 140 | ## Functional interface 141 | 142 | This interface requires a bit more manual setup. We must first set up the metric 143 | density: 144 | 145 | ```python 146 | from diffbank.metric import get_density 147 | 148 | density_fun = lambda theta: get_density(theta, amp, Psi, fs, Sn) 149 | ``` 150 | 151 | A match function is also required to check whether effectualness points are covered. 152 | This requires defining padding arrays that make the IFFT maximization over the difference 153 | in time of coalescence for the two waveforms work correctly: 154 | 155 | ```python 156 | from diffbank.utils import get_match 157 | 158 | eff_pad_low, eff_pad_high = get_eff_pads(fs) 159 | match_fun = lambda theta1, theta2: get_match( 160 | theta1, theta2, amp, Psi, amp, Psi, fs, Sn, eff_pad_low, eff_pad_high, 161 | ) 162 | ``` 163 | 164 | Finally we need to define the density of the base sampler, which we can set to 1 165 | since it need not be normalized. Also, we need the ratio of the metric density to 166 | this one: 167 | 168 | ```python 169 | key, subkey = random.split(key) 170 | density_fun_base = lambda _: jnp.array(1.0) 171 | ratio_max = est_ratio_max(subkey, density_fun, sample_base, density_fun_base) 172 | ``` 173 | 174 | Then we can generate some templates: 175 | 176 | ```python 177 | key, subkey = random.split(key) 178 | templates, eff_pts = gen_bank_random( 179 | subkey, 180 | 1 - m_star, 181 | eta, 182 | match_fun, 183 | ratio_max, 184 | density_fun, 185 | sample_base, 186 | density_fun_base 187 | ) 188 | ``` 189 | 190 | This is the function wrapped by `Bank.fill_bank`. 191 | 192 | For stochastic bank generation, you must pass in a sampler that proposes templates. 193 | This is because the choice of this sampler does not make a huge difference in bank 194 | generation time. We can set up a rejection sampler to draw from the metric density 195 | with 196 | 197 | ```python 198 | from diffbank.utils import gen_template_rejection 199 | 200 | gen_template = lambda key: gen_template_rejection( 201 | key, ratio_max, density_fun, sampler, density_fun_base 202 | ) 203 | ``` 204 | 205 | and then generate the bank: 206 | 207 | ```python 208 | templates, eff_pts = gen_bank_stochastic( 209 | key, minimum_match, eta, match_fun, propose_template, eff_pt_sampler, n_eff 210 | ) 211 | ``` 212 | 213 | To do: explain how to check coverage. 214 | -------------------------------------------------------------------------------- /docs/source/utils.md: -------------------------------------------------------------------------------- 1 | # diffbank.utils 2 | 3 | ```{eval-rst} 4 | .. automodule:: diffbank.utils 5 | :members: 6 | :undoc-members: 7 | :show-inheritance: 8 | ``` 9 | -------------------------------------------------------------------------------- /docs/source/waveforms.md: -------------------------------------------------------------------------------- 1 | # diffbank.waveforms 2 | 3 | All waveforms are defined in the frequency domain in terms of an amplitude and a 4 | phase function. 5 | 6 | ## Waveform models 7 | 8 | ### diffbank.waveforms.kappa4D 9 | 10 | ```{eval-rst} 11 | .. automodule:: diffbank.waveforms.kappa4D 12 | :members: 13 | :undoc-members: 14 | :show-inheritance: 15 | ``` 16 | 17 | ### diffbank.waveforms.kappa5D 18 | 19 | ```{eval-rst} 20 | .. automodule:: diffbank.waveforms.kappa5D 21 | :members: 22 | :undoc-members: 23 | :show-inheritance: 24 | ``` 25 | 26 | ### diffbank.waveforms.kappa6D 27 | 28 | ```{eval-rst} 29 | .. automodule:: diffbank.waveforms.kappa6D 30 | :members: 31 | :undoc-members: 32 | :show-inheritance: 33 | ``` 34 | 35 | ### diffbank.waveforms.taylorF2 36 | 37 | ```{eval-rst} 38 | .. automodule:: diffbank.waveforms.taylorF2 39 | :members: 40 | :undoc-members: 41 | :show-inheritance: 42 | ``` 43 | 44 | ### diffbank.waveforms.taylorf2reducedspin 45 | 46 | ```{eval-rst} 47 | .. automodule:: diffbank.waveforms.taylorf2reducedspin 48 | :members: 49 | :undoc-members: 50 | :show-inheritance: 51 | ``` 52 | 53 | ### diffbank.waveforms.threePN_simple 54 | 55 | ```{eval-rst} 56 | .. automodule:: diffbank.waveforms.threePN_simple 57 | :members: 58 | :undoc-members: 59 | :show-inheritance: 60 | ``` 61 | 62 | ### diffbank.waveforms.twoPN_chirptimes 63 | 64 | ```{eval-rst} 65 | .. automodule:: diffbank.waveforms.twoPN_chirptimes 66 | :members: 67 | :undoc-members: 68 | :show-inheritance: 69 | ``` 70 | 71 | ### diffbank.waveforms.twoPN_simple 72 | 73 | ```{eval-rst} 74 | .. automodule:: diffbank.waveforms.twoPN_simple 75 | :members: 76 | :undoc-members: 77 | :show-inheritance: 78 | ``` 79 | -------------------------------------------------------------------------------- /plot_style.mplstyle: -------------------------------------------------------------------------------- 1 | font.size : 18 2 | legend.fontsize : 18 3 | legend.frameon : False 4 | axes.labelsize : 18 5 | axes.titlesize : 18 6 | xtick.labelsize: 18 7 | ytick.labelsize : 18 8 | figure.figsize : 7, 5 9 | xtick.top : True 10 | ytick.right : True 11 | xtick.bottom : True 12 | ytick.left : True 13 | xtick.major.size : 8 14 | xtick.minor.size : 4 15 | ytick.major.size : 8 16 | ytick.minor.size : 4 17 | xtick.direction : in 18 | ytick.direction : in 19 | text.usetex : True 20 | font.family : serif 21 | axes.linewidth : 1.5 22 | 23 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = [ 3 | "setuptools>=42", 4 | "wheel" 5 | ] 6 | build-backend = "setuptools.build_meta" 7 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | black 2 | click 3 | diffjeom 4 | furo 5 | isort 6 | jax 7 | jaxlib 8 | jupyter 9 | matplotlib 10 | myst-parser 11 | numpy 12 | pytest 13 | scipy 14 | Sphinx 15 | sphinx-autodoc-typehints 16 | sphinxcontrib-napoleon 17 | tqdm 18 | -------------------------------------------------------------------------------- /scripts/README.md: -------------------------------------------------------------------------------- 1 | # Scripts 2 | 3 | These are for creating figures and running analysis. They should be run from and 4 | produce outputs in this directory. 5 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-1-mm=0.8-eta_star=0.8-n_eff=10.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:2e5acea423e6221c3e667fff2915199f485f374498fa8c188bc88ec43ed82ffc 3 | size 83132 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-1-mm=0.95-eta_star=0.9-n_eff=1000.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:af5d188e3b014d005c1c93e22db9278152052596fa736ca993ae31d582ddb4a4 3 | size 279484 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-100-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:cd04c7a2eaab45a8051d66a40b3e8c7d15aa65163b72ef24639c2a07124e16df 3 | size 148493 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-101-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:440057acfe141a5342fca3c2084532f5749e2da99c101756137e9be4bd234d2f 3 | size 147309 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-102-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:788a941d10dd8fbe147e3cab5f0b447254fa9699ff0ff80b511e11c05d8c6468 3 | size 145981 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-103-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:c128714d7114b59cee14c916673dec8cb79ae14bddaf776875d528aa58016b9b 3 | size 147645 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-104-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:4d89a36e1edb5ea1fc98ab23b07233b9eb58553271dc8e78d54b6231ef6c6fe0 3 | size 146173 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-105-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:9621451c986a47c59e10df46b6363e427573d98149be65cde16a31be4103e84e 3 | size 139869 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-106-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:9e6b82766b0674f7baf7057e34f63236e72a30c78bdbbfb12bf52672653ef60f 3 | size 144173 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-107-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:057d61c3ceed8c26d982277ee32b2702c800886a26cbbe2e0cb7ae9d81a2059c 3 | size 144029 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-108-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:e77b3d62bb75255222e396ecd3568f285d2dd3b56b99d52bb758959fee72520b 3 | size 146573 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-109-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:2ca310ef30d25895f4b3d31dfe2d199aabc90eff1bc60bdae010d037e298016e 3 | size 144045 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-110-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:aa9d66e9ff431c5f7e7f2987c5832ccf9bada5fb4fcab27d8f46a371db6ea218 3 | size 147325 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-111-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:467569dd135559afdcd2953e3d8b770c231e2ea86c781961cccae6f399dc891b 3 | size 140125 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-112-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:f51d70d33228218081cdab7704740fa99ba2ab7b1045da04cc21efc9c6caff7b 3 | size 145085 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-113-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:677a49b67fd8312021fe6efc1495e92cf6ee019f4fef4784df8f83ad5f28a00f 3 | size 142781 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-114-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:05970914543943907e2c44b19d1afda4ec30b3854bcf17feee40072757382e6b 3 | size 150701 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-115-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:abf104b8a5959e63576232eff9c8bbedacdba587efa026466c04a4cd26b5f524 3 | size 147933 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-116-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:68831d2fd7a46bf659ab16bf4857ec14e96afde0fce2ebd30de6dd1beffc7c05 3 | size 145133 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-117-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:0a94cf0f201bf1ff2df07ff969fc0a612d98ab6bf7683fffaccbdf4dbfa42dac 3 | size 148093 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-118-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:d43be6d49d3f8e414759139109d9522d9e1f9e8ca90a21384dc85951721a9579 3 | size 136813 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-120-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:a87bd7dfbc69eedaf5f5aec9cf601ff0e8667dc4251cd7a31e269ffe2c8af5c3 3 | size 145053 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-121-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:4f0f5ff256fe14ccf443a24871536d959baf9bc8bca0aae7b8d692aa4c2ed7aa 3 | size 146141 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-122-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:91c5e13d01f1061922eaf54dea91af0fcd2e4151e4817e17a38707f1b2c8f190 3 | size 145741 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-123-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:74718f034b92450ed0b426c6da698ab168c6f0c914be2149e2989a830393a0dd 3 | size 139085 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-124-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:744e742f5bea4b62c55983d26c715d3eaaacabd911288c10be1d1068e275d949 3 | size 151053 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-125-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:e26fc21844b9e6f641af0acb5a999b8abe4704e37fe215c6077ab4c264044b0d 3 | size 153277 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-126-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:db5b0edaf1340e0317f96435f38046cca56d8a1ec9f38760f69434847f4caa2d 3 | size 147197 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-127-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:8d658caf865b4aeebc02c6c3bb855826f93bba94258d8ef8d312f5267eb62785 3 | size 141693 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-128-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:d788e28a6c2592ec23eacab033eb39560679be2095a78efb2c6ede97604c0129 3 | size 142589 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-129-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:c457ea9cd977d887782871b79689e8dded2164d837086523f821d7588b9d6fd2 3 | size 146109 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-130-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:f0a5cffde59a3781aa890894c1c970fd931f705c06f2781c0f474c648d05859a 3 | size 146877 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-501-mm=0.95-eta_star=0.9-n_eff=10.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:b48dadca339cfd651083f334dd3b9bb7cc4fa433d45ad6f9f0ab591422dde9e5 3 | size 120334 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-502-mm=0.95-eta_star=0.9-n_eff=11.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:21479576a2455b691bcaa901bde88b0d128601ecf0ba8c885bfa98a75f7917e2 3 | size 266252 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-503-mm=0.95-eta_star=0.9-n_eff=13.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:1dd2a2ef206b5849f96f718ba538554e602d9866c8e5aa214455a90b22d771d5 3 | size 223372 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-504-mm=0.95-eta_star=0.9-n_eff=15.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:b829e2d6fce51a3eb47807957d012f5157abbbe28ef0cb2eced88a6397017da2 3 | size 205916 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-505-mm=0.95-eta_star=0.9-n_eff=17.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:8a3bc8beec656d60defee7160323d429df3573f66721f8a61c1f0a1b5ccfe463 3 | size 204236 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-506-mm=0.95-eta_star=0.9-n_eff=20.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:8440d86d2eb2126a9b07eafc54938ed78a263705fe330fb0b56ed820f1ee1998 3 | size 223580 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-507-mm=0.95-eta_star=0.9-n_eff=23.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:a252570b8a4ad905eb2d0694fdaa66a76c603e8800c6a4dfd0916257297ac880 3 | size 219180 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-508-mm=0.95-eta_star=0.9-n_eff=26.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:0b02f2772e74052272989bfb7067ccfd4b4a5ff60b01ef9e82e4080dd717d06f 3 | size 214508 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-509-mm=0.95-eta_star=0.9-n_eff=30.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:09858f417acf7c9b4cf08108d8e3784c44882328c845eb2149045dd8c63a733f 3 | size 207692 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-510-mm=0.95-eta_star=0.9-n_eff=35.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:2f6dc8eaabdc6c8ec6c8784d758dbcda1900d3a606249017df72d2dae9839cf4 3 | size 234828 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-511-mm=0.95-eta_star=0.9-n_eff=40.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:ec9fe4c5a7ad3aaadc493e281b80387385452382dcc6753e972e27edefe3cb90 3 | size 243676 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-512-mm=0.95-eta_star=0.9-n_eff=47.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:4e5c7139608cff22d60eb6a5f8a463612c0be716509415265c30d8a5169273cc 3 | size 227740 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-513-mm=0.95-eta_star=0.9-n_eff=54.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:e31b4da05311528caeedd66097dc04661e6e005d711a9838129d3bf3a258b1df 3 | size 230268 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-514-mm=0.95-eta_star=0.9-n_eff=62.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:cbc6bbc797e1a6ac001a7ac19fbdc8611e827ce7ceac73b76da6b90f86570eb1 3 | size 211356 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-515-mm=0.95-eta_star=0.9-n_eff=71.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:61e1d2625177f3697c83c518cff46b2c4cb2353a57ff3a68fb49f08ad3fa5968 3 | size 208380 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-516-mm=0.95-eta_star=0.9-n_eff=82.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:adbdcaea7dafc82312ccdd96c9d16e85bff08ebe1db38dee2f4c75a535b5337f 3 | size 233196 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-517-mm=0.95-eta_star=0.9-n_eff=95.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:0fc45cec9ddf5683ac6358fc4a744cc47314f4941d5a4c74098972b3e51b1ad2 3 | size 215100 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-518-mm=0.95-eta_star=0.9-n_eff=109.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:c769ae8a7557e8e7807c50719789999ce9c48e5178b19e819e8a4e9d9d474bd0 3 | size 217085 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-519-mm=0.95-eta_star=0.9-n_eff=126.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:424a30ad214beb035f3305fa7d3d4e8144e6ed9b2a5ff54e13e943751b9c0f85 3 | size 219469 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-520-mm=0.95-eta_star=0.9-n_eff=145.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:74c306840a4c3cb6a6fa2ec2ececd03677676a1d37fb1da5b32cb7642bd94ed7 3 | size 221597 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-521-mm=0.95-eta_star=0.9-n_eff=167.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:2c6a3dc282004bac498f1542e6b863fa4a5f384969d3d1190eb2d204cb463c99 3 | size 214125 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-522-mm=0.95-eta_star=0.9-n_eff=193.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:6581c87f2cc73cddaea349a475e6e004a6b170fbfcf4529d0dac8d0e196fb846 3 | size 217053 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-523-mm=0.95-eta_star=0.9-n_eff=222.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:048d5cc235cb4b2800958103758ec36a3094f55ca1d8cc120360a9f28e917eb3 3 | size 221661 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-524-mm=0.95-eta_star=0.9-n_eff=255.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:1a416220f2f5e897915810e3e1204c737a62a5c3ee871d7da0c06ea99f16c5a2 3 | size 218701 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-525-mm=0.95-eta_star=0.9-n_eff=294.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:8d14350ff821526837c9b105c9b0bee376b1215935a9dd7fbfe4b6803fd00120 3 | size 225741 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-526-mm=0.95-eta_star=0.9-n_eff=339.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:374d2bacafde99fdf4e6d1644ebecb9c25a56fc50849eb71a68b1488d46ce471 3 | size 217965 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-527-mm=0.95-eta_star=0.9-n_eff=390.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:acc0d5e2de7e548edf543c8757d00e844d79a68d20a611c6926cb49271efde40 3 | size 217757 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-528-mm=0.95-eta_star=0.9-n_eff=449.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:764f05e1b6612a2b100bb1c4ef88cab1764710a1adf144b5298d4d921fce71c8 3 | size 209837 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-529-mm=0.95-eta_star=0.9-n_eff=517.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:240fd8715b8fd6299b0d426b02349d494a1afc96fb58d3d57a9b541c455f63f0 3 | size 214669 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-530-mm=0.95-eta_star=0.9-n_eff=596.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:389944c0fd1964b07c24eb01dbd70c13f4b99cb77a5ab69589031497a67af792 3 | size 219485 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-531-mm=0.95-eta_star=0.9-n_eff=686.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:90fd89d63307a9f09a9008a098f6fe0bbfa35a2970299fb698b8cdee3049d756 3 | size 225277 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-532-mm=0.95-eta_star=0.9-n_eff=790.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:bcceb633a529c926cefc9ed5b7404d3dab5fe8b09627542f5c77f3992540da03 3 | size 217789 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-533-mm=0.95-eta_star=0.9-n_eff=910.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:d62b38cad1663ab18e38c5de126055c7dd4d1d3e3dea757839ef0150ac753099 3 | size 217885 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-534-mm=0.95-eta_star=0.9-n_eff=1048.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:987ac1213e177d8f63fb5168f9c53fe2c9f06ce019447e3cba2623af06fc8d27 3 | size 220030 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-535-mm=0.95-eta_star=0.9-n_eff=1206.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:1a288d1796594e4825217a7b8dfd6bc49529b6303df3873494bcef3ccf3df465 3 | size 217470 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-536-mm=0.95-eta_star=0.9-n_eff=1389.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:10916a4732b8ff643639ae597ba6c83cf84d65d85b851858d354b8f012cdcb03 3 | size 215374 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-537-mm=0.95-eta_star=0.9-n_eff=1599.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:2535bd5e57510b2466e8f69051d4654336d9410365935c40b9ea15f82e32517f 3 | size 217470 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-538-mm=0.95-eta_star=0.9-n_eff=1842.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:978166046be54e2797047dab79602eced54113262b1db02483b85c4b41f4a178 3 | size 219262 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-539-mm=0.95-eta_star=0.9-n_eff=2120.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:cf525d4d5844428f271aa4f5c6d8a8e56934227c06cde57e2a9574739bf23baa 3 | size 219310 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-540-mm=0.95-eta_star=0.9-n_eff=2442.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:d956ea8b036468ffc8273f5e20a898d9c7c64e0bd1aeb9e99e5b2226d0821488 3 | size 218302 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-541-mm=0.95-eta_star=0.9-n_eff=2811.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:2969069853b013debda2aa0ee1d22d5312dfa5e402e540e0f802e0f8e25fc008 3 | size 219774 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-542-mm=0.95-eta_star=0.9-n_eff=3237.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:d1228ff17f82246e835fe1418dbbf63b66218367c8e8be5fce9cbf9443d3f3f7 3 | size 217422 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-543-mm=0.95-eta_star=0.9-n_eff=3727.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:99df30fe00034a5762cb156595b4877f889dddb3164937c078af23e795735110 3 | size 220126 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-544-mm=0.95-eta_star=0.9-n_eff=4291.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:e39c3c4c5f1df06c2b87c518a3b9d85b19a8a582b20f7e4610d3f00325fabac0 3 | size 218734 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-545-mm=0.95-eta_star=0.9-n_eff=4941.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:05cedb14372c45e07030e2c76667f1f5ee05f852dd6da5f7618cbaf6a5f2c022 3 | size 217662 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-546-mm=0.95-eta_star=0.9-n_eff=5689.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:1b6cd7ffd45da801a1dab8c791cf5aeb08440de148afa28905fcbadb8fb6ca2a 3 | size 219406 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-547-mm=0.95-eta_star=0.9-n_eff=6551.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:1609ab1831e0c90db341b5f62361cd0d3f8d89b829a10cf443bf0bce4d567156 3 | size 217214 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-548-mm=0.95-eta_star=0.9-n_eff=7543.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:bd456c786ef982e641533f68a0022a2f8f04c92b7fcd6ba3a844c8e7953e584a 3 | size 219598 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-549-mm=0.95-eta_star=0.9-n_eff=8685.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:fb921f39c27757bf448cd1f4b561cdc3ab018ae8a47695eef7b62f792e4dfe6b 3 | size 217134 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-random-550-mm=0.95-eta_star=0.9-n_eff=10000.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:aea32ad52d5fa1ba1b8a475c7de459c5f9479ed28f781c48e7a0209a01d21808 3 | size 219743 4 | -------------------------------------------------------------------------------- /scripts/banks/3pn-stochastic-1000-mm=0.95-eta_star=0.9-n_eff=500.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:863d4a79f5e77c027abca6a4a6a2a9f14e8c35192ed06f97b6e9bbb6e71f6962 3 | size 129154 4 | -------------------------------------------------------------------------------- /scripts/banks/README.md: -------------------------------------------------------------------------------- 1 | # Banks 2 | 3 | Various template banks. The file names are of the form `WAVEFORM-KIND-SEED.npz`. 4 | -------------------------------------------------------------------------------- /scripts/banks/tf2-random-1-mm=0.96-eta_star=0.9-n_eff=1000.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:e97c1aa67c2a89a885068de243af68a378fe9a7cbaf7acb2e523a950c60f8be7 3 | size 9113952 4 | -------------------------------------------------------------------------------- /scripts/banks/tf2rs-random-1-mm=0.95-eta_star=0.993-n_eff=1300.npz: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:b4077bec76ee7e1c6737b8fe6090b3df395cc6fac2b9e41a8ea155c3fbc82974 3 | size 50220645 4 | -------------------------------------------------------------------------------- /scripts/calc_bank_eff.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import click 4 | import jax 5 | import jax.numpy as jnp 6 | from genbank_3D_taylorf2reducedspin import sampler 7 | from jax import random 8 | 9 | from diffbank.bank import Bank 10 | from diffbank.noise import Sn_aLIGOZeroDetHighPower as Sn 11 | from diffbank.waveforms.taylorf2reducedspin import Psi, amp 12 | 13 | """ 14 | Update eta estimate for an existing bank using more points and resave. 15 | """ 16 | 17 | 18 | @click.command() 19 | @click.option("--path", type=str, help="path to bank") 20 | @click.option("--seed", type=int) 21 | @click.option( 22 | "--n-eta", type=int, help="number of new points at which to compute effectualnesses" 23 | ) 24 | @click.option("--device", default="cpu", help="device to run on") 25 | def run(path, seed, n_eta, device): 26 | jax.config.update("jax_platform_name", device) 27 | 28 | kind = os.path.split(path)[-1].split("-")[0] 29 | if kind != "tf2rs": 30 | raise ValueError("wrong waveform model") 31 | 32 | bank = Bank.load(path, amp, Psi, Sn, sampler) 33 | key = random.PRNGKey(seed) 34 | bank.calc_bank_effectualness(key, n_eta) 35 | jnp.savez( 36 | f"banks/effs/{os.path.splitext(os.path.split(path)[-1])[0]}-effs-{seed}.npz", 37 | eff_pts=bank.effectualness_points, 38 | effs=bank.effectualnesses, 39 | ) 40 | 41 | 42 | if __name__ == "__main__": 43 | run() 44 | -------------------------------------------------------------------------------- /scripts/figures/.gitignore: -------------------------------------------------------------------------------- 1 | # *.pdf 2 | # *.png 3 | -------------------------------------------------------------------------------- /scripts/figures/README.md: -------------------------------------------------------------------------------- 1 | # figures 2 | 3 | Figures produced by scripts. 4 | -------------------------------------------------------------------------------- /scripts/figures/bank-effs.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/figures/bank-effs.pdf -------------------------------------------------------------------------------- /scripts/figures/density.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/figures/density.pdf -------------------------------------------------------------------------------- /scripts/figures/eta-dist.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/figures/eta-dist.pdf -------------------------------------------------------------------------------- /scripts/figures/neff-scaling.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/figures/neff-scaling.pdf -------------------------------------------------------------------------------- /scripts/figures/scalar-curvature.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/figures/scalar-curvature.pdf -------------------------------------------------------------------------------- /scripts/figures/scaling.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/figures/scaling.pdf -------------------------------------------------------------------------------- /scripts/figures/taylorf2reducedspin_density.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/figures/taylorf2reducedspin_density.pdf -------------------------------------------------------------------------------- /scripts/figures/threePN-eff-cdfs.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/figures/threePN-eff-cdfs.pdf -------------------------------------------------------------------------------- /scripts/genbank_2D_threePN.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import click 4 | import jax 5 | import jax.numpy as jnp 6 | from jax import random 7 | from scipy.optimize import minimize_scalar 8 | 9 | from diffbank.bank import Bank 10 | from diffbank.utils import get_m1_m2_sampler 11 | from diffbank.waveforms.threePN_simple import Psi, amp 12 | 13 | """ 14 | Generates a 3.5PN-2D bank. 15 | 16 | To generate two banks for the paper, use 17 | 18 | >>> python genbank_2D_threePN.py --seed 1 --kind random 19 | >>> python genbank_2D_threePN.py --seed 2 --kind stochastic 20 | 21 | """ 22 | 23 | ##### Frequency settings 24 | f_u = 512.0 # Hz 25 | f_l = 24.0 # Hz 26 | N_fbins = 4880 27 | ##### 28 | 29 | 30 | @click.command() 31 | @click.option("--seed", default=1, help="PRNG seed") 32 | @click.option("--kind", default="random", help="kind of bank: 'random' or 'stochastic'") 33 | @click.option( 34 | "--n-eta", 35 | default=1000, 36 | type=int, 37 | help="number of new points at which to compute effectualnesses", 38 | ) 39 | @click.option("--mm", default=0.95, help="minimum match") 40 | @click.option("--eta-star", default=0.9, help="eta*") 41 | @click.option("--n-eff", default=1000) 42 | @click.option("--savedir", default="banks", help="directory in which to save the bank") 43 | @click.option("--device", default="cpu", help="device to run on") 44 | @click.option( 45 | "--noise", 46 | default="interpolated", 47 | help="noise curve: 'analytic' (LIGO-I) or 'interpolated' (Livingston O3a)", 48 | ) 49 | def gen_2D_threePN(seed, kind, n_eta, mm, eta_star, n_eff, savedir, device, noise): 50 | jax.config.update("jax_platform_name", device) 51 | 52 | key = random.PRNGKey(seed) 53 | m_star = 1 - mm 54 | 55 | fs = jnp.linspace(f_l, f_u, N_fbins) 56 | m_range = (1.0, 3.0) 57 | sampler = get_m1_m2_sampler(m_range, m_range) 58 | if noise == "interpolated": 59 | from diffbank.noise import Sn_O3a as Sn 60 | elif noise == "analytic": 61 | from diffbank.noise import Sn_LIGOI as Sn 62 | else: 63 | raise ValueError("invalid 'noise' argument") 64 | 65 | bank = Bank( 66 | amp, 67 | Psi, 68 | fs, 69 | Sn, 70 | m_star, 71 | eta_star, 72 | sampler, 73 | name=f"3pn-{kind}-{seed}-mm={mm}-eta_star={eta_star}-n_eff={n_eff}", 74 | ) 75 | 76 | # Get max density 77 | fun = lambda m1: -bank.density_fun(jnp.array([m1, m_range[0]])) 78 | res = minimize_scalar(fun, bracket=m_range, bounds=m_range, method="bounded") 79 | assert res.success 80 | theta_dmax = jnp.array([res.x, m_range[0]]) 81 | ratio_max = bank.density_fun(theta_dmax) 82 | bank.ratio_max = ratio_max 83 | 84 | # Fill bank 85 | key, subkey = random.split(key) 86 | bank.fill_bank(subkey, kind, n_eff) 87 | bank.save(savedir) 88 | print(f"Saved bank to {os.path.join(savedir, bank.name + '.npz')}") 89 | 90 | # Get effectualnesses 91 | if n_eta > 0: 92 | key, subkey = random.split(key) 93 | bank.calc_bank_effectualness(subkey, n_eta) 94 | bank.save(savedir) 95 | else: 96 | print("Skipping effectualnesses calculation") 97 | 98 | 99 | if __name__ == "__main__": 100 | gen_2D_threePN() 101 | -------------------------------------------------------------------------------- /scripts/genbank_3D_taylorf2reducedspin.py: -------------------------------------------------------------------------------- 1 | import os 2 | from math import pi 3 | from typing import Tuple 4 | 5 | import click 6 | import jax 7 | import jax.numpy as jnp 8 | from jax import random 9 | from scipy.optimize import minimize_scalar 10 | 11 | from diffbank.bank import Bank 12 | from diffbank.constants import MSUN, C, G 13 | from diffbank.utils import Array, PRNGKeyArray 14 | from diffbank.waveforms.taylorf2reducedspin import Psi, amp, get_th_boundary_interps 15 | 16 | """ 17 | Generate a TaylorF2ReducedSpin bank for comparison with Ajith et al 2014, 18 | https://arxiv.org/abs/1210.6666. 19 | 20 | To reproduce the bank in the paper, run 21 | 22 | >>> python genbank_3D_taylorf2reducedspin.py 23 | 24 | """ 25 | 26 | ##### Frequency settings 27 | # Since the lowest BH mass for this bank is 1 * MSUN, need to go up to its ISCO 28 | # frequency 29 | f_u = 2200.0 # Hz 30 | f_0 = f_l = 20.0 # Hz 31 | df = 0.1 32 | N_fbins = int((f_u - f_l) / df) 33 | ##### 34 | 35 | m_range = (1 * MSUN, 20 * MSUN) 36 | m_ns_thresh = 2 * MSUN 37 | M_tot_max = m_range[0] + m_range[1] 38 | chi_bh_max = 0.98 39 | chi_ns_max = 0.4 40 | 41 | th0_range, th3_interp_low, th3_interp_high = get_th_boundary_interps(*m_range, f_0) 42 | # Figure out where th3 attains its maximum 43 | def get_th3S_max(th0, th3): 44 | """ 45 | Gets max value of th3S at a given `(th0, th3)` point. This computes the 46 | component masses, gets the corresponding `chi1`, `chi2` values, computes 47 | the max value `chi` can take and converts this to a max value for `th3S`. 48 | """ 49 | M_chirp = 1 / (16 * pi * f_0) * (125 / (2 * th0 ** 3)) ** (1 / 5) * C ** 3 / G 50 | eta = (16 * pi ** 5 / 25 * th0 ** 2 / th3 ** 5) ** (1 / 3) 51 | q = (1 + jnp.sqrt(1 - 4 * eta) - 2 * eta) / (2 * eta) 52 | m2 = (1 + q) ** (1 / 5) / q ** (3 / 5) * M_chirp 53 | m1 = q * m2 54 | delta = (m1 - m2) / (m1 + m2) 55 | chi1_max = jnp.where(m1 > m_ns_thresh, chi_bh_max, chi_ns_max) 56 | chi2_max = jnp.where(m2 > m_ns_thresh, chi_bh_max, chi_ns_max) 57 | chi_s_max = (chi1_max + chi2_max) / 2 58 | chi_a_max = (chi1_max - chi2_max) / 2 59 | chi_max = chi_s_max * (1 - 76 * eta / 113) + delta * chi_a_max 60 | th3S_max = 113 * th3 * chi_max / (48 * pi) 61 | return th3S_max 62 | 63 | 64 | def get_M_tot(th0, th3): 65 | M_chirp = 1 / (16 * pi * f_0) * (125 / (2 * th0 ** 3)) ** (1 / 5) * C ** 3 / G 66 | eta = (16 * pi ** 5 / 25 * th0 ** 2 / th3 ** 5) ** (1 / 3) 67 | q = (1 + jnp.sqrt(1 - 4 * eta) - 2 * eta) / (2 * eta) 68 | m2 = (1 + q) ** (1 / 5) / q ** (3 / 5) * M_chirp 69 | m1 = q * m2 70 | return m1 + m2 71 | 72 | 73 | def is_in_bounds(theta: Array) -> Array: 74 | """ 75 | Checks if a point is in bounds using the `th` values and total mass. 76 | """ 77 | th0, th3, th3S = theta[..., 0], theta[..., 1], theta[..., 2] 78 | return jnp.logical_and( 79 | th3 > th3_interp_low(th0), 80 | jnp.logical_and( 81 | th3 < th3_interp_high(th0), 82 | jnp.logical_and( 83 | jnp.abs(th3S) < get_th3S_max(th0, th3), get_M_tot(th0, th3) < M_tot_max 84 | ), 85 | ), 86 | ) 87 | 88 | 89 | def base_sample_1( 90 | key: PRNGKeyArray, 91 | th0_range: Tuple[float, float], 92 | th3_range: Tuple[float, float], 93 | th3S_max: float, 94 | ) -> Array: 95 | """ 96 | Sample uniformly over maximum parameter ranges. 97 | """ 98 | return random.uniform( 99 | key, 100 | (3,), 101 | minval=jnp.array([th0_range[0], th3_range[0], -th3S_max]), 102 | maxval=jnp.array([th0_range[1], th3_range[1], th3S_max]), 103 | ) 104 | 105 | 106 | @jax.jit 107 | def sample_1( 108 | key: PRNGKeyArray, 109 | th0_range: Tuple[float, float], 110 | th3_range: Tuple[float, float], 111 | th3S_max: float, 112 | ) -> Array: 113 | """ 114 | Samples a single point with rejection sampling. 115 | """ 116 | cond_fun = lambda val: jnp.logical_not(is_in_bounds(val[1])) 117 | 118 | def body_fun(val): 119 | key = val[0] 120 | key, subkey = random.split(key) 121 | return (key, base_sample_1(subkey, th0_range, th3_range, th3S_max)) 122 | 123 | key, subkey = random.split(key) 124 | init_val = (key, base_sample_1(subkey, th0_range, th3_range, th3S_max)) 125 | return jax.lax.while_loop(cond_fun, body_fun, init_val)[1] 126 | 127 | 128 | def _sampler( 129 | key: PRNGKeyArray, 130 | n: int, 131 | th0_range: Tuple[float, float], 132 | th3_range: Tuple[float, float], 133 | th3S_max: float, 134 | ) -> Array: 135 | return jax.lax.map( 136 | lambda key: sample_1(key, th0_range, th3_range, th3S_max), random.split(key, n) 137 | ) 138 | 139 | 140 | # Define sampling bounds 141 | bracket = (th0_range[0], 5e3) # NOTE: need to change if m_range changes! 142 | res = minimize_scalar(lambda th0: -th3_interp_high(th0), bracket, bracket) 143 | assert res.success 144 | th0_th3_max = res.x 145 | th3_max = -res.fun 146 | th3_range = (th3_interp_low(th0_range[0]), th3_max) 147 | # Maximum value of th3 148 | th3S_max = get_th3S_max(th0_th3_max, th3_max) 149 | 150 | 151 | # Capture globals 152 | def sampler(key: PRNGKeyArray, n: int) -> Array: 153 | return _sampler(key, n, th0_range, th3_range, th3S_max) 154 | 155 | 156 | @click.command() 157 | @click.option("--seed", default=1, help="PRNG seed") 158 | @click.option("--kind", default="random", help="kind of bank: 'random' or 'stochastic'") 159 | @click.option( 160 | "--n-eta", 161 | default=0, 162 | type=int, 163 | help="number of new points at which to compute effectualnesses", 164 | ) 165 | @click.option( 166 | "--mm", default=0.95, help="minimum match, chosen to match arXiv:1210.6666" 167 | ) 168 | @click.option("--eta-star", default=0.993, help="eta, chosen to match arXiv:1210.6666") 169 | @click.option("--n-eff", default=1300) 170 | @click.option("--savedir", default="banks", help="directory in which to save the bank") 171 | @click.option("--device", default="cpu", help="device to run on") 172 | @click.option( 173 | "--noise", 174 | default="interpolated", 175 | help="noise curve: 'analytic' (LIGO-I) or 'interpolated' (aLIGOZeroDetHighPower from pycbc)", 176 | ) 177 | def gen_3D_tf2rs(seed, kind, n_eta, mm, eta_star, n_eff, savedir, device, noise): 178 | jax.config.update("jax_platform_name", device) 179 | 180 | key = random.PRNGKey(seed) 181 | m_star = 1 - mm 182 | fs = jnp.linspace(f_l, f_u, N_fbins) 183 | if noise == "interpolated": 184 | from diffbank.noise import Sn_aLIGOZeroDetHighPower as Sn 185 | elif noise == "analytic": 186 | from diffbank.noise import Sn_LIGOI as Sn 187 | else: 188 | raise ValueError("invalid 'noise' argument") 189 | 190 | bank = Bank( 191 | amp, 192 | Psi, 193 | fs, 194 | Sn, 195 | m_star, 196 | eta_star, 197 | sampler, 198 | name=f"tf2rs-{kind}-{seed}-mm={mm}-eta_star={eta_star}-n_eff={n_eff}", 199 | ) 200 | 201 | # Get max density 202 | # NOTE: need to change if m_range changes! 203 | th0s = jnp.linspace(1.0001 * th0_range[0], 0.9999 * th0_range[1], 500) 204 | th3s = th3_interp_high(th0s) * 0.99999 205 | th3Ss = -get_th3S_max(th0s, th3s) 206 | boundary_densities = jax.lax.map( 207 | bank.density_fun, jnp.stack([th0s, th3s, th3Ss], -1) 208 | ) 209 | bank.ratio_max = jnp.nanmax(boundary_densities) 210 | 211 | # Fill bank 212 | key, subkey = random.split(key) 213 | bank.fill_bank(subkey, kind, n_eff) 214 | bank.save(savedir) 215 | print(f"Saved bank to {os.path.join(savedir, bank.name + '.npz')}") 216 | 217 | # Get effectualnesses 218 | if n_eta > 0: 219 | key, subkey = random.split(key) 220 | bank.calc_bank_effectualness(subkey, n_eta) 221 | bank.save(savedir) 222 | else: 223 | print("Skipping effectualnesses calculation") 224 | 225 | 226 | if __name__ == "__main__": 227 | gen_3D_tf2rs() 228 | -------------------------------------------------------------------------------- /scripts/genbank_4D_taylorf2.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import click 4 | import jax 5 | import jax.numpy as jnp 6 | from jax import random 7 | 8 | from diffbank.bank import Bank 9 | from diffbank.waveforms import taylorF2 10 | 11 | """ 12 | Generate a TaylorF2 bank which can be compared with the BNS section of https://arxiv.org/abs/1904.01683 13 | 14 | To reproduce the bank in the paper, run 15 | 16 | >>> python genbank_4D_taylorf2.py 17 | 18 | """ 19 | 20 | ##### Frequency settings 21 | f_u = 512.0 # Hz 22 | f_l = 24.0 # Hz 23 | N_fbins = 4880 24 | ##### 25 | 26 | m1_range = (1.0001, 3.0) 27 | m2_range = (1.0, 3.0) 28 | 29 | chi1_range = (-0.99, 0.99) 30 | chi2_range = (-0.99, 0.99) 31 | 32 | 33 | def sampler(key, n): 34 | ms = random.uniform( 35 | key, 36 | minval=jnp.array([m1_range[0], m2_range[0]]), 37 | maxval=jnp.array([m1_range[1], m2_range[1]]), 38 | shape=(n, 2), 39 | ) 40 | key, subkey = random.split(key) 41 | chi1s = random.uniform( 42 | subkey, 43 | minval=jnp.array(chi1_range[0]), 44 | maxval=jnp.array(chi1_range[1]), 45 | shape=(n, 1), 46 | ) 47 | key, subkey = random.split(key) 48 | chi2s = random.uniform( 49 | subkey, 50 | minval=jnp.array(chi2_range[0]), 51 | maxval=jnp.array(chi2_range[1]), 52 | shape=(n, 1), 53 | ) 54 | ms_correct = jnp.stack( 55 | [ 56 | ms.max(axis=1), 57 | ms.min(axis=1), 58 | ] 59 | ).T 60 | return jnp.hstack((ms_correct, chi1s, chi2s)) 61 | 62 | 63 | @click.command() 64 | @click.option("--seed", default=1, help="PRNG seed") 65 | @click.option("--kind", default="random", help="kind of bank: 'random' or 'stochastic'") 66 | @click.option( 67 | "--n-eta", 68 | default=1000, 69 | type=int, 70 | help="number of new points at which to compute effectualnesses", 71 | ) 72 | @click.option("--mm", default=0.96, help="minimum match") 73 | @click.option("--eta-star", default=0.9, help="eta*") 74 | @click.option("--n-eff", default=1000) 75 | @click.option("--savedir", default="banks", help="directory in which to save the bank") 76 | @click.option("--device", default="cpu", help="device to run on") 77 | @click.option( 78 | "--noise", 79 | default="interpolated", 80 | help="noise curve: 'analytic' (LIGO-I) or 'interpolated' (LIGO O2)", 81 | ) 82 | def gen_4D_taylorf2bank(seed, kind, n_eta, mm, eta_star, n_eff, savedir, device, noise): 83 | jax.config.update("jax_platform_name", device) 84 | 85 | key = random.PRNGKey(seed) 86 | fs = jnp.linspace(f_l, f_u, N_fbins) 87 | if noise == "interpolated": 88 | from diffbank.noise import Sn_O2 as Sn 89 | elif noise == "analytic": 90 | from diffbank.noise import Sn_LIGOI as Sn 91 | else: 92 | raise ValueError("invalid 'noise' argument") 93 | 94 | bank = Bank( 95 | taylorF2.Amp, 96 | taylorF2.Psi, 97 | fs, 98 | Sn, 99 | 1 - mm, 100 | eta_star, 101 | sampler, 102 | name=f"tf2-{kind}-{seed}-mm={mm}-eta_star={eta_star}-n_eff={n_eff}", 103 | ) 104 | 105 | theta_max = jnp.array([m1_range[0], m2_range[0], chi1_range[0], chi2_range[1]]) 106 | bank.ratio_max = bank.density_fun(theta_max) 107 | 108 | key, subkey = random.split(key) 109 | bank.fill_bank(subkey, kind, n_eff) 110 | bank.save(savedir) 111 | print(f"Saved bank to {os.path.join(savedir, bank.name + '.npz')}") 112 | 113 | # Get effectualnesses 114 | if n_eta > 0: 115 | key, subkey = random.split(key) 116 | bank.calc_bank_effectualness(subkey, n_eta) 117 | bank.save(savedir) 118 | else: 119 | print("Skipping effectualnesses calculation") 120 | 121 | 122 | if __name__ == "__main__": 123 | gen_4D_taylorf2bank() 124 | -------------------------------------------------------------------------------- /scripts/job-2D-threePN.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | #SBATCH --job-name=threePN 3 | #SBATCH --account=rrg-lplevass 4 | #SBATCH --time=00-48:00:00 5 | #SBATCH --gres=gpu:1 6 | #SBATCH --cpus-per-task=1 7 | #SBATCH --mem=47000M 8 | 9 | source ~/.virtualenvs/diffbank-3.9.6/bin/activate 10 | 11 | python genbank_2D_threePN.py --seed 1 --kind random --device gpu 12 | python genbank_2D_threePN.py --seed 2 --kind stochastic --device gpu 13 | -------------------------------------------------------------------------------- /scripts/job-3D-taylorfsreducedspin.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | #SBATCH --job-name=tf2rs 3 | #SBATCH --account=rrg-lplevass 4 | #SBATCH --time=07-00:00:00 5 | #SBATCH --gres=gpu:1 6 | #SBATCH --cpus-per-task=1 7 | #SBATCH --mem=47000M 8 | 9 | source ~/.virtualenvs/diffbank-3.9.6/bin/activate 10 | 11 | python genbank_3D_taylorf2reducedspin.py --seed 1 --kind random --device gpu --n-eta 0 12 | -------------------------------------------------------------------------------- /scripts/job-4D-taylorf2.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | #SBATCH --job-name=tf2 3 | #SBATCH --account=rrg-lplevass 4 | #SBATCH --time=07-00:00:00 5 | #SBATCH --gres=gpu:1 6 | #SBATCH --cpus-per-task=1 7 | #SBATCH --mem=47000M 8 | 9 | source ~/.virtualenvs/diffbank-3.9.6/bin/activate 10 | 11 | python genbank_4D_taylorf2.py --seed 1 --kind random --device gpu 12 | -------------------------------------------------------------------------------- /scripts/job-threePN-coverage.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | #SBATCH --job-name=3pn-neff 3 | #SBATCH --account=rrg-lplevass 4 | #SBATCH --time=00-48:00:00 5 | #SBATCH --gres=gpu:a100:1 6 | #SBATCH --mem=200G 7 | #SBATCH --array=1-50 8 | 9 | source ~/.envs/diffbank-3.9.6/bin/activate 10 | 11 | neffs=(10 11 13 15 17 20 23 26 30 35 40 47 54 62 71 82 95 109 126 145 167 193 222 255 294 339 390 449 517 596 686 790 910 1048 1206 1389 1599 1842 2120 2442 2811 3237 3727 4291 4941 5689 6551 7543 8685 10000) 12 | seeds=(500 501 502 503 504 505 506 507 508 509 510 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550) 13 | 14 | # Get neff and seed for this job 15 | python genbank_2D_threePN.py \ 16 | --n-eff ${neffs[$SLURM_ARRAY_TASK_ID]} \ 17 | --seed ${seeds[$SLURM_ARRAY_TASK_ID]} \ 18 | --kind random \ 19 | --device gpu \ 20 | --n-eta 4000 \ 21 | --noise analytic 22 | 23 | wait 24 | -------------------------------------------------------------------------------- /scripts/job-threePN-est-p-variation.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | #SBATCH --job-name=p 3 | #SBATCH --account=rrg-lplevass 4 | #SBATCH --time=00-04:00:00 5 | #SBATCH --gres=gpu:1 6 | #SBATCH --cpus-per-task=1 7 | #SBATCH --mem=47000M 8 | 9 | source ~/.virtualenvs/diffbank-3.9.6/bin/activate 10 | 11 | python threePN_est_p_variation.py --mm 0.95 & 12 | python threePN_est_p_variation.py --mm 0.8 & 13 | 14 | wait 15 | -------------------------------------------------------------------------------- /scripts/job-threePN-scaling.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | #SBATCH --job-name=scaling 3 | #SBATCH --time=72:00:00 4 | #SBATCH -n 13 5 | #SBATCH -p normal 6 | 7 | savedir="$TMPDIR/threePN-banks-scaling" 8 | outdir="$TMPDIR/threePN-outputs-scaling" 9 | mkdir -p $savedir # scratch for saving results 10 | mkdir -p $outdir # scratch for saving results 11 | 12 | echo "Random, varying mm" 13 | python genbank_2D_threePN.py --seed 5 --kind random --mm 0.95 --eta-star 0.9 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-random-5-mm=0.95-eta_star=0.9.txt & 14 | python genbank_2D_threePN.py --seed 6 --kind random --mm 0.90 --eta-star 0.9 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-random-6-mm=0.90-eta_star=0.9.txt & 15 | python genbank_2D_threePN.py --seed 7 --kind random --mm 0.85 --eta-star 0.9 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-random-7-mm=0.85-eta_star=0.9.txt & 16 | python genbank_2D_threePN.py --seed 8 --kind random --mm 0.80 --eta-star 0.9 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-random-8-mm=0.80-eta_star=0.9.txt & 17 | python genbank_2D_threePN.py --seed 9 --kind random --mm 0.75 --eta-star 0.9 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-random-9-mm=0.75-eta_star=0.9.txt & 18 | echo "Stochastic, varying mm" 19 | python genbank_2D_threePN.py --seed 10 --kind stochastic --mm 0.95 --eta-star 0.9 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-stochastic-10-mm=0.95-eta_star=0.9.txt & 20 | python genbank_2D_threePN.py --seed 11 --kind stochastic --mm 0.90 --eta-star 0.9 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-stochastic-11-mm=0.90-eta_star=0.9.txt & 21 | python genbank_2D_threePN.py --seed 12 --kind stochastic --mm 0.85 --eta-star 0.9 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-stochastic-12-mm=0.85-eta_star=0.9.txt & 22 | python genbank_2D_threePN.py --seed 13 --kind stochastic --mm 0.80 --eta-star 0.9 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-stochastic-13-mm=0.80-eta_star=0.9.txt & 23 | python genbank_2D_threePN.py --seed 14 --kind stochastic --mm 0.75 --eta-star 0.9 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-stochastic-14-mm=0.75-eta_star=0.9.txt & 24 | wait 25 | echo "Done varying mm" 26 | 27 | # Retrieve final results 28 | cp -r $savedir . 29 | cp -r $outdir . 30 | 31 | echo "Random, varying eta_star" 32 | python genbank_2D_threePN.py --seed 15 --kind random --mm 0.90 --eta-star 0.975 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-random-15-mm=0.90-eta_star=0.975.txt & 33 | python genbank_2D_threePN.py --seed 16 --kind random --mm 0.90 --eta-star 0.950 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-random-16-mm=0.90-eta_star=0.950.txt & 34 | python genbank_2D_threePN.py --seed 17 --kind random --mm 0.90 --eta-star 0.925 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-random-17-mm=0.90-eta_star=0.925.txt & 35 | python genbank_2D_threePN.py --seed 18 --kind random --mm 0.90 --eta-star 0.900 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-random-18-mm=0.90-eta_star=0.900.txt & 36 | python genbank_2D_threePN.py --seed 19 --kind random --mm 0.90 --eta-star 0.875 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-random-19-mm=0.90-eta_star=0.875.txt & 37 | python genbank_2D_threePN.py --seed 20 --kind random --mm 0.90 --eta-star 0.850 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-random-20-mm=0.90-eta_star=0.850.txt & 38 | echo "Stochastic, varying eta_star" 39 | python genbank_2D_threePN.py --seed 21 --kind stochastic --mm 0.90 --eta-star 0.975 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-stochastic-21-mm=0.90-eta_star=0.975.txt & 40 | python genbank_2D_threePN.py --seed 22 --kind stochastic --mm 0.90 --eta-star 0.950 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-stochastic-22-mm=0.90-eta_star=0.950.txt & 41 | python genbank_2D_threePN.py --seed 23 --kind stochastic --mm 0.90 --eta-star 0.925 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-stochastic-23-mm=0.90-eta_star=0.925.txt & 42 | python genbank_2D_threePN.py --seed 24 --kind stochastic --mm 0.90 --eta-star 0.900 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-stochastic-24-mm=0.90-eta_star=0.900.txt & 43 | python genbank_2D_threePN.py --seed 25 --kind stochastic --mm 0.90 --eta-star 0.875 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-stochastic-25-mm=0.90-eta_star=0.875.txt & 44 | python genbank_2D_threePN.py --seed 26 --kind stochastic --mm 0.90 --eta-star 0.850 --n-eff 1000 --noise analytic --n-eta 0 --savedir $savedir 2>&1 | tee -a $outdir/threePN-stochastic-26-mm=0.90-eta_star=0.850.txt & 45 | wait 46 | echo "Done with varying eta_star" 47 | 48 | # Retrieve final results 49 | cp -r $savedir . 50 | cp -r $outdir . 51 | -------------------------------------------------------------------------------- /scripts/minimal_example.py: -------------------------------------------------------------------------------- 1 | from jax import random 2 | import jax.numpy as jnp 3 | 4 | from diffbank.bank import Bank 5 | from diffbank.noise import Sn_LIGOI as Sn 6 | from diffbank.utils import get_m1_m2_sampler 7 | from diffbank.waveforms.threePN_simple import Psi, amp 8 | 9 | """ 10 | Minimal example of bank generation. Creates a very coarse 3.5PN bank. 11 | """ 12 | 13 | 14 | # Set seed 15 | key = random.PRNGKey(10) 16 | 17 | # Set up generation 18 | m_range = (2.5, 3.0) 19 | sampler = get_m1_m2_sampler(m_range, m_range) 20 | fs = jnp.linspace(24.0, 512.0, 4880) 21 | minimum_match = 0.8 22 | m_star = 1 - minimum_match 23 | eta = 0.8 24 | n_eff = 100 25 | bank = Bank(amp, Psi, fs, Sn, m_star, eta, sampler, name=f"3pn-bank") 26 | 27 | # Estimate max ratio between metric density and base density (required for rejection 28 | # sampling from metric density) 29 | key, subkey = random.split(key) 30 | bank.ratio_max = bank.est_ratio_max(subkey)[0] 31 | 32 | # Generate the bank 33 | key, subkey = random.split(key) 34 | bank.fill_bank(subkey, "random", n_eff) 35 | 36 | # Estimate its effectualness 37 | key, subkey = random.split(key) 38 | bank.calc_bank_effectualness(subkey, 100) 39 | 40 | # Save 41 | bank.save() 42 | 43 | # Example of loading bank 44 | Bank.load("3pn-bank.npz", amp, Psi, Sn, sampler) 45 | -------------------------------------------------------------------------------- /scripts/plot_bank_effs.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import click 4 | import jax.numpy as jnp 5 | import matplotlib.pyplot as plt 6 | 7 | plt.style.use("../plot_style.mplstyle") 8 | 9 | """ 10 | Plots effectualness histograms for the main three banks in the diffbank paper. 11 | 12 | The banks are generated by running the following scripts: 13 | 14 | >>> python genbank_2D_threePN.py --seed 1 --kind random 15 | >>> python genbank_3D_taylorf2reducedspin.py --seed 1 --kind random 16 | >>> python genbank_4D_taylorf2.py --seed 1 --kind random 17 | """ 18 | 19 | 20 | @click.command() 21 | def run(): 22 | bank_dir = "banks" 23 | 24 | effs_3pn = jnp.load( 25 | os.path.join(bank_dir, "3pn-random-1-mm=0.95-eta_star=0.9-n_eff=1000.npz"), 26 | allow_pickle=True, 27 | )["bank"].item()["effectualnesses"] 28 | 29 | effs_tf2rs = jnp.load( 30 | os.path.join(bank_dir, "tf2rs-random-1-mm=0.95-eta_star=0.993-n_eff=1300.npz"), 31 | allow_pickle=True, 32 | )["bank"].item()["effectualnesses"] 33 | 34 | effs_tf2 = jnp.load( 35 | os.path.join(bank_dir, "tf2-random-1-mm=0.96-eta_star=0.9-n_eff=1000.npz"), 36 | allow_pickle=True, 37 | )["bank"].item()["effectualnesses"] 38 | 39 | hist_kwargs = dict( 40 | cumulative=True, density=True, histtype="step", bins=jnp.linspace(0, 1.0, 101) 41 | ) 42 | plt.hist(effs_3pn, label="3.5PN-2D", **hist_kwargs) 43 | plt.hist(effs_tf2rs, label="2.5PN-3D", **hist_kwargs) 44 | plt.hist(effs_tf2, label="3.5PN-4D", **hist_kwargs) 45 | 46 | # Minimum mismatches 47 | for i, mm in enumerate([0.95, 0.95, 0.96]): 48 | plt.axvline(mm, color=f"C{i}", linestyle="--", linewidth=1) 49 | 50 | # Target coverage 51 | # plt.axhline(1 - 0.9, color="C0", linestyle="--", linewidth=1) 52 | # plt.axhline(1 - 0.993, color="C1", linestyle="--", linewidth=1) 53 | # plt.axhline(1 - 0.9, color="C2", linestyle="--", linewidth=1) 54 | for i, eta in enumerate([0.9, 0.993, 0.9]): 55 | eta_err = jnp.sqrt((eta * (1 - eta)) / 1000) 56 | plt.axhspan( 57 | 1 - (eta + 2 * eta_err), 1 - (eta - 2 * eta_err), color=f"C{i}", alpha=0.1 58 | ) 59 | 60 | eta = 0.9 61 | eta_err = jnp.sqrt((eta * (1 - eta)) / 1000) 62 | plt.text( 63 | 0.905, 64 | 9.8e-2, 65 | r"$\hat{\eta} = %.3f \pm %.3f$" % (eta, 2 * eta_err), 66 | fontsize=12, 67 | va="center", 68 | color="C0", 69 | ha="right", 70 | alpha=0.5, 71 | ) 72 | plt.text( 73 | 0.905, 74 | 9.8e-2, 75 | r"$\hat{\eta} = %.3f \pm %.3f$" % (eta, 2 * eta_err), 76 | fontsize=12, 77 | va="center", 78 | color="C2", 79 | ha="right", 80 | alpha=0.5, 81 | ) 82 | eta = 0.993 83 | eta_err = jnp.sqrt((eta * (1 - eta)) / 1000) 84 | plt.text( 85 | 0.905, 86 | 4.5e-3, 87 | r"$\hat{\eta} = %.3f \pm %.3f$" % (eta, 2 * eta_err), 88 | fontsize=12, 89 | va="center", 90 | color="C1", 91 | ha="right", 92 | ) 93 | 94 | # plt.text(0.954, 1e-2, r"$m_*$") 95 | # plt.text(0.87, 1.3e-1, r"$1 - \eta$") 96 | 97 | plt.xlim(0.7, 1.0) 98 | plt.ylim(1e-3, 1.2) 99 | plt.yscale("log") 100 | plt.legend(loc="upper left", frameon=False) 101 | plt.xlabel("Effectualness") 102 | plt.ylabel("CDF") 103 | 104 | # plt.tight_layout() 105 | plt.savefig("figures/bank-effs.pdf", bbox_inches="tight") 106 | 107 | 108 | if __name__ == "__main__": 109 | run() 110 | -------------------------------------------------------------------------------- /scripts/plot_density.py: -------------------------------------------------------------------------------- 1 | from itertools import product 2 | from math import log10 3 | from operator import itemgetter 4 | 5 | import click 6 | import jax 7 | import jax.numpy as jnp 8 | import matplotlib.cm as cm 9 | import matplotlib.pyplot as plt 10 | from matplotlib import colors 11 | from scipy.spatial import Voronoi, voronoi_plot_2d 12 | from tqdm import tqdm 13 | 14 | from diffbank.constants import MSUN 15 | from diffbank.metric import get_density 16 | from diffbank.noise import Sn_O3a as Sn 17 | from diffbank.utils import ms_to_Mc_eta 18 | from diffbank.waveforms.twoPN_chirptimes import ( 19 | Amp, 20 | Psi, 21 | analytic_metric, 22 | get_th_boundary_interps, 23 | phys_to_th, 24 | ) 25 | 26 | plt.style.use("../plot_style.mplstyle") 27 | 28 | """ 29 | Compares the analytic metric from arXiv:gr-qc/0604037 with the one computed with 30 | automatic differentiation. 31 | 32 | To reproduce the plot: 33 | 34 | >>> python plot_density.py 35 | 36 | """ 37 | 38 | ##### Frequency settings 39 | f_u = 512.0 # Hz 40 | f_l = 10.0 # Hz 41 | N_fbins = 1000 42 | ##### 43 | 44 | 45 | @click.command() 46 | @click.option("--n-m1s", type=int, default=50) 47 | @click.option("--n-m2s", type=int, default=48) 48 | @click.option("--fig-path", default="figures/density.pdf") 49 | def run(n_m1s, n_m2s, fig_path): 50 | fs = jnp.linspace(f_l, f_u, N_fbins) 51 | 52 | density_fun = lambda theta: get_density(theta, Amp, Psi, fs, Sn) 53 | 54 | # Set parameter grid 55 | m_min = jnp.array(1.0) * MSUN 56 | m_max = jnp.array(20.0) * MSUN 57 | M_max = m_min + m_max 58 | m1s = jnp.geomspace(m_min, m_max, n_m1s) 59 | m2s = jnp.geomspace(m_min, m_max, n_m2s) 60 | m1s, m2s = jnp.array(list(product(m1s, m2s))).T 61 | m1s, m2s = m1s[m1s >= m2s], m2s[m1s >= m2s] # remove redundant systems 62 | m1s, m2s = m1s[m1s + m2s <= M_max], m2s[m1s + m2s <= M_max] 63 | M_chirps, etas = ms_to_Mc_eta(jnp.stack([m1s, m2s])) 64 | 65 | # (th0, th3) boundaries 66 | (th0_min, th0_max), th3_interp_low, th3_interp_high = get_th_boundary_interps( 67 | m_min, m_max, f_l 68 | ) 69 | 70 | thetas = phys_to_th(jnp.stack([M_chirps, etas]), f_l).T # type: ignore 71 | densities = jax.lax.map(density_fun, thetas) 72 | densities_analytic = [] 73 | for i in tqdm(range(thetas.shape[0])): 74 | densities_analytic.append( 75 | jnp.sqrt(jnp.linalg.det(analytic_metric(fs, thetas[i], Sn))) 76 | ) 77 | densities_analytic = jnp.array(densities_analytic) 78 | 79 | # Quantify difference in densities 80 | diffs = jnp.log10(jnp.abs((densities - densities_analytic) / densities_analytic)) 81 | 82 | # Plot! 83 | th0_scale = 1e4 84 | th3_scale = 1e2 85 | thetas_scaled = thetas / jnp.array([th0_scale, th3_scale]) 86 | vor = Voronoi(thetas_scaled) 87 | norm = colors.Normalize(-13.5, -12.5) 88 | cmap = cm.ScalarMappable(norm=norm) 89 | voronoi_plot_2d(vor, show_vertices=False, line_width=0, point_size=0) 90 | for r in range(len(vor.point_region)): 91 | region = vor.regions[vor.point_region[r]] 92 | if -1 not in region: 93 | polygon = [vor.vertices[i] for i in region] 94 | plt.fill(*zip(*polygon), color=cmap.to_rgba(diffs[r])) 95 | 96 | plt.colorbar( 97 | cmap, 98 | label=r"$\log_{10} \left| \frac{\sqrt{|g_\mathrm{AD}|} - \sqrt{|g_\mathrm{ref}|}}{\sqrt{|g_\mathrm{ref}|}} \right|$", 99 | ) 100 | 101 | # Mask outside boundaries 102 | th0_grid = jnp.linspace(th0_min, th0_max, 200) 103 | plt.fill_between( 104 | th0_grid / th0_scale, 105 | th3_interp_low(th0_grid) / th3_scale, 106 | jnp.full_like(th0_grid, -1e3) / th3_scale, 107 | where=jnp.full_like(th0_grid, True), 108 | color="w", 109 | zorder=1.5, 110 | ) 111 | plt.fill_between( 112 | th0_grid / th0_scale, 113 | th3_interp_high(th0_grid) / th3_scale, 114 | jnp.full_like(th0_grid, 1e3) / th3_scale, 115 | where=jnp.full_like(th0_grid, True), 116 | color="w", 117 | zorder=1.5, 118 | ) 119 | 120 | plt.xlabel(r"$\theta_0 / 10^{%i}$" % log10(th0_scale)) 121 | plt.ylabel(r"$\theta_3 / 10^{%i}$" % log10(th3_scale)) 122 | plt.xlim(0.1, 10) 123 | plt.ylim(1, 8) 124 | 125 | # plt.tight_layout() 126 | plt.savefig(fig_path, bbox_inches="tight") 127 | 128 | 129 | if __name__ == "__main__": 130 | run() 131 | -------------------------------------------------------------------------------- /scripts/plot_eta_dist.py: -------------------------------------------------------------------------------- 1 | from math import log 2 | from operator import itemgetter 3 | 4 | import jax.numpy as jnp 5 | import matplotlib.pyplot as plt 6 | from plot_scaling import parse_ps 7 | 8 | plt.style.use("../plot_style.mplstyle") 9 | 10 | """ 11 | Plots distribution of `eta` (coverage) values for a set of 3.5PN-2D banks. 12 | 13 | The requisite banks are generated by running `job-threePN-coverage.sh` and 14 | `job-threePN-stochastic.sh`. 15 | 16 | To reproduce the plot: 17 | 18 | >>> python plot_eta_dist.py 19 | 20 | """ 21 | 22 | mm = 0.95 23 | eta = 0.9 24 | eta_est_s, eta_est_err_s, n_templates_s = itemgetter( 25 | "eta_est", "eta_est_err", "n_templates" 26 | )( 27 | jnp.load( 28 | f"banks/3pn-stochastic-1000-mm={mm}-eta_star={eta}-n_eff=500.npz", 29 | allow_pickle=True, 30 | )["bank"].item() 31 | ) 32 | 33 | 34 | def plot_eta_dist(): 35 | # Load random banks 36 | eta_ests = [] 37 | eta_est_errs = [] 38 | n_templatess = [] 39 | for i in range(100, 131): 40 | if i == 119: # this bank got interrupted 41 | continue 42 | 43 | eta_est, eta_est_err, n_templates = itemgetter( 44 | "eta_est", "eta_est_err", "n_templates" 45 | )( 46 | jnp.load( 47 | f"banks/3pn-random-{i}-mm={mm}-eta_star={eta}-n_eff=500.npz", 48 | allow_pickle=True, 49 | )["bank"].item() 50 | ) 51 | eta_ests.append(eta_est) 52 | eta_est_errs.append(eta_est_err) 53 | n_templatess.append(n_templates) 54 | 55 | eta_ests = jnp.array(eta_ests) 56 | eta_est_errs = jnp.array(eta_est_errs) 57 | n_templatess = jnp.array(n_templatess) 58 | 59 | # Random banks 60 | plt.errorbar( 61 | n_templatess, 62 | eta_ests, 63 | yerr=2 * eta_est_errs, 64 | fmt=".", 65 | c="C0", 66 | label=r"Random ($\pm 2\sigma$)", 67 | ) 68 | # Stochastic bank 69 | plt.errorbar( 70 | n_templates_s, 71 | eta_est_s, 72 | yerr=2 * eta_est_err_s, 73 | fmt=".", 74 | c="C1", 75 | label=r"Stochastic ($\pm 2\sigma$)", 76 | ) 77 | 78 | # Target eta 79 | plt.axhline(0.9, color="k", linestyle="--", linewidth=1) 80 | 81 | plt.xlabel(r"$N_T$") 82 | plt.ylabel(r"$\eta$") 83 | plt.legend(loc="upper left", frameon=False) 84 | # plt.tight_layout() 85 | plt.savefig("figures/eta-dist.pdf", bbox_inches="tight") 86 | 87 | 88 | if __name__ == "__main__": 89 | plot_eta_dist() 90 | -------------------------------------------------------------------------------- /scripts/plot_neff_scaling.py: -------------------------------------------------------------------------------- 1 | import os 2 | from operator import itemgetter 3 | 4 | import jax.numpy as jnp 5 | import matplotlib.pyplot as plt 6 | 7 | from diffbank.bank import Bank 8 | from diffbank.noise import Sn_LIGOI as Sn 9 | from diffbank.utils import get_m1_m2_sampler 10 | from diffbank.waveforms.threePN_simple import Psi, amp 11 | 12 | plt.style.use("../plot_style.mplstyle") 13 | 14 | BANK_DIR = "../scripts/banks/" 15 | START_SEED = 501 16 | N_BANKS = 50 17 | ETA = 0.9 18 | N_EFFS = jnp.array( 19 | [ 20 | 10, 21 | 11, 22 | 13, 23 | 15, 24 | 17, 25 | 20, 26 | 23, 27 | 26, 28 | 30, 29 | 35, 30 | 40, 31 | 47, 32 | 54, 33 | 62, 34 | 71, 35 | 82, 36 | 95, 37 | 109, 38 | 126, 39 | 145, 40 | 167, 41 | 193, 42 | 222, 43 | 255, 44 | 294, 45 | 339, 46 | 390, 47 | 449, 48 | 517, 49 | 596, 50 | 686, 51 | 790, 52 | 910, 53 | 1048, 54 | 1206, 55 | 1389, 56 | 1599, 57 | 1842, 58 | 2120, 59 | 2442, 60 | 2811, 61 | 3237, 62 | 3727, 63 | 4291, 64 | 4941, 65 | 5689, 66 | 6551, 67 | 7543, 68 | 8685, 69 | 10000, 70 | ] 71 | ) 72 | sampler = get_m1_m2_sampler((1.0, 3.0), (1.0, 3.0)) 73 | 74 | 75 | def run(): 76 | # Load 77 | ns = [] 78 | eta_ests = [] 79 | for i in range(N_BANKS): 80 | bank = Bank.load( 81 | os.path.join( 82 | BANK_DIR, 83 | f"3pn-random-{START_SEED + i}-mm=0.95-eta_star={ETA}-n_eff={N_EFFS[i]}.npz", 84 | ), 85 | amp, 86 | Psi, 87 | Sn, 88 | sampler, 89 | ) 90 | ns.append(bank.n_templates) 91 | eta_ests.append(bank.eta_est) 92 | 93 | eta_est_errs = jnp.sqrt(ETA * (1 - ETA) / (N_EFFS - 1)) 94 | 95 | n_s = Bank.load( 96 | os.path.join( 97 | BANK_DIR, "3pn-stochastic-1000-mm=0.95-eta_star=0.9-n_eff=500.npz" 98 | ), 99 | amp, 100 | Psi, 101 | Sn, 102 | sampler, 103 | ).n_templates 104 | 105 | # Plot 106 | fig, axs = plt.subplots(1, 2, figsize=(10, 4.5)) 107 | 108 | ax = axs[0] 109 | ax.scatter(N_EFFS, eta_ests) 110 | ax.axhline(ETA, color="k", linestyle="--") 111 | for n_sigma in [1, 2]: 112 | ax.fill_between( 113 | N_EFFS, 114 | ETA - n_sigma * eta_est_errs, 115 | ETA + n_sigma * eta_est_errs, 116 | color="k", 117 | alpha=0.1, 118 | linewidth=0, 119 | ) 120 | ax.set_ylabel(r"$\eta$") 121 | ax.set_ylim(0.68, 1.0) 122 | 123 | ax = axs[1] 124 | ax.scatter(N_EFFS, ns) 125 | ax.axhline(n_s, color="C1") 126 | ax.text(1.3e3, 1500, "Stochastic", color="C1") 127 | ax.set_ylabel(r"$N_T$") 128 | ax.set_ylim(900) 129 | 130 | for ax in axs: 131 | ax.set_xscale("log") 132 | ax.set_xlabel(r"$n_\mathrm{eff}$") 133 | ax.set_xlim(N_EFFS[0], N_EFFS[-1]) 134 | 135 | fig.tight_layout() 136 | fig.savefig("figures/neff-scaling.pdf") 137 | 138 | 139 | if __name__ == "__main__": 140 | run() 141 | -------------------------------------------------------------------------------- /scripts/plot_scalar_curvature.py: -------------------------------------------------------------------------------- 1 | from itertools import product 2 | 3 | import click 4 | from diffjeom import get_ricci_scalar 5 | import jax 6 | import jax.numpy as jnp 7 | import matplotlib.pyplot as plt 8 | from scipy.interpolate import griddata 9 | 10 | from diffbank.constants import MSUN 11 | from diffbank.metric import get_g 12 | from diffbank.noise import Sn_O3a as Sn 13 | from diffbank.utils import ms_to_Mc_eta 14 | from diffbank.waveforms.twoPN_chirptimes import ( 15 | Amp, 16 | Psi, 17 | get_th_boundary_interps, 18 | phys_to_th, 19 | ) 20 | 21 | plt.style.use("../plot_style.mplstyle") 22 | 23 | """ 24 | Plots the scalar curvature (Ricci scalar) for the metric for the waveform in 25 | https://arxiv.org/abs/gr-qc/0604037. 26 | 27 | To reproduce: 28 | 29 | >>> python plot_scalar_curvature.py 30 | 31 | """ 32 | 33 | ##### Frequency settings 34 | f_u = 512.0 # Hz 35 | f_l = 10.0 # Hz 36 | N_fbins = 1000 37 | ##### 38 | 39 | 40 | @click.command() 41 | @click.option("--n-m1s", type=int, default=200) 42 | @click.option("--n-m2s", type=int, default=100) 43 | @click.option("--n-th0s", type=int, default=200) 44 | @click.option("--n-th3s", type=int, default=100) 45 | @click.option("--fig-path", default="figures/scalar-curvature.pdf") 46 | def run(n_m1s, n_m2s, n_th0s, n_th3s, fig_path): 47 | fs = jnp.linspace(f_l, f_u, N_fbins) 48 | 49 | g_fun = lambda theta: get_g(theta, Amp, Psi, fs, Sn) 50 | 51 | # Set parameter grid 52 | m_min = jnp.array(1.0) * MSUN 53 | m_max = jnp.array(20.0) * MSUN 54 | M_max = m_min + m_max 55 | m1s = jnp.geomspace(m_min, m_max, n_m1s) 56 | m2s = jnp.geomspace(m_min, m_max, n_m2s) 57 | m1s, m2s = jnp.array(list(product(m1s, m2s))).T 58 | m1s, m2s = m1s[m1s >= m2s], m2s[m1s >= m2s] # remove redundant systems 59 | m1s, m2s = m1s[m1s + m2s <= M_max], m2s[m1s + m2s <= M_max] 60 | M_chirps, etas = ms_to_Mc_eta(jnp.stack([m1s, m2s])) 61 | 62 | # (th0, th3) boundaries 63 | (th0_min, th0_max), bound_itp_low, bound_itp_high = get_th_boundary_interps( 64 | m_min, m_max, f_l 65 | ) 66 | 67 | # Plotting configuration 68 | vmin = -10.0 69 | vmax = -7.5 70 | levels = jnp.linspace(vmin, vmax, 60) 71 | cbar_ticks = jnp.arange(vmin, vmax + 0.05, 0.5) 72 | 73 | thetas = phys_to_th(jnp.stack([M_chirps, etas]), f_l).T # type: ignore 74 | Rss = jax.vmap(lambda x: get_ricci_scalar(x, g_fun))(thetas) 75 | 76 | # Plot! 77 | th0s = jnp.linspace(thetas[:, 0].min(), thetas[:, 0].max(), n_th0s) 78 | th3s = jnp.linspace(thetas[:, 1].min(), thetas[:, 1].max(), n_th3s) 79 | 80 | cs = plt.contourf( 81 | th0s / 1e4, 82 | th3s / 1e2, 83 | jnp.clip( 84 | griddata( 85 | thetas[:, :2], 86 | jnp.log10(jnp.abs(Rss)), 87 | jnp.stack(jnp.meshgrid(th0s, th3s)).reshape([2, -1]).T, 88 | ).reshape([len(th3s), len(th0s)]), 89 | vmin, 90 | vmax, 91 | ), 92 | levels=levels, 93 | cmap="viridis", 94 | ) 95 | plt.colorbar(cs, label=r"$\log_{10}(|R|)$", ticks=cbar_ticks) 96 | 97 | # Mask outside boundaries 98 | th0_grid = jnp.linspace(th0_min, th0_max, 200) 99 | plt.fill_between( 100 | th0_grid / 1e4, 101 | bound_itp_low(th0_grid) / 1e2, 102 | jnp.full_like(th0_grid, -1e3) / 1e2, 103 | where=jnp.full_like(th0_grid, True), 104 | color="w", 105 | ) 106 | plt.fill_between( 107 | th0_grid / 1e4, 108 | bound_itp_high(th0_grid) / 1e2, 109 | jnp.full_like(th0_grid, 1e3) / 1e2, 110 | where=jnp.full_like(th0_grid, True), 111 | color="w", 112 | ) 113 | 114 | plt.xlabel(r"$\theta_0 / 10^4$") 115 | plt.ylabel(r"$\theta_3 / 10^2$") 116 | plt.xlim(0.0, 10.0) 117 | plt.ylim(0.8, 8) 118 | # plt.tight_layout() 119 | plt.savefig(fig_path, bbox_inches="tight") 120 | 121 | 122 | if __name__ == "__main__": 123 | run() 124 | -------------------------------------------------------------------------------- /scripts/plot_scaling.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | from math import log 4 | from typing import Dict, Tuple 5 | 6 | import jax.numpy as jnp 7 | import matplotlib.pyplot as plt 8 | 9 | from diffbank.bank import Bank 10 | from diffbank.noise import Sn_LIGOI 11 | from diffbank.utils import get_m1_m2_sampler 12 | from diffbank.waveforms.threePN_simple import Psi, amp 13 | 14 | plt.style.use("../plot_style.mplstyle") 15 | 16 | # scatter kwargs 17 | kwargs_s = dict(marker="X", c="C2", s=25, label="Stochastic") 18 | kwargs_r = dict(c="C0", s=25, label="Random") 19 | kwargs_r_est = dict(c="none", edgecolor="C1", s=25, label="Random est. (const. $q$)") 20 | kwargs_s_est = dict( 21 | marker="X", c="none", edgecolor="C3", s=25, label="Stochastic est. (const. $q$)" 22 | ) 23 | 24 | 25 | """ 26 | Plots scaling properties of the 3.5PN-2D bank. 27 | 28 | The requisite banks and data files are generated by running the following 29 | scripts: 30 | - job-threePN-est-p.sh: estimates the covering probability `p`. 31 | - job-threePN-scaling.sh: generates banks with a range of different `mm` 32 | and `eta` values. 33 | 34 | Generates the following plots: 35 | - figures/threePN-n_templates-scaling.pdf 36 | - figures/threePN-time-scaling.pdf 37 | 38 | To reproduce the plots: 39 | 40 | >>> python plot_scaling.py 41 | 42 | """ 43 | 44 | 45 | def parse_ps() -> Dict[float, Tuple[float, float]]: 46 | """ 47 | Parses output from job-threePN-est-p.sh. 48 | 49 | Returns 50 | Dict mapping `mm` to an MC estimate of the covering probability `p` and 51 | associated error `p_err`. 52 | """ 53 | with open("../scripts/threePN-p.txt", "r") as f: 54 | lines = f.readlines() 55 | 56 | ps = {} 57 | for i in range(len(lines) // 2): 58 | mm = float(re.search("mm=(\d\.\d*)", lines[2 * i]).group(1)) # type: ignore 59 | 60 | match = re.search("p = (\d\.\d*) \+\/\- (\d\.\d*)", lines[2 * i + 1]) # type: ignore 61 | p = float(match.group(1)) # type: ignore 62 | p_err = float(match.group(2)) # type: ignore 63 | 64 | ps[mm] = (p, p_err) 65 | 66 | return ps 67 | 68 | 69 | PS = parse_ps() 70 | FS = jnp.linspace(20.0, 2000.0, 1000) 71 | M_RANGE = (1.4, 5.0) 72 | sampler = get_m1_m2_sampler(M_RANGE, M_RANGE) 73 | ETA_REF = 0.9 74 | MM_REF = 0.9 75 | MMS = [0.95, 0.90, 0.85, 0.80, 0.75] 76 | ETAS = [0.975, 0.95, 0.925, 0.900, 0.875, 0.850] 77 | 78 | 79 | def cr_pred(p, nr, n_eff=1000): 80 | cr = n_eff * (1 - (1 - p) ** nr) / p 81 | return cr 82 | 83 | 84 | def cs_pred(p, ns, n_eff=1000): 85 | cr = cr_pred(p, ns, n_eff) 86 | rejection_cost = (1 - p) ** (1 - ns) * (ns * p + (1 - p) ** ns - 1) / p ** 2 87 | return cr + rejection_cost 88 | 89 | 90 | def predict_ns(): 91 | n_ests, n_est_errs = {}, {} 92 | 93 | for mm in MMS: 94 | p, p_err = PS[mm] 95 | n_ests[(mm, ETA_REF)] = log(1 - ETA_REF) / log(1 - p) 96 | n_est_errs[(mm, ETA_REF)] = ( 97 | log(1 - ETA_REF) / ((1 - p) * log(1 - p) ** 2) * p_err 98 | ) 99 | 100 | for eta in ETAS: 101 | p, p_err = PS[MM_REF] 102 | n_ests[(MM_REF, eta)] = log(1 - eta) / log(1 - p) 103 | n_est_errs[(MM_REF, eta)] = log(1 - eta) / ((1 - p) * log(1 - p) ** 2) * p_err 104 | 105 | return n_ests, n_est_errs 106 | 107 | 108 | def load_ns(): 109 | ns, ns_s = {}, {} 110 | 111 | # Varying mm 112 | for seed, mm in enumerate(MMS, 5): 113 | # Random 114 | path = os.path.join( 115 | "../scripts/threePN-banks-scaling", 116 | f"3pn-random-{seed}-mm={mm}-eta_star={ETA_REF}-n_eff=1000.npz", 117 | ) 118 | ns[(mm, ETA_REF)] = Bank.load(path, amp, Psi, Sn_LIGOI, sampler).n_templates 119 | # Stochastic 120 | path = os.path.join( 121 | "../scripts/threePN-banks-scaling", 122 | f"3pn-stochastic-{seed + 5}-mm={mm}-eta_star={ETA_REF}-n_eff=1000.npz", 123 | ) 124 | ns_s[(mm, ETA_REF)] = Bank.load(path, amp, Psi, Sn_LIGOI, sampler).n_templates 125 | 126 | # Varying eta 127 | for seed, eta in enumerate(ETAS, 15): 128 | # Random 129 | path = os.path.join( 130 | "../scripts/threePN-banks-scaling", 131 | f"3pn-random-{seed}-mm={MM_REF}-eta_star={eta}-n_eff=1000.npz", 132 | ) 133 | ns[(MM_REF, eta)] = Bank.load(path, amp, Psi, Sn_LIGOI, sampler).n_templates 134 | # Stochastic 135 | path = os.path.join( 136 | "../scripts/threePN-banks-scaling", 137 | f"3pn-stochastic-{seed + 6}-mm={MM_REF}-eta_star={eta}-n_eff=1000.npz", 138 | ) 139 | ns_s[(MM_REF, eta)] = Bank.load(path, amp, Psi, Sn_LIGOI, sampler).n_templates 140 | 141 | return ns, ns_s 142 | 143 | 144 | def load_runtimes(): 145 | runtimes = {} 146 | filenames = os.listdir("../scripts/threePN-outputs-scaling/") 147 | 148 | # Parse tqdm output to get timing information 149 | for fn in filenames: 150 | kind = fn.split("-")[1] 151 | mm = float(re.search("mm=(\d*\.\d*)-", fn).group(1)) # type: ignore 152 | eta = float(re.search("eta_star=(\d*\.\d*)\.txt", fn).group(1)) # type: ignore 153 | 154 | with open(os.path.join("../scripts/threePN-outputs-scaling", fn)) as f: 155 | raw = f.read() 156 | last_pbar = raw.split("\n")[-4] 157 | 158 | # Convert to runtime [s] 159 | search_result = re.search("\[(.*)<", last_pbar) # type: ignore 160 | if search_result is not None: 161 | raw_time = search_result.group(1) 162 | else: 163 | raw_time = re.search("\[(.*?),", last_pbar).group(1) # type: ignore 164 | raw_time = raw_time.split(":") 165 | assert len(raw_time) <= 3 166 | time = sum([float(rt) * 60 ** i for i, rt in enumerate(reversed(raw_time))]) 167 | 168 | runtimes[(kind, mm, eta)] = time 169 | 170 | return runtimes 171 | 172 | 173 | def plot_n_templates_scaling(axs): 174 | ns, ns_s = load_ns() 175 | n_ests = predict_ns()[0] 176 | 177 | ax = axs[0] 178 | ax.scatter(MMS, [ns[(mm, ETA_REF)] for mm in MMS], **kwargs_r) 179 | ax.scatter(MMS, [n_ests[(mm, ETA_REF)] for mm in MMS], **kwargs_r_est) 180 | ax.scatter(MMS, [ns_s[(mm, ETA_REF)] for mm in MMS], **kwargs_s) 181 | 182 | ax = axs[1] 183 | ax.scatter(ETAS, [ns[(MM_REF, eta)] for eta in ETAS], **kwargs_r) 184 | ax.scatter(ETAS, [n_ests[(MM_REF, eta)] for eta in ETAS], **kwargs_r_est) 185 | ax.scatter(ETAS, [ns_s[(MM_REF, eta)] for eta in ETAS], **kwargs_s) 186 | ax.set_ylim(0, 2100) 187 | 188 | 189 | def plot_time_scaling(axs): 190 | runtimes = load_runtimes() 191 | ns, ns_s = load_ns() 192 | 193 | # mm scaling 194 | mms = jnp.array(MMS[::-1]) 195 | p_ests = jnp.array([PS[mm][0] for mm in mms]) 196 | # eta scaling 197 | etas = jnp.array(ETAS[::-1]) 198 | p_est_ref = PS[MM_REF][0] 199 | 200 | ax = axs[0] 201 | c_ss = jnp.array([runtimes[("stochastic", mm, ETA_REF)] for mm in mms]) 202 | c_rs = jnp.array([runtimes[("random", mm, ETA_REF)] for mm in mms]) 203 | nr_arr = jnp.array([ns[(mm, ETA_REF)] for mm in mms]) 204 | ns_arr = jnp.array([ns_s[(mm, ETA_REF)] for mm in mms]) 205 | ax.scatter(mms, c_rs, **kwargs_r) 206 | ax.scatter(mms, c_ss, **kwargs_s) 207 | c_rs_pred = cr_pred(p_ests, nr_arr) 208 | c_ss_pred = cs_pred(p_ests, ns_arr) 209 | ax.scatter(mms, c_rs_pred * c_rs[0] / c_rs_pred[0], **kwargs_r_est) 210 | ax.scatter(mms, c_ss_pred * c_ss[0] / c_ss_pred[0], **kwargs_s_est) 211 | 212 | ax = axs[1] 213 | c_ss = jnp.array([runtimes[("stochastic", MM_REF, eta)] for eta in etas]) 214 | c_rs = jnp.array([runtimes[("random", MM_REF, eta)] for eta in etas]) 215 | nr_arr = jnp.array([ns[(MM_REF, eta)] for eta in etas]) 216 | ns_arr = jnp.array([ns_s[(MM_REF, eta)] for eta in etas]) 217 | ax.scatter(etas, c_rs, **kwargs_r) 218 | ax.scatter(etas, c_ss, **kwargs_s) 219 | c_rs_pred = cr_pred(p_est_ref, nr_arr) 220 | c_ss_pred = cs_pred(p_est_ref, ns_arr) 221 | ax.scatter(etas, c_rs_pred * c_rs[0] / c_rs_pred[0], **kwargs_r_est) 222 | ax.scatter(etas, c_ss_pred * c_ss[0] / c_ss_pred[0], **kwargs_s_est) 223 | 224 | ax.set_yscale("log") 225 | 226 | 227 | def run(): 228 | fig, axes = plt.subplots(2, 2, sharey="row", figsize=(10, 6)) 229 | 230 | plot_n_templates_scaling([axes[0, 0], axes[0, 1]]) 231 | plot_time_scaling([axes[1, 0], axes[1, 1]]) 232 | 233 | axes[1, 0].set_xlabel(r"$\mathrm{m}_\ast = 1 - \mathrm{m}_{\mathrm{mis},\ast}$") 234 | axes[1, 1].set_xlabel(r"$\eta$") 235 | axes[0, 0].set_ylabel(r"$N_T$") 236 | axes[1, 0].set_ylabel(r"$C$ [arb.]") 237 | axes[0, 0].set_title(r"$\eta = %g$" % ETA_REF) 238 | axes[0, 1].set_title( 239 | r"$\mathrm{m}_\ast = 1 - \mathrm{m}_{\mathrm{mis},\ast} = %g$" % MM_REF 240 | ) 241 | 242 | # Customize legend 243 | ax = axes[1, 0] 244 | handles, labels = ax.get_legend_handles_labels() 245 | axes[0, 0].legend(handles, labels, loc="upper left", fontsize=14) 246 | 247 | fig.tight_layout() 248 | fig.savefig("figures/scaling.pdf") 249 | 250 | 251 | if __name__ == "__main__": 252 | run() 253 | -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-random-15-mm=0.9-eta_star=0.975-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-random-15-mm=0.9-eta_star=0.975-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-random-16-mm=0.9-eta_star=0.95-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-random-16-mm=0.9-eta_star=0.95-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-random-17-mm=0.9-eta_star=0.925-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-random-17-mm=0.9-eta_star=0.925-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-random-18-mm=0.9-eta_star=0.9-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-random-18-mm=0.9-eta_star=0.9-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-random-19-mm=0.9-eta_star=0.875-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-random-19-mm=0.9-eta_star=0.875-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-random-20-mm=0.9-eta_star=0.85-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-random-20-mm=0.9-eta_star=0.85-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-random-5-mm=0.95-eta_star=0.9-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-random-5-mm=0.95-eta_star=0.9-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-random-6-mm=0.9-eta_star=0.9-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-random-6-mm=0.9-eta_star=0.9-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-random-7-mm=0.85-eta_star=0.9-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-random-7-mm=0.85-eta_star=0.9-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-random-8-mm=0.8-eta_star=0.9-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-random-8-mm=0.8-eta_star=0.9-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-random-9-mm=0.75-eta_star=0.9-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-random-9-mm=0.75-eta_star=0.9-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-stochastic-10-mm=0.95-eta_star=0.9-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-stochastic-10-mm=0.95-eta_star=0.9-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-stochastic-11-mm=0.9-eta_star=0.9-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-stochastic-11-mm=0.9-eta_star=0.9-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-stochastic-12-mm=0.85-eta_star=0.9-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-stochastic-12-mm=0.85-eta_star=0.9-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-stochastic-13-mm=0.8-eta_star=0.9-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-stochastic-13-mm=0.8-eta_star=0.9-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-stochastic-14-mm=0.75-eta_star=0.9-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-stochastic-14-mm=0.75-eta_star=0.9-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-stochastic-21-mm=0.9-eta_star=0.975-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-stochastic-21-mm=0.9-eta_star=0.975-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-stochastic-22-mm=0.9-eta_star=0.95-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-stochastic-22-mm=0.9-eta_star=0.95-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-stochastic-23-mm=0.9-eta_star=0.925-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-stochastic-23-mm=0.9-eta_star=0.925-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-stochastic-24-mm=0.9-eta_star=0.9-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-stochastic-24-mm=0.9-eta_star=0.9-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-stochastic-25-mm=0.9-eta_star=0.875-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-stochastic-25-mm=0.9-eta_star=0.875-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-banks-scaling/3pn-stochastic-26-mm=0.9-eta_star=0.85-n_eff=1000.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/scripts/threePN-banks-scaling/3pn-stochastic-26-mm=0.9-eta_star=0.85-n_eff=1000.npz -------------------------------------------------------------------------------- /scripts/threePN-outputs-scaling/threePN-stochastic-14-mm=0.75-eta_star=0.9.txt: -------------------------------------------------------------------------------- 1 | 2021-11-12 17:57:08.147969: W external/org_tensorflow/tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /sw/arch/Debian10/EB_production/2019/software/CUDA/10.1.243-GCC-8.3.0/nvvm/lib64:/sw/arch/Debian10/EB_production/2019/software/CUDA/10.1.243-GCC-8.3.0/extras/CUPTI/lib64:/sw/arch/Debian10/EB_production/2019/software/CUDA/10.1.243-GCC-8.3.0/lib64:/sw/arch/Debian10/EB_production/2019/software/binutils/2.32-GCCcore-8.3.0/lib:/sw/arch/Debian10/EB_production/2019/software/zlib/1.2.11-GCCcore-8.3.0/lib:/sw/arch/Debian10/EB_production/2019/software/GCCcore/8.3.0/lib64:/sw/arch/Debian10/EB_production/2019/software/GCCcore/8.3.0/lib 2 | 2021-11-12 17:57:08.921020: W external/org_tensorflow/tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /sw/arch/Debian10/EB_production/2019/software/CUDA/10.1.243-GCC-8.3.0/nvvm/lib64:/sw/arch/Debian10/EB_production/2019/software/CUDA/10.1.243-GCC-8.3.0/extras/CUPTI/lib64:/sw/arch/Debian10/EB_production/2019/software/CUDA/10.1.243-GCC-8.3.0/lib64:/sw/arch/Debian10/EB_production/2019/software/binutils/2.32-GCCcore-8.3.0/lib:/sw/arch/Debian10/EB_production/2019/software/zlib/1.2.11-GCCcore-8.3.0/lib:/sw/arch/Debian10/EB_production/2019/software/GCCcore/8.3.0/lib64:/sw/arch/Debian10/EB_production/2019/software/GCCcore/8.3.0/lib 3 | 2021-11-12 17:57:08.937011: W external/org_tensorflow/tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /sw/arch/Debian10/EB_production/2019/software/CUDA/10.1.243-GCC-8.3.0/nvvm/lib64:/sw/arch/Debian10/EB_production/2019/software/CUDA/10.1.243-GCC-8.3.0/extras/CUPTI/lib64:/sw/arch/Debian10/EB_production/2019/software/CUDA/10.1.243-GCC-8.3.0/lib64:/sw/arch/Debian10/EB_production/2019/software/binutils/2.32-GCCcore-8.3.0/lib:/sw/arch/Debian10/EB_production/2019/software/zlib/1.2.11-GCCcore-8.3.0/lib:/sw/arch/Debian10/EB_production/2019/software/GCCcore/8.3.0/lib64:/sw/arch/Debian10/EB_production/2019/software/GCCcore/8.3.0/lib 4 | 2021-11-12 17:57:09.884192: W external/org_tensorflow/tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /sw/arch/Debian10/EB_production/2019/software/CUDA/10.1.243-GCC-8.3.0/nvvm/lib64:/sw/arch/Debian10/EB_production/2019/software/CUDA/10.1.243-GCC-8.3.0/extras/CUPTI/lib64:/sw/arch/Debian10/EB_production/2019/software/CUDA/10.1.243-GCC-8.3.0/lib64:/sw/arch/Debian10/EB_production/2019/software/binutils/2.32-GCCcore-8.3.0/lib:/sw/arch/Debian10/EB_production/2019/software/zlib/1.2.11-GCCcore-8.3.0/lib:/sw/arch/Debian10/EB_production/2019/software/GCCcore/8.3.0/lib64:/sw/arch/Debian10/EB_production/2019/software/GCCcore/8.3.0/lib 5 | 2021-11-12 17:57:09.884259: W external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_driver.cc:269] failed call to cuInit: UNKNOWN ERROR (303) 6 | WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.) 7 | 0%| | 0/900 [00:00 mm 73 | ) 74 | 75 | p, p_err = covereds.mean(), covereds.std() / sqrt(len(covereds)) 76 | print(f"Bank: {bank.name}") 77 | print(f"\tp = {p} +/- {p_err}") 78 | 79 | 80 | if __name__ == "__main__": 81 | run() 82 | -------------------------------------------------------------------------------- /scripts/threePN_est_p_variation.py: -------------------------------------------------------------------------------- 1 | from math import sqrt 2 | from operator import itemgetter 3 | 4 | import click 5 | import jax 6 | import jax.numpy as jnp 7 | import matplotlib.pyplot as plt 8 | from jax import random 9 | from scipy.optimize import minimize_scalar 10 | from tqdm.auto import tqdm 11 | 12 | from diffbank.bank import Bank 13 | from diffbank.noise import Sn_LIGOI 14 | from diffbank.utils import ( 15 | gen_template_rejection, 16 | gen_templates_rejection, 17 | get_m1_m2_sampler, 18 | ) 19 | from diffbank.waveforms.threePN_simple import Psi, amp 20 | 21 | """ 22 | Estimates variation of p over a 3.5PN bank. 23 | """ 24 | 25 | 26 | @click.command() 27 | @click.option("--seed", default=1000, help="PRNG seed") 28 | @click.option("--mm", default=0.95, help="minimum match") 29 | @click.option("--n-pts", default=200, help="number of points at which to estimate p") 30 | @click.option("--n-templates", default=10000, help="number of templates per point") 31 | def run(seed, mm, n_pts, n_templates): 32 | key = random.PRNGKey(seed) 33 | m_star = 1 - mm 34 | 35 | fs = jnp.linspace(20.0, 2000.0, 1000) 36 | m_range = (1.4, 5.0) 37 | sampler = get_m1_m2_sampler(m_range, m_range) 38 | 39 | bank = Bank( 40 | amp, 41 | Psi, 42 | fs, 43 | Sn_LIGOI, 44 | m_star, 45 | 0.9, 46 | sampler, 47 | name=f"3pn-???-mm={mm}-eta_star=???-n_eff=???", 48 | ) 49 | 50 | # Get max density 51 | fun = lambda m1: -bank.density_fun(jnp.array([m1, m_range[0]])) 52 | res = minimize_scalar(fun, bracket=m_range, bounds=m_range) 53 | assert res.success 54 | theta_dmax = jnp.array([res.x, m_range[0]]) 55 | ratio_max = bank.density_fun(theta_dmax) 56 | bank.ratio_max = ratio_max 57 | 58 | gen_templates = jax.jit( 59 | lambda key: gen_templates_rejection( 60 | key, 61 | n_templates, 62 | bank.ratio_max, 63 | bank.density_fun, 64 | bank.sample_base, 65 | bank.density_fun_base, 66 | ) 67 | ) 68 | 69 | gen_template = jax.jit( 70 | lambda key: gen_template_rejection( 71 | key, 72 | bank.ratio_max, 73 | bank.density_fun, 74 | bank.sample_base, 75 | bank.density_fun_base, 76 | ) 77 | ) 78 | 79 | eff_fun = jax.jit(lambda template, ep: bank.match_fun(template, ep)) 80 | 81 | # Generate points 82 | key, pt_keys = itemgetter(0, slice(1, -1))(random.split(key, n_pts + 1)) 83 | eff_pts = jax.lax.map(gen_template, pt_keys) 84 | 85 | # Estimate p for each of them 86 | ps = [] 87 | p_errs = [] 88 | 89 | with tqdm(eff_pts) as pbar: 90 | for ep in pbar: 91 | key, template_key = random.split(key) 92 | templates = gen_templates(template_key) 93 | covereds = ( 94 | jax.lax.map(lambda template: eff_fun(template, ep), templates) > mm 95 | ) 96 | p, p_err = covereds.mean(), covereds.std() / sqrt(n_templates) 97 | ps.append(p) 98 | p_errs.append(p_err) 99 | pbar.set_postfix(dict(p=p, p_err=p_err)) 100 | 101 | ps = jnp.array(ps) 102 | p_errs = jnp.array(p_errs) 103 | 104 | print("min, max, mean:", ps.min(), ps.max(), ps.mean()) 105 | 106 | # Save results 107 | jnp.savez( 108 | f"threePN-est-p-variation-{mm}.npz", eff_pts=eff_pts, ps=ps, p_errs=p_errs 109 | ) 110 | 111 | plt.scatter(*eff_pts.T, c=ps) 112 | plt.colorbar() 113 | plt.savefig(f"threePN-est-p-variation-{mm}.pdf") 114 | 115 | 116 | if __name__ == "__main__": 117 | run() 118 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [metadata] 2 | name = diffbank 3 | version = 0.0.1 4 | description = Automatic gravitational wave template bank generator 5 | long_description = file: README.md 6 | long_description_content_type = text/markdown 7 | keywords = gravitational waves, template banks, jax 8 | url = https://diffbank.readthedocs.io/en/latest/ 9 | project_urls = 10 | Bug Tracker = https://github.com/adam-coogan/diffbank 11 | classifiers = 12 | Programming Language :: Python :: 3 13 | License :: OSI Approved :: MIT License 14 | Operating System :: OS Independent 15 | 16 | [options] 17 | package_dir = 18 | = src 19 | packages = find: 20 | python_requires = >=3.6 21 | install_requires = 22 | jax 23 | jaxlib 24 | jaxtyping 25 | numpy 26 | tqdm 27 | 28 | [options.package_data] 29 | * = *.dat 30 | 31 | [options.packages.find] 32 | where = src 33 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import setuptools 2 | 3 | setuptools.setup() 4 | -------------------------------------------------------------------------------- /src/diffbank/__init__.py: -------------------------------------------------------------------------------- 1 | from jax.config import config 2 | 3 | config.update("jax_enable_x64", True) 4 | -------------------------------------------------------------------------------- /src/diffbank/bank.py: -------------------------------------------------------------------------------- 1 | r""" 2 | Defines ``Bank``, a class-based interface for managing template banks. 3 | """ 4 | import os 5 | from typing import Callable, Optional, Set, Tuple, Union 6 | 7 | import jax 8 | import jax.numpy as jnp 9 | from jax import random 10 | 11 | from .metric import get_density, get_g, get_gam 12 | from .utils import ( 13 | Array, 14 | PRNGKeyArray, 15 | est_ratio_max, 16 | gen_bank_random, 17 | gen_bank_stochastic, 18 | gen_template_rejection, 19 | gen_templates_rejection, 20 | get_bank_effectualness, 21 | get_eff_pads, 22 | get_match, 23 | get_top_matches_templates, 24 | ) 25 | 26 | 27 | class Bank: 28 | r""" 29 | Class representing a template bank. 30 | """ 31 | 32 | computed_vars: Set[str] = set( 33 | [ 34 | "ratio_max", 35 | "n_templates", 36 | "templates", 37 | "effectualness_points", 38 | "effectualnesses", 39 | "eta_est", 40 | "eta_est_err", 41 | "_eff_pad_low", 42 | "_eff_pad_high", 43 | "_dim", 44 | ] 45 | ) 46 | """ 47 | The names of attributes that are computed in the course of working with a 48 | bank. For example, this includes ``"templates"``. 49 | """ 50 | provided_vars: Set[str] = set( 51 | ["fs", "m_star", "eta", "name",] 52 | ) 53 | """ 54 | The names of attributes containing data the user provides when initialing a 55 | bank 56 | """ 57 | 58 | def __init__( 59 | self, 60 | amp: Callable[[Array, Array], Array], 61 | Psi: Callable[[Array, Array], Array], 62 | fs: Array, 63 | Sn: Callable[[Array], Array], 64 | m_star: Union[float, Array], 65 | eta: Union[float, Array], 66 | sample_base: Callable[[PRNGKeyArray, int], Array], 67 | density_fun_base: Callable[[Array], Array] = lambda _: jnp.array(1.0), 68 | name: str = "test", 69 | ): 70 | r""" 71 | Initializes a :class:`Bank` instance. 72 | 73 | Args: 74 | amp: function returning the amplitude of the Fourier-domain waveform 75 | as a function of frequency and parameters 76 | Psi: function returning the phase of the Fourier-domain waveform as 77 | a function of frequency and parameters 78 | fs: uniformly-spaced grid of frequencies used to compute the 79 | noise-weighted inner product or match 80 | Sn: power spectral density of the detector noise 81 | m_star: the target maximum mismatch for the bank (equal to 82 | ``1 - minimum_match``) 83 | eta: target fraction of parameter space to cover with the bank 84 | sample_base: ``jit``-able base sampler for rejection sampling from 85 | the metric density. Takes a key and number of samples as 86 | arguments. 87 | density_fun_base: ``jit``-able density function for ``sample_base``. 88 | Need not be normalized. 89 | name: an identifier for the bank. 90 | 91 | Raises: 92 | ValueError: if ``len(fs) <= 2`` or ``m_star`` or ``eta`` are outside 93 | the range (0, 1) 94 | """ 95 | # Validation 96 | if len(fs) <= 2: 97 | # Required for padding to work 98 | raise ValueError("length of frequency array must be at least three") 99 | if m_star > 1 or m_star < 0: 100 | raise ValueError("m_star must be in (0, 1)") 101 | if eta > 1 or eta < 0: 102 | raise ValueError("eta must be in (0, 1)") 103 | 104 | self.amp = amp 105 | self.Psi = Psi 106 | self.fs = fs 107 | self.Sn = Sn 108 | self.m_star = m_star 109 | self.eta = eta 110 | self.sample_base = sample_base 111 | self.density_fun_base = density_fun_base 112 | self.name = name 113 | 114 | self.ratio_max: Optional[Array] = None 115 | self.n_templates: Optional[int] = None 116 | self.templates: Optional[Array] = None 117 | self.effectualness_points: Optional[Array] = None 118 | self.effectualnesses: Optional[Array] = None 119 | self.eta_est: Optional[Array] = None 120 | self.eta_est_err: Optional[Array] = None 121 | 122 | # Padding for accurate effectualness calculation 123 | # Length of padded array 124 | self._eff_pad_low, self._eff_pad_high = get_eff_pads(self.fs) 125 | 126 | # Key doesn't matter 127 | self._dim = self.sample_base(random.PRNGKey(1), 1).shape[-1] 128 | 129 | def __str__(self): 130 | r"""Gets a string representation of the bank.""" 131 | return f"Bank(m_star={float(self.m_star)}, eta={float(self.eta)}, dim={self.dim}, name='{self.name}')" 132 | 133 | def __repr__(self): 134 | return str(self) # for now 135 | 136 | @property 137 | def dim(self) -> int: 138 | r""" 139 | The dimensionality of the bank's parameter space. 140 | """ 141 | return self._dim 142 | 143 | def match_fun( 144 | self, 145 | theta1: Array, 146 | theta2: Array, 147 | amp2: Optional[Callable[[Array, Array], Array]] = None, 148 | Psi2: Optional[Callable[[Array, Array], Array]] = None, 149 | ) -> Array: 150 | r""" 151 | Gets the match between a point from the bank's waveform model and a different 152 | point, potentially with a distinct waveform model. See :meth:`diffbank.utils.get_match`. 153 | 154 | Args: 155 | theta1: the point from the same waveform model as the bank 156 | theta2: the other point 157 | amp2: amplitude function for the ``theta2``. If not provided, it is 158 | assumed ``theta2`` has the same waveform model as the bank. 159 | Psi2: phase function for the ``theta2``. If not provided, it is 160 | assumed ``theta2`` has the same waveform model as the bank. 161 | 162 | Returns: 163 | The match between the waveforms for the two points. 164 | """ 165 | return get_match( 166 | theta1, 167 | theta2, 168 | self.amp, 169 | self.Psi, 170 | self.amp if amp2 is None else amp2, 171 | self.Psi if Psi2 is None else Psi2, 172 | self.fs, 173 | self.Sn, 174 | self._eff_pad_low, 175 | self._eff_pad_high, 176 | ) 177 | 178 | def density_fun(self, theta: Array) -> Array: 179 | """ 180 | Computes the determinant of the metric over the intrinsic of the binary, 181 | maximized over the time and phase at coalescence. 182 | See :meth:`diffbank.metric.get_density`. 183 | """ 184 | return get_density(theta, self.amp, self.Psi, self.fs, self.Sn) 185 | 186 | def g_fun(self, theta) -> Array: 187 | """ 188 | Computes the metric for the intrinsic binary parameters of the binary, 189 | maximized over the time and phase at coalescence. See 190 | :meth:`diffbank.metric.get_density`. 191 | """ 192 | return get_g(theta, self.amp, self.Psi, self.fs, self.Sn) 193 | 194 | def gam_fun(self, theta) -> Array: 195 | """ 196 | Computes the metric for the time of coalescence and intrinsic parameters 197 | of the binary. See :meth:`diffbank.metric.get_gam`. 198 | """ 199 | return get_gam(theta, self.amp, self.Psi, self.fs, self.Sn) 200 | 201 | def est_ratio_max( 202 | self, 203 | key: PRNGKeyArray, 204 | n_iter: int = 1000, 205 | n_init: int = 200, 206 | show_progress: bool = True, 207 | ) -> Tuple[Array, Array]: 208 | r""" 209 | Estimate maximum of the ratio of target to base density using empirical 210 | supremum rejection sampling. See :meth:`diffbank.utils.est_ratio_max`. 211 | """ 212 | return est_ratio_max( 213 | key, 214 | self.density_fun, 215 | self.sample_base, 216 | self.density_fun_base, 217 | n_iter, 218 | n_init, 219 | show_progress, 220 | ) 221 | 222 | def gen_template_rejection(self, key: PRNGKeyArray) -> Array: 223 | r""" 224 | Generates a single template using rejection sampling. See 225 | :meth:`diffbank.utils.gen_template_rejection`. 226 | 227 | Raises: 228 | RuntimeError: if ``ratio_max`` is not set 229 | """ 230 | if self.ratio_max is None: 231 | raise RuntimeError( 232 | "must set bank's 'ratio_max' attribute to an estimate of the" 233 | " maximum value of sqrt(|g|)" 234 | ) 235 | return gen_template_rejection( 236 | key, 237 | self.ratio_max, 238 | self.density_fun, 239 | self.sample_base, 240 | self.density_fun_base, 241 | ) 242 | 243 | def gen_templates_rejection(self, key: PRNGKeyArray, n_templates) -> Array: 244 | r""" 245 | Generates multiple templates using rejection sampling. See 246 | :meth:`diffbank.utils.gen_template_rejection`. 247 | 248 | Raises: 249 | RuntimeError: if ``ratio_max`` is not set 250 | """ 251 | if self.ratio_max is None: 252 | raise RuntimeError( 253 | "must set bank's 'ratio_max' attribute to an estimate of the" 254 | " maximum value of sqrt(|g|)" 255 | ) 256 | return gen_templates_rejection( 257 | key, n_templates, self.ratio_max, self.density_fun, self.sample_base 258 | ) 259 | 260 | def fill_bank( 261 | self, 262 | key: PRNGKeyArray, 263 | method="random", 264 | n_eff: int = 1000, 265 | effs=None, 266 | show_progress: bool = True, 267 | save_interval: Optional[int] = 20, 268 | save_path: str = "", 269 | ): 270 | """ 271 | Fills the bank with the required number of templates using the random or 272 | stochastic sampling methods from Coogan et al 2022. Updates the 273 | ``templates`` and ``n_templates`` attributes. See 274 | :meth:`diffbank.utils.gen_bank_random` and 275 | :meth:`diffbank.utils.gen_bank_stochastic`. 276 | 277 | Raises: 278 | RuntimeError: if ``ratio_max`` is not set 279 | """ 280 | if self.ratio_max is None: 281 | raise RuntimeError( 282 | "must set bank's 'ratio_max' attribute to an estimate of the" 283 | " maximum value of sqrt(|g|)" 284 | ) 285 | 286 | save_callback = lambda t, ep, e, k: jnp.savez( 287 | os.path.join(save_path, f"{self.name}-checkpoint.npz"), 288 | templates=t, 289 | eff_pts=ep, 290 | eff=e, 291 | key=k, 292 | ) 293 | 294 | if method == "random": 295 | self.templates, _ = gen_bank_random( 296 | key, 297 | 1 - self.m_star, 298 | self.eta, 299 | self.match_fun, 300 | self.ratio_max, 301 | self.density_fun, 302 | self.sample_base, 303 | self.density_fun_base, 304 | n_eff=n_eff, 305 | show_progress=show_progress, 306 | callback_interval=save_interval, 307 | callback_fn=save_callback, 308 | templates=self.templates, 309 | eff_pts=self.effectualness_points, 310 | effs=effs, 311 | ) 312 | self.n_templates = len(self.templates) 313 | elif method == "stochastic": 314 | propose_template = jax.jit(self.gen_template_rejection) 315 | self.templates, _ = gen_bank_stochastic( 316 | key, 317 | 1 - self.m_star, 318 | self.eta, 319 | self.match_fun, 320 | propose_template, 321 | propose_template, 322 | n_eff=n_eff, 323 | show_progress=show_progress, 324 | callback_interval=save_interval, 325 | callback_fn=save_callback, 326 | ) 327 | self.n_templates = len(self.templates) 328 | 329 | def get_top_matches_templates( 330 | self, 331 | theta2s: Array, 332 | amp2: Optional[Callable[[Array, Array], Array]] = None, 333 | Psi2: Optional[Callable[[Array, Array], Array]] = None, 334 | n_top: int = 1, 335 | show_progress: bool = True, 336 | ) -> Tuple[Array, Array]: 337 | """ 338 | Finds best-matching templates for a set of points for a waveform model distinct 339 | from the bank's. See :meth:`diffbank.utils.get_top_matches_templates`. 340 | 341 | Raises: 342 | RuntimeError: if the bank does not contain templates 343 | """ 344 | if self.templates is None: 345 | raise RuntimeError("cannot calculate effectualness of an empty bank") 346 | 347 | match_fun = lambda template, point: self.match_fun(template, point, amp2, Psi2) 348 | return get_top_matches_templates( 349 | self.templates, match_fun, theta2s, n_top, show_progress 350 | ) 351 | 352 | def calc_bank_effectualness( 353 | self, key: PRNGKeyArray, n: int, show_progress: bool = True 354 | ): 355 | """ 356 | Computes effectualness of bank at points sampled from the metric 357 | density. See :meth:`diffbank.utils.get_bank_effectualness`. 358 | 359 | Raises: 360 | RuntimeError: if the bank does not contain templates 361 | """ 362 | if self.templates is None: 363 | raise RuntimeError("cannot calculate effectualness of an empty bank") 364 | 365 | sample_eff_pt = jax.jit(self.gen_template_rejection) 366 | ( 367 | self.effectualnesses, 368 | self.effectualness_points, 369 | self.eta_est, 370 | self.eta_est_err, 371 | ) = get_bank_effectualness( 372 | key, 373 | 1 - self.m_star, 374 | self.templates, 375 | self.match_fun, 376 | sample_eff_pt, 377 | n, 378 | show_progress, 379 | ) 380 | 381 | def save(self, path: str = ""): 382 | """ 383 | Saves template bank non-function attributes to a ``npz`` file. The 384 | file name is the bank's name plus the ``.npz`` extension. 385 | 386 | Args: 387 | path: directory in which to save bank 388 | """ 389 | d = {k: getattr(self, k) for k in self.provided_vars | self.computed_vars} 390 | jnp.savez(os.path.join(path, f"{self.name}.npz"), bank=d) 391 | 392 | @classmethod 393 | def load( 394 | cls, 395 | path: str, 396 | amp: Callable[[Array, Array], Array], 397 | Psi: Callable[[Array, Array], Array], 398 | Sn: Callable[[Array], Array], 399 | sample_base: Callable[[PRNGKeyArray, int], Array], 400 | ignore_key_errors: bool = False, 401 | ): 402 | """ 403 | Loads template bank's non-function attributes from a ``npz`` file. The 404 | other required attributes must be provided as arguments to this 405 | function. 406 | 407 | Args: 408 | path: path to ``npz`` file generated with ``Bank.save`` 409 | amp: function returning the amplitude of the Fourier-domain waveform 410 | as a function of frequency and parameters 411 | Psi: function returning the phase of the Fourier-domain waveform as 412 | a function of frequency and parameters 413 | Sn: power spectral density of the detector noise 414 | sample_base: ``jit``-able base sampler for rejection sampling from 415 | the metric density. Takes a key and number of samples as 416 | arguments. 417 | ignore_key_errors: if ``True``, will try to load the bank even if 418 | there are extra or missing keys in the ``npz`` file. 419 | 420 | Raises: 421 | ValueError: if the bank being loaded is missing or contains extra keys, 422 | and ``ignore_key_errors`` is ``False``. This should only happen 423 | when trying to load old banks after ``diffbank`` goes through major 424 | updates. 425 | 426 | Returns: 427 | A new ``Bank`` 428 | """ 429 | d = jnp.load(path, allow_pickle=True)["bank"].item() 430 | if d.keys() != cls.provided_vars | cls.computed_vars and not ignore_key_errors: 431 | raise ValueError("missing or extra keys in bank file") 432 | 433 | # Instantiate with provided variables and functions 434 | fn_kwargs = { 435 | "amp": amp, 436 | "Psi": Psi, 437 | "Sn": Sn, 438 | "sample_base": sample_base, 439 | } 440 | 441 | try: 442 | bank = cls(**{**fn_kwargs, **{name: d[name] for name in cls.provided_vars}}) 443 | 444 | # Set computed variables 445 | for name in cls.computed_vars: 446 | setattr(bank, name, d[name]) 447 | except KeyError as e: 448 | if not ignore_key_errors: 449 | raise e 450 | 451 | return bank # type: ignore 452 | -------------------------------------------------------------------------------- /src/diffbank/constants.py: -------------------------------------------------------------------------------- 1 | """ 2 | Various constants. ``diffbank`` uses SI units throughout. 3 | """ 4 | 5 | MSUN = 1.98855e30 # kg 6 | """Solar mass""" 7 | 8 | G = 6.674e-11 # m^3 / kg / s^2 9 | """Newton's gravitational constant""" 10 | 11 | C = 299792458.0 # m / s 12 | """Speed of light""" 13 | -------------------------------------------------------------------------------- /src/diffbank/metric.py: -------------------------------------------------------------------------------- 1 | """ 2 | Functions related to the parameter space metric. 3 | """ 4 | # from functools import partial 5 | from typing import Callable 6 | 7 | import jax.numpy as jnp 8 | from jax import hessian # , jit 9 | 10 | from .utils import Array, get_phase_maximized_inner_product 11 | 12 | 13 | # @partial(jit, static_argnames=("amp", "Psi", "Sn")) 14 | def get_gam( 15 | theta: Array, 16 | amp: Callable[[Array, Array], Array], 17 | Psi: Callable[[Array, Array], Array], 18 | fs: Array, 19 | Sn: Callable[[Array], Array], 20 | ) -> Array: 21 | """ 22 | Computes the metric for the time of coalescence and intrinsic parameters of 23 | the binary. 24 | 25 | Args: 26 | theta: point at which to compute metric 27 | amp: function returning the amplitude of the Fourier-domain waveform as 28 | a function of frequency and parameters 29 | Psi: function returning the phase of the Fourier-domain waveform as a 30 | function of frequency and parameters 31 | fs: uniformly-spaced grid of frequencies used to compute the 32 | noise-weighted inner product 33 | Sn: power spectral density of the detector noise 34 | 35 | Return: 36 | The metric. The first value of the index corresponds to the time of 37 | coalescence. The other values of the index correspond to the intrinsic 38 | parameters, in the same order they appear in ``theta``. 39 | """ 40 | hess_func = lambda delta: get_phase_maximized_inner_product( 41 | theta, theta + delta[1:], delta[0], amp, Psi, amp, Psi, fs, Sn 42 | ) 43 | del_theta = jnp.zeros(theta.size + 1) 44 | return -1 / 2 * hessian(hess_func)(del_theta) 45 | 46 | 47 | # @partial(jit, static_argnames=("amp", "Psi", "Sn")) 48 | def get_g( 49 | theta: Array, 50 | amp: Callable[[Array, Array], Array], 51 | Psi: Callable[[Array, Array], Array], 52 | fs: Array, 53 | Sn: Callable[[Array], Array], 54 | ) -> Array: 55 | """ 56 | Computes the metric for the intrinsic binary parameters of the binary, 57 | maximized over the time and phase at coalescence. 58 | 59 | Args: 60 | theta: point at which to compute metric 61 | amp: function returning the amplitude of the Fourier-domain waveform as 62 | a function of frequency and parameters 63 | Psi: function returning the phase of the Fourier-domain waveform as a 64 | function of frequency and parameters 65 | fs: uniformly-spaced grid of frequencies used to compute the 66 | noise-weighted inner product 67 | Sn: power spectral density of the detector noise 68 | 69 | Return: 70 | The metric, where the values of the index correspond to the intrinsic 71 | parameters in the same order they appear in ``theta``. 72 | """ 73 | gam = get_gam(theta, amp, Psi, fs, Sn) 74 | # Maximize over Delta t_0 75 | return gam[1:, 1:] - jnp.outer(gam[0, 1:], gam[0, 1:]) / gam[0, 0] 76 | 77 | 78 | # @partial(jit, static_argnames=("amp", "Psi", "Sn")) 79 | def get_density( 80 | theta: Array, 81 | amp: Callable[[Array, Array], Array], 82 | Psi: Callable[[Array, Array], Array], 83 | fs: Array, 84 | Sn: Callable[[Array], Array], 85 | ) -> Array: 86 | r""" 87 | Computes the determinant of the metric over the intrinsic of the binary, 88 | maximized over the time and phase at coalescence. 89 | 90 | Args: 91 | theta: point at which to compute metric density 92 | amp: function returning the amplitude of the Fourier-domain waveform as 93 | a function of frequency and parameters 94 | Psi: function returning the phase of the Fourier-domain waveform as a 95 | function of frequency and parameters 96 | fs: uniformly-spaced grid of frequencies used to compute the 97 | noise-weighted inner product 98 | Sn: power spectral density of the detector noise 99 | 100 | Return: 101 | The determinant of the metric over the intrinsic parameters of the 102 | binary. 103 | """ 104 | return jnp.sqrt(jnp.linalg.det(get_g(theta, amp, Psi, fs, Sn))) 105 | 106 | 107 | # @partial(jit, static_argnames=("amp", "Psi", "Sn")) 108 | def get_metric_ellipse( 109 | theta: Array, 110 | amp: Callable[[Array, Array], Array], 111 | Psi: Callable[[Array, Array], Array], 112 | fs: Array, 113 | Sn: Callable[[Array], Array], 114 | ) -> Array: 115 | r""" 116 | Gets the ellipse representation of the metric. 117 | 118 | Warning: 119 | Only works in two dimensions. 120 | 121 | Args: 122 | theta: point at which to find metric ellipse 123 | amp: function returning the amplitude of the Fourier-domain waveform as 124 | a function of frequency and parameters 125 | Psi: function returning the phase of the Fourier-domain waveform as a 126 | function of frequency and parameters 127 | fs: uniformly-spaced grid of frequencies used to compute the 128 | noise-weighted inner product 129 | Sn: power spectral density of the detector noise 130 | 131 | Returns: 132 | Major and minor axes and orientation of the metric ellipse. 133 | """ 134 | g = get_g(theta, amp, Psi, fs, Sn) 135 | eigval, norm_eigvec = jnp.linalg.eig(g) 136 | r_major, r_minor = 1 / jnp.sqrt(eigval) 137 | U = jnp.linalg.inv(norm_eigvec) 138 | ang = jnp.arccos(U[0, 0] / jnp.linalg.norm(U[:, 0])) 139 | 140 | return jnp.array([r_major, r_minor, ang]) 141 | -------------------------------------------------------------------------------- /src/diffbank/noise.py: -------------------------------------------------------------------------------- 1 | """ 2 | Detector noise power spectral densities. 3 | """ 4 | try: 5 | import importlib.resources as pkg_resources 6 | except ImportError: 7 | # Try backported to PY<37 `importlib_resources`. 8 | import importlib_resources as pkg_resources 9 | 10 | from typing import Callable 11 | 12 | import jax 13 | import jax.numpy as jnp 14 | import numpy as np 15 | 16 | from . import noise_resources 17 | from .utils import Array 18 | 19 | 20 | def Sn_LIGOI(f: Array) -> Array: 21 | r""" 22 | LIGO-I power spectral density. 23 | 24 | References: 25 | ``_ 26 | 27 | Args: 28 | f: frequency 29 | 30 | Returns: 31 | The noise power spectral density 32 | """ 33 | fs = 40 # Hz 34 | f_theta = 150 # Hz 35 | x = f / f_theta 36 | normalization = 1e-46 37 | return jnp.where( 38 | f > fs, 39 | normalization 40 | * 9 41 | * ((4.49 * x) ** (-56) + 0.16 * x ** (-4.52) + 0.52 + 0.32 * x ** 2), 42 | jnp.inf, 43 | ) 44 | 45 | 46 | def _load_noise(name: str, asd: bool = False) -> Callable[[Array], Array]: 47 | r""" 48 | Loads noise curve from text data file into an interpolator. The file's 49 | columns must contain the frequencies and corresponding noise power spectral 50 | density or amplitude spectral density values. 51 | 52 | Args: 53 | name: name of data file in ``noise_resources`` without the `.dat` 54 | extension 55 | asd: ``True`` if the file contains ASD rather than PSD data 56 | 57 | Returns 58 | Interpolator for noise curve returning ``inf`` above and below the 59 | frequency range in the data file 60 | """ 61 | path_context = pkg_resources.path(noise_resources, f"{name}.dat") 62 | with path_context as path: 63 | fs, Sns = np.loadtxt(path, unpack=True) 64 | 65 | if asd: 66 | Sns = Sns ** 2 67 | 68 | Sns[Sns == 0.0] = np.inf 69 | 70 | fs = jnp.array(fs) 71 | Sns = jnp.array(Sns) 72 | 73 | return jax.jit(lambda f: jnp.interp(f, fs, Sns, left=jnp.inf, right=jnp.inf)) 74 | 75 | 76 | Sn_aLIGOZeroDetHighPower = _load_noise("aLIGOZeroDetHighPower") 77 | r"""The aLIGOZeroDetHighPower noise curve from pycbc. 78 | 79 | References: 80 | ??? 81 | 82 | Args: 83 | f: frequency 84 | 85 | Returns: 86 | The noise power spectral density 87 | """ 88 | 89 | Sn_O3a = _load_noise("O3a_Livingston_ASD", asd=True) 90 | r"""The LIGO O3a Livingston noise curve. 91 | 92 | References: 93 | ??? 94 | 95 | Args: 96 | f: frequency 97 | 98 | Returns: 99 | The noise power spectral density 100 | """ 101 | 102 | Sn_O2 = _load_noise("O2_ASD", asd=True) 103 | r"""The LIGO O2 noise curve. 104 | 105 | References: 106 | ``_ 107 | 108 | Args: 109 | f: frequency 110 | 111 | Returns: 112 | The noise power spectral density 113 | """ 114 | -------------------------------------------------------------------------------- /src/diffbank/noise_resources/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/src/diffbank/noise_resources/__init__.py -------------------------------------------------------------------------------- /src/diffbank/noise_resources/aLIGOZeroDetHighPower.dat: -------------------------------------------------------------------------------- 1 | # aLIGOZeroDetHighPower noise curve from pycbc 2 | # Columns: f [Hz], PSD 3 | 1.999999999999999289e+01 4.295614257695998751e-46 4 | 2.047803355839261741e+01 3.883262108739285580e-46 5 | 2.096749292093271677e+01 3.513718228164062869e-46 6 | 2.146865118351100676e+01 3.182498850154123068e-46 7 | 2.198178796946818636e+01 2.880126867743783032e-46 8 | 2.250718958561203920e+01 2.614543966611146638e-46 9 | 2.304514918196342776e+01 2.372114560063095778e-46 10 | 2.359596691532055601e+01 2.155170502585085538e-46 11 | 2.415995011673283699e+01 1.960936421848462479e-46 12 | 2.473741346297734012e+01 1.784082755163435794e-46 13 | 2.532867915213417476e+01 1.625956980185866039e-46 14 | 2.593407708335816153e+01 1.482266568676072662e-46 15 | 2.655394504094748953e+01 1.353904816091892958e-46 16 | 2.718862888281180545e+01 1.239137492461325411e-46 17 | 2.783848273344515789e+01 1.134930128538470590e-46 18 | 2.850386918151116333e+01 1.040470502965065315e-46 19 | 2.918515948215095790e+01 9.561357470835292138e-47 20 | 2.988273376412640303e+01 8.797322030924402362e-47 21 | 3.059698124191464075e+01 8.114996412781482270e-47 22 | 3.132830043287189881e+01 7.496948954201411366e-47 23 | 3.207709937958782831e+01 6.937678170846969859e-47 24 | 3.284379587755474006e+01 6.432012741433469664e-47 25 | 3.362881770827817718e+01 5.980551488576292971e-47 26 | 3.443260287795941821e+01 5.567299806810546355e-47 27 | 3.525559986188298467e+01 5.198482893458040958e-47 28 | 3.609826785464509413e+01 4.865054514491053515e-47 29 | 3.696107702636341230e+01 4.560392545284464964e-47 30 | 3.784450878501023396e+01 4.288399269355341107e-47 31 | 3.874905604501618228e+01 4.042555501459501061e-47 32 | 3.967522350229393879e+01 3.818041551264217365e-47 33 | 4.062352791583514744e+01 3.615329151236163749e-47 34 | 4.159449839603856702e+01 3.432215181051643205e-47 35 | 4.258867669992929450e+01 3.266709501161945412e-47 36 | 4.360661753343433134e+01 3.115622573325289432e-47 37 | 4.464888886088301945e+01 2.979045946041466017e-47 38 | 4.571607222190525732e+01 2.854370438427165394e-47 39 | 4.680876305590383168e+01 2.741583097426302991e-47 40 | 4.792757103428237997e+01 2.638559315494744999e-47 41 | 4.907312040061404446e+01 2.544461618046305996e-47 42 | 5.024605031894083140e+01 2.459193619247115917e-47 43 | 5.144701523039773861e+01 2.381189363043735399e-47 44 | 5.267668521836103679e+01 2.309798942174515183e-47 45 | 5.393574638232409768e+01 2.244421134205164751e-47 46 | 5.522490122070935570e+01 2.184502077255961518e-47 47 | 5.654486902283019845e+01 2.129533428089485501e-47 48 | 5.789638627022161188e+01 2.079050154261648106e-47 49 | 5.928020704756301029e+01 2.032309041853823840e-47 50 | 6.069710346342291274e+01 1.989593506941756919e-47 51 | 6.214786608106017951e+01 1.949939442594162849e-47 52 | 6.363330435952208575e+01 1.913103891336993162e-47 53 | 6.515424710528525054e+01 1.879070934999045561e-47 54 | 6.671154293469183472e+01 1.847373767155782349e-47 55 | 6.830606074743845113e+01 1.817995380634058553e-47 56 | 6.993869021138253572e+01 1.790549028991190350e-47 57 | 7.161034225893585869e+01 1.764726897969965730e-47 58 | 7.332194959532350254e+01 1.740686844714436162e-47 59 | 7.507446721899034969e+01 1.718131897011704024e-47 60 | 7.686887295444657298e+01 1.696937326798950626e-47 61 | 7.870616799784880868e+01 1.676884627580561435e-47 62 | 8.058737747562180687e+01 1.657990769821645743e-47 63 | 8.251355101643186174e+01 1.640069826203928320e-47 64 | 8.448576333683168116e+01 1.623136730172874432e-47 65 | 8.650511484090282011e+01 1.607033731662399678e-47 66 | 8.857273223423078434e+01 1.591704022183909060e-47 67 | 9.068976915255510107e+01 1.577161635106469385e-47 68 | 9.285740680544518000e+01 1.563220781157180861e-47 69 | 9.507685463536115833e+01 1.549908571319502947e-47 70 | 9.734935099246710877e+01 1.537184504409409697e-47 71 | 9.967616382557417865e+01 1.524960884767620215e-47 72 | 1.020585913895975096e+02 1.513261232635641342e-47 73 | 1.044979629699228525e+02 1.502009638385919966e-47 74 | 1.069956396240875591e+02 1.491185938410482850e-47 75 | 1.095530149411873708e+02 1.480771489777426882e-47 76 | 1.121715158194362374e+02 1.470749127653420709e-47 77 | 1.148526032623091595e+02 1.461103104269037922e-47 78 | 1.175977731937161082e+02 1.451819015241578186e-47 79 | 1.204085572926582159e+02 1.442852749560879994e-47 80 | 1.232865238478347294e+02 1.434197725063887395e-47 81 | 1.262332786326766723e+02 1.425847812276182391e-47 82 | 1.292504658012938989e+02 1.417797339951893022e-47 83 | 1.323397688058388155e+02 1.410065446377679317e-47 84 | 1.355029113357944937e+02 1.402574281802054387e-47 85 | 1.387416582797149260e+02 1.395370995926732523e-47 86 | 1.420578167099522773e+02 1.388451415643993878e-47 87 | 1.454532368909194986e+02 1.381773697581351845e-47 88 | 1.489298133114541542e+02 1.375377193549870756e-47 89 | 1.524894857418552476e+02 1.369225051639712864e-47 90 | 1.561342403161873449e+02 1.363320679708802508e-47 91 | 1.598661106404510690e+02 1.357681863232999836e-47 92 | 1.636871789272433375e+02 1.352294993963371125e-47 93 | 1.675995771575354354e+02 1.347163011975847783e-47 94 | 1.716054882702211444e+02 1.342276981126124541e-47 95 | 1.757071473800971262e+02 1.337653074780006550e-47 96 | 1.799068430249532753e+02 1.333283780559908832e-47 97 | 1.842069184424732669e+02 1.329173973193900118e-47 98 | 1.886097728776531426e+02 1.325328296656800898e-47 99 | 1.931178629214696798e+02 1.321743631217759769e-47 100 | 1.977337038815460346e+02 1.318426512454222172e-47 101 | 2.024598711855786064e+02 1.315389262614555628e-47 102 | 2.072990018183062091e+02 1.312625029055572594e-47 103 | 2.122537957928285266e+02 1.310141739664987726e-47 104 | 2.173270176570877652e+02 1.307951006808051662e-47 105 | 2.225214980363616064e+02 1.306051563991806524e-47 106 | 2.278401352126206802e+02 1.304455208505190366e-47 107 | 2.332858967416378277e+02 1.303168879794697938e-47 108 | 2.388618211087489271e+02 1.302196476702934481e-47 109 | 2.445710194241866873e+02 1.301549082064601399e-47 110 | 2.504166771589395921e+02 1.301234913735752045e-47 111 | 2.564020559220969631e+02 1.301262704777242715e-47 112 | 2.625304952806780534e+02 1.301642904621162975e-47 113 | 2.688054146229582102e+02 1.302384973950636150e-47 114 | 2.752303150663289557e+02 1.303502272508718913e-47 115 | 2.818087814107631175e+02 1.305002189063655940e-47 116 | 2.885444841389667658e+02 1.306903111752953732e-47 117 | 2.954411814643426624e+02 1.309215261771127716e-47 118 | 3.025027214278985639e+02 1.311950779602568480e-47 119 | 3.097330440452802236e+02 1.315126394629239216e-47 120 | 3.171361835051176286e+02 1.318755686583544664e-47 121 | 3.247162704199178620e+02 1.322857866815496677e-47 122 | 3.324775341307584426e+02 1.327447985276665075e-47 123 | 3.404243050670651769e+02 1.332547758675028700e-47 124 | 3.485610171627926093e+02 1.338173760522960194e-47 125 | 3.568922103303565905e+02 1.344342610936038526e-47 126 | 3.654225329936982121e+02 1.351079148304737763e-47 127 | 3.741567446818892222e+02 1.358401641559892918e-47 128 | 3.830997186847336025e+02 1.366337602653216175e-47 129 | 3.922564447718372094e+02 1.374916646860576798e-47 130 | 4.016320319766734883e+02 1.384160420524311990e-47 131 | 4.112317114471871946e+02 1.394090772516928144e-47 132 | 4.210608393645363208e+02 1.404740860609382529e-47 133 | 4.311248999315972696e+02 1.416134495567597898e-47 134 | 4.414295084328953180e+02 1.428307842333892054e-47 135 | 4.519804143676797707e+02 1.441311927160016671e-47 136 | 4.627835046578772449e+02 1.455149877059953634e-47 137 | 4.738448069327282042e+02 1.469875466674718859e-47 138 | 4.851704928919243685e+02 1.485532848684031857e-47 139 | 4.967668817491356208e+02 1.502139514937478542e-47 140 | 5.086404437578432294e+02 1.519771801622899860e-47 141 | 5.207978038214412209e+02 1.538434697833357796e-47 142 | 5.332457451896328848e+02 1.558210331822043703e-47 143 | 5.459912132431688860e+02 1.579137753958695972e-47 144 | 5.590413193690560547e+02 1.601256712287663552e-47 145 | 5.724033449283806476e+02 1.624625438626640335e-47 146 | 5.860847453189786620e+02 1.649305114763409575e-47 147 | 6.000931541352016438e+02 1.675359963827378849e-47 148 | 6.144363874271172108e+02 1.702837901023425593e-47 149 | 6.291224480615022685e+02 1.731827817429941545e-47 150 | 6.441595301870778485e+02 1.762362413494983074e-47 151 | 6.595560238064706482e+02 1.794557895280629090e-47 152 | 6.753205194574451298e+02 1.828473334073254540e-47 153 | 6.914618130060354133e+02 1.864169024637050364e-47 154 | 7.079889105542296193e+02 1.901729552614225047e-47 155 | 7.249110334649677725e+02 1.941290676360087956e-47 156 | 7.422376235072340478e+02 1.982875596965642422e-47 157 | 7.599783481241368008e+02 2.026628760062163837e-47 158 | 7.781431058268937022e+02 2.072652773579457067e-47 159 | 7.967420317177492279e+02 2.121028295991562481e-47 160 | 8.157855031448998488e+02 2.171864787506997775e-47 161 | 8.352841454925730886e+02 2.225304081670298579e-47 162 | 8.552488381095113255e+02 2.281438704174093456e-47 163 | 8.756907203791432721e+02 2.340450900014992684e-47 164 | 8.966211979348558998e+02 2.402414100271541986e-47 165 | 9.180519490238082199e+02 2.467493244975090920e-47 166 | 9.399949310228654440e+02 2.535830279892053312e-47 167 | 9.624623871102603516e+02 2.607605158686158696e-47 168 | 9.854668530967287552e+02 2.682940626832571488e-47 169 | 1.009021164419919842e+03 2.762064208247527015e-47 170 | 1.033138463305975392e+03 2.845077041658915360e-47 171 | 1.057832206102298187e+03 2.932258544534424184e-47 172 | 1.083116170785567647e+03 3.023722301919536846e-47 173 | 1.109004464649229021e+03 3.119731875991590385e-47 174 | 1.135511532174708691e+03 3.220488439436866813e-47 175 | 1.162652163090774820e+03 3.326200705681810297e-47 176 | 1.190441500625533763e+03 3.437125547396485822e-47 177 | 1.218895049955646982e+03 3.553490400170400144e-47 178 | 1.248028686857520370e+03 3.675573448453915646e-47 179 | 1.277858666565248313e+03 3.803665604434692079e-47 180 | 1.308401632840301090e+03 3.938070974817924013e-47 181 | 1.339674627257968723e+03 4.079016104001712699e-47 182 | 1.371695098715791346e+03 4.226872985527600248e-47 183 | 1.404480913169232281e+03 4.381984673914090683e-47 184 | 1.438050363600073524e+03 4.544660221414686047e-47 185 | 1.472422180223051782e+03 4.715269711015587348e-47 186 | 1.507615540936463503e+03 4.894199176110498279e-47 187 | 1.543650082022558308e+03 5.081903859864711013e-47 188 | 1.580545909103673011e+03 5.278699232040103498e-47 189 | 1.618323608360260550e+03 5.485129299907780764e-47 190 | 1.657004258017022039e+03 5.701596946886741810e-47 191 | 1.696609440103603220e+03 5.928576619231173705e-47 192 | 1.737161252496364341e+03 6.166681051683314898e-47 193 | 1.778682321247996470e+03 6.416311211934339037e-47 194 | 1.821195813211809082e+03 6.678064412816786238e-47 195 | 1.864725448967777766e+03 6.952565412690708204e-47 196 | 1.909295516057546138e+03 7.240402094905998063e-47 197 | 1.954930882535748651e+03 7.542252504037334565e-47 198 | 2.001657010845259492e+03 7.858688015342971260e-47 199 | 2.049499972024053477e+03 8.190508879183589559e-47 200 | 2.098486460251664539e+03 8.538482416495913391e-47 201 | 2.148643807743308116e+03 8.903335931910260582e-47 202 | 2.199999999999997272e+03 9.285824434920114223e-47 203 | -------------------------------------------------------------------------------- /src/diffbank/waveforms/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/adam-coogan/diffbank/0a38c0f1253af8ffcc0fd491c3f9ebea172569d8/src/diffbank/waveforms/__init__.py -------------------------------------------------------------------------------- /src/diffbank/waveforms/taylorF2.py: -------------------------------------------------------------------------------- 1 | from math import pi 2 | 3 | import jax.numpy as jnp 4 | 5 | # from ..constants import C, G, MSUN 6 | 7 | 8 | def Psi(f: jnp.ndarray, theta: jnp.ndarray) -> jnp.ndarray: 9 | """ 10 | Computes the phase of the waveform. Sets time and phase of coealence to be zero. 11 | Returns: 12 | -------- 13 | phase (array): Phase of the GW as a function of frequency 14 | """ 15 | 16 | # ( 17 | # Mt, 18 | # eta, 19 | # chi1, 20 | # chi2, 21 | # ) = theta 22 | 23 | m1, m2, chi1, chi2 = theta 24 | Mt = m1 + m2 25 | eta = m1 * m2 / (m1 + m2) ** 2 26 | 27 | kappa1 = 1.0 28 | kappa2 = 1.0 29 | 30 | # Mt, eta, chi_1, chi_2, kappa1, kappa2 = theta 31 | 32 | # th0, th3, chi_1, chi_2, kappa1, kappa2 = theta 33 | # M_chirp = ( 34 | # 1 / (16 * pi * f[0]) * (125 / (2 * th0 ** 3)) ** (1 / 5) * C ** 3 / G 35 | # ) / MSUN 36 | # eta = (16 * pi ** 5 / 25 * th0 ** 2 / th3 ** 5) ** (1 / 3) 37 | # Mt = M_chirp / eta ** (3 / 5) 38 | 39 | gt = 4.92549094830932e-6 # GN*Msun/c^3 in seconds 40 | EulerGamma = 0.57721566490153286060 41 | vlso = 1.0 / jnp.sqrt(6.0) 42 | 43 | chi_s = 0.5 * (chi1 + chi2) 44 | chi_a = 0.5 * (chi1 - chi2) 45 | k_s = 0.5 * (kappa1 + kappa2) 46 | k_a = 0.5 * (kappa1 - kappa2) 47 | lambda1 = 1.0 48 | lambda2 = 1.0 49 | 50 | lambda_s = 0.5 * (lambda1 + lambda2) 51 | lambda_a = 0.5 * (lambda1 - lambda2) 52 | L = 1.0 53 | delta = jnp.sqrt(1.0 - 4.0 * eta) 54 | 55 | v = (pi * Mt * (f + 1e-100) * gt) ** (1.0 / 3.0) 56 | v2 = v * v 57 | v3 = v2 * v 58 | v4 = v2 * v2 59 | v5 = v4 * v 60 | v6 = v3 * v3 61 | v7 = v3 * v4 62 | v10 = v5 * v5 63 | v12 = v10 * v2 64 | eta2 = eta ** 2 65 | eta3 = eta ** 3 66 | 67 | # # ------------------------- Non spinning part of the waveform 68 | psi_NS_0PN = 1.0 69 | psi_NS_1PN = (3715.0 / 756.0 + 55.0 * eta / 9.0) * v2 70 | psi_NS_15PN = -16.0 * pi * v3 71 | psi_NS_2PN = ( 72 | 15293365.0 / 508032.0 + 27145.0 * eta / 504.0 + 3085.0 * eta2 / 72.0 73 | ) * v4 74 | psi_NS_25PN = ( 75 | pi * (38645.0 / 756.0 - 65.0 * eta / 9.0) * (1 + 3.0 * jnp.log(v / vlso)) * v5 76 | ) 77 | psi_NS_3PN = ( 78 | ( 79 | 11583231236531.0 / 4694215680.0 80 | - 640.0 * pi ** 2 / 3.0 81 | - 6848.0 * EulerGamma / 21.0 82 | ) 83 | + (2255.0 * pi ** 2 / 12.0 - 15737765635.0 / 3048192.0) * eta 84 | + 76055.0 * eta2 / 1728.0 85 | - 127825.0 * eta3 / 1296.0 86 | - 6848.0 * jnp.log(4.0 * v) / 21.0 87 | ) * v6 88 | psi_NS_35PN = ( 89 | pi 90 | * (77096675.0 / 254016.0 + 378515.0 * eta / 1512.0 - 74045.0 * eta2 / 756.0) 91 | * v7 92 | ) 93 | 94 | # # ------------------------- Spining part of the waveform 95 | psi_S_15PN = ( 96 | (113.0 / 3.0 - 76.0 * eta / 3.0) * jnp.dot(chi_s, L) 97 | + 113.0 * delta * jnp.dot(chi_a, L) / 3.0 98 | ) * v3 99 | 100 | psi_S_2PN = ( 101 | -(5.0 / 8.0) 102 | * (1.0 + 156.0 * eta + 80.0 * delta * k_a + 80.0 * (1.0 - 2.0 * eta) * k_s) 103 | * (jnp.dot(chi_s, L) ** 2) 104 | ) 105 | psi_S_2PN -= ( 106 | (5.0 / 8.0) 107 | * (1.0 - 160.0 * eta + 80.0 * delta * k_a + 80.0 * (1.0 - 2.0 * eta) * k_s) 108 | * (jnp.dot(chi_a, L) ** 2) 109 | ) 110 | psi_S_2PN -= ( 111 | (5.0 / 4.0) 112 | * (delta + 80.0 * delta * k_s + 80.0 * (1.0 - 2.0 * eta) * k_a) 113 | * jnp.dot(chi_s, L) 114 | * jnp.dot(chi_a, L) 115 | ) 116 | psi_S_2PN *= v4 117 | 118 | psi_S_25PN = ( 119 | -(732985.0 / 2268.0 - 24260.0 * eta / 81.0 - 340.0 * eta2 / 9.0) 120 | * jnp.dot(chi_s, L) 121 | - (732985.0 / 2268.0 + 140.0 * eta / 9.0) * delta * jnp.dot(chi_a, L) 122 | ) * v5 123 | psi_S_25PN_log = 3.0 * psi_S_25PN * jnp.log(v / vlso) 124 | 125 | psi_S_3PN = (2270.0 / 3.0 - 520.0 * eta) * pi * jnp.dot(chi_s, L) + ( 126 | 2270.0 * pi / 3.0 127 | ) * delta * jnp.dot(chi_a, L) 128 | psi_S_3PN += ( 129 | ( 130 | (26015.0 / 14.0 - 88510.0 * eta / 21.0 - 480.0 * eta2) * k_a 131 | + delta 132 | * ( 133 | -1344475.0 / 1008.0 134 | + 745.0 * eta / 18.0 135 | + (26015.0 / 14.0 - 1495.0 * eta / 3.0) * k_s 136 | ) 137 | ) 138 | * jnp.dot(chi_s, L) 139 | * jnp.dot(chi_a, L) 140 | ) 141 | psi_S_3PN += ( 142 | -1344475.0 / 2016.0 143 | + 829705.0 * eta / 504.0 144 | + 3415.0 * eta2 / 9.0 145 | + (26015.0 / 28.0 - 44255.0 * eta / 21.0 - 240.0 * eta2) * k_s 146 | + delta * (26015.0 / 28.0 - 1495.0 * eta / 6.0) * k_a 147 | ) * (jnp.dot(chi_s, L)) ** 2 148 | psi_S_3PN += ( 149 | -1344475.0 / 2016.0 150 | + 267815.0 * eta / 252.0 151 | - 240.0 * eta2 152 | + (26015.0 / 28.0 - 44255.0 * eta / 21.0 - 240.0 * eta2) * k_s 153 | + delta * (26015.0 / 28.0 - 1495.0 * eta / 6.0) * k_a 154 | ) * (jnp.dot(chi_a, L)) ** 2 155 | psi_S_3PN *= v6 156 | 157 | psi_S_35PN = ( 158 | -25150083775.0 / 3048192.0 159 | + 10566655595.0 * eta / 762048.0 160 | - 1042165 * eta2 / 3024.0 161 | + 5345.0 * eta3 / 36.0 162 | ) * jnp.dot(chi_s, L) 163 | psi_S_35PN += ( 164 | (-25150083775.0 / 3048192.0 + 26804935.0 * eta / 6048.0 - 1985.0 * eta2 / 48.0) 165 | * delta 166 | * jnp.dot(chi_a, L) 167 | ) 168 | psi_S_35PN += ( 169 | 265.0 / 24.0 170 | + 4035.0 * eta / 2.0 171 | - 20.0 * eta2 / 3.0 172 | + (3110.0 / 3.0 - 10250.0 * eta / 3.0 + 40.0 * eta2) * k_s 173 | - 440.0 * (1.0 - 3.0 * eta) * lambda_s 174 | + delta 175 | * ((3110.0 / 3.0 - 4030.0 * eta / 3.0) * k_a - 440.0 * (1.0 - eta) * lambda_a) 176 | ) * (jnp.dot(chi_s, L)) ** 3 177 | psi_S_35PN += ( 178 | (3110.0 / 3.0 - 8470.0 * eta / 3.0) * k_a 179 | - 440.0 * (1.0 - 3.0 * eta) * lambda_a 180 | + delta 181 | * ( 182 | 265.0 / 24.0 183 | - 2070.0 * eta 184 | + (3110.0 / 3.0 - 750.0 * eta) * k_s 185 | - 440.0 * (1 - eta) * lambda_s 186 | ) 187 | ) * (jnp.dot(chi_a, L)) ** 3 188 | psi_S_35PN += ( 189 | (3110.0 - 28970.0 * eta / 3.0 + 80.0 * eta2) * k_a 190 | - 1320.0 * (1.0 - 3.0 * eta) * lambda_a 191 | + delta 192 | * ( 193 | 265.0 / 8.0 194 | + 12055.0 * eta / 6.0 195 | + (3110.0 - 10310.0 * eta / 3.0) * k_s 196 | - 1320.0 * (1.0 - eta) * lambda_s 197 | ) 198 | ) * (jnp.dot(chi_s, L) ** 2 * jnp.dot(chi_a, L)) 199 | psi_S_35PN += ( 200 | 265.0 / 8.0 201 | - 6500.0 * eta / 3.0 202 | + 40.0 * eta2 203 | + (3110.0 - 27190.0 * eta / 3.0 + 40.0 * eta2) * k_s 204 | - 1320.0 * (1.0 - 3 * eta) * lambda_s 205 | + delta 206 | * ((3110.0 - 8530.0 * eta / 3.0) * k_a - 1320.0 * (1.0 - eta) * lambda_a) 207 | ) * (jnp.dot(chi_a, L) ** 2 * jnp.dot(chi_s, L)) 208 | psi_S_35PN *= v7 209 | 210 | psi_NS = ( 211 | psi_NS_0PN 212 | + psi_NS_1PN 213 | + psi_NS_15PN 214 | + psi_NS_2PN 215 | + psi_NS_25PN 216 | + psi_NS_3PN 217 | + psi_NS_35PN 218 | ) 219 | psi_S = ( 220 | psi_S_15PN + psi_S_2PN + psi_S_25PN + psi_S_25PN_log + psi_S_3PN + psi_S_35PN 221 | ) 222 | 223 | return 3.0 / 128.0 / eta / v5 * (psi_NS + psi_S) 224 | 225 | 226 | def Amp(f: jnp.ndarray, theta: jnp.ndarray) -> jnp.ndarray: 227 | """ 228 | Computes the Taylor F2 Frequency domain strain waveform with non-standard 229 | spin induced quadrupoole moment for object two. 230 | 231 | Note that this waveform assumes object 1 is a BH and therefore uses the 232 | chi * M_total relation to find C 233 | 234 | Note that this waveform also assumes that object one is the more massive. 235 | Therefore the more massive object is always considered a BH 236 | 237 | Returns: 238 | Strain (array): 239 | """ 240 | # ( 241 | # th0, 242 | # th3, 243 | # _, 244 | # _, 245 | # _, 246 | # _, 247 | # ) = theta 248 | # M_chirp = ( 249 | # 1 / (16 * pi * f[0]) * (125 / (2 * th0 ** 3)) ** (1 / 5) * C ** 3 / G 250 | # ) / MSUN 251 | # eta = (16 * pi ** 5 / 25 * th0 ** 2 / th3 ** 5) ** (1 / 3) 252 | # Mt = M_chirp / eta ** (3 / 5) 253 | 254 | # ( 255 | # Mt, 256 | # eta, 257 | # _, 258 | # _, 259 | # ) = theta 260 | 261 | m1, m2, _, _ = theta 262 | 263 | Mt = m1 + m2 264 | eta = m1 * m2 / (m1 + m2) ** 2 265 | distance = 1.0 266 | pre = 3.6686934875530996e-19 # (GN*Msun/c^3)^(5/6)/Hz^(7/6)*c/Mpc/sec 267 | Mchirp = Mt * eta ** 0.6 268 | A0 = ( 269 | Mchirp ** (5.0 / 6.0) 270 | / (f + 1e-100) ** (7.0 / 6.0) 271 | / distance 272 | / pi ** (2.0 / 3.0) 273 | * jnp.sqrt(5.0 / 24.0) 274 | ) 275 | 276 | return pre * A0 277 | -------------------------------------------------------------------------------- /src/diffbank/waveforms/taylorf2reducedspin.py: -------------------------------------------------------------------------------- 1 | from math import pi 2 | 3 | import jax 4 | import jax.numpy as jnp 5 | 6 | from ..constants import C, G 7 | from ..utils import get_f_isco, ms_to_Mc_eta 8 | 9 | """ 10 | TaylorF2ReducedSpin waveform, in terms of the dimensionless chirp times (th0, 11 | th3, th3S). 12 | 13 | Reference: Ajith et al 2014, https://arxiv.org/abs/1210.6666 14 | """ 15 | 16 | 17 | def th_to_phys(th: jnp.ndarray, f_0=20.0) -> jnp.ndarray: 18 | """ 19 | Converts (th0, th3, th3S) -> (M_chirp, eta, chi). 20 | """ 21 | th0, th3, th3S = th 22 | M_chirp = 1 / (16 * pi * f_0) * (125 / (2 * th0 ** 3)) ** (1 / 5) * C ** 3 / G 23 | eta = (16 * pi ** 5 / 25 * th0 ** 2 / th3 ** 5) ** (1 / 3) 24 | chi = 48 * pi * th3S / (113 * th3) 25 | return jnp.stack([M_chirp, eta, chi]) 26 | 27 | 28 | def phys_to_th(phys: jnp.ndarray, f_0=20.0) -> jnp.ndarray: 29 | """ 30 | Converts (M_chirp, eta, chi) -> (th0, th3, th3S). 31 | """ 32 | _M_chirp, eta, chi = phys 33 | M_chirp = _M_chirp / (C ** 3 / G) 34 | th0 = 5 / (128 * (f_0 * M_chirp * pi) ** (5 / 3)) 35 | th3 = pi ** (1 / 3) / (4 * f_0 ** (2 / 3) * M_chirp ** (2 / 3) * eta ** (3 / 5)) 36 | th3S = ( 37 | 113 38 | * chi 39 | / (192 * f_0 ** (2 / 3) * M_chirp ** (2 / 3) * pi ** (2 / 3) * eta ** (3 / 5)) 40 | ) 41 | return jnp.stack([th0, th3, th3S]) 42 | 43 | 44 | def th_to_f_isco(theta, f_0=20.0): 45 | """ 46 | Computes f_isco from (th0, th3, th3S). 47 | """ 48 | M_chirp, eta, _ = th_to_phys(theta, f_0) 49 | M = M_chirp / eta ** (3 / 5) # total mass 50 | return get_f_isco(M) 51 | 52 | 53 | def get_th_boundary_interps(m_min, m_max, f_0=20.0, n=1000): 54 | """ 55 | Given a range of BH masses, returns corresponding range of th0 and 56 | interpolators for the minimum and maximum corresponding values of th3. 57 | """ 58 | chis = jnp.zeros((n,)) # value doesn't matter 59 | 60 | # Lower boundary 61 | ms = jnp.linspace(m_min, m_max, n) 62 | M_chirps = ms_to_Mc_eta(jnp.stack([ms, ms]))[0] 63 | etas = jnp.full_like(M_chirps, 0.25) 64 | th0_lows, th3_lows = phys_to_th(jnp.stack([M_chirps, etas, chis]), f_0)[:2, ::-1] 65 | interp_low = lambda th0: jnp.interp( 66 | th0, th0_lows, th3_lows, left=jnp.nan, right=jnp.nan 67 | ) 68 | 69 | # Upper boundary 70 | M_chirps, etas = ms_to_Mc_eta(jnp.stack([ms, jnp.full_like(ms, m_max)])) 71 | th0_highs_1, th3_highs_1 = phys_to_th(jnp.stack([M_chirps, etas, chis]), f_0)[ 72 | :2, ::-1 73 | ] 74 | 75 | M_chirps, etas = ms_to_Mc_eta(jnp.stack([ms, jnp.full_like(ms, m_min)])) 76 | th0_highs_2, th3_highs_2 = phys_to_th(jnp.stack([M_chirps, etas, chis]), f_0)[ 77 | :2, ::-1 78 | ] 79 | 80 | th0_highs = jnp.concatenate([th0_highs_1, th0_highs_2]) 81 | th3_highs = jnp.concatenate([th3_highs_1, th3_highs_2]) 82 | interp_high = lambda th0: jnp.interp( 83 | th0, th0_highs, th3_highs, left=jnp.nan, right=jnp.nan 84 | ) 85 | 86 | th0_min = min(th0_lows.min(), th0_highs_1.min()) 87 | th0_max = max(th0_lows.max(), th0_highs_2.max()) 88 | 89 | return (th0_min, th0_max), interp_low, interp_high 90 | 91 | 92 | @jax.jit 93 | def amp(f: jnp.ndarray, theta: jnp.ndarray, f_0=20.0) -> jnp.ndarray: 94 | """ 95 | Amplitude, truncated at the ISCO frequency for a BH with mass equal to the 96 | total mass of the system. 97 | """ 98 | return jnp.where( 99 | f <= th_to_f_isco(jax.lax.stop_gradient(theta), f_0), f ** (-7 / 6), 0.0 100 | ) 101 | 102 | 103 | @jax.jit 104 | def Psi(f: jnp.ndarray, theta: jnp.ndarray, f_0=20.0) -> jnp.ndarray: 105 | th0, th3, th3S = theta 106 | 107 | phi_0 = 0.0 108 | t_0 = 0.0 109 | 110 | psi_0 = 3 * th0 / 5 111 | psi_1 = 0.0 112 | psi_2 = 743 / 2016 * (25 / (2 * pi ** 2)) ** (1 / 3) * th0 ** (1 / 3) * th3 ** ( 113 | 2 / 3 114 | ) + 11 * pi * th0 / (12 * th3) 115 | psi_3 = -3 / 2 * (th3 - th3S) 116 | psi_4 = ( 117 | ( 118 | 675 119 | * th3 120 | * th3S ** 2 121 | * ( 122 | 8 * 10 ** (2 / 3) * pi ** (7 / 3) * th0 ** (2 / 3) 123 | - 405 * 10 ** (1 / 3) * pi ** (2 / 3) * th3 ** (5 / 3) 124 | ) 125 | / ( 126 | 4 127 | * th0 ** (1 / 3) 128 | * ( 129 | 152 * 10 ** (1 / 3) * pi ** (5 / 3) * th0 ** (2 / 3) 130 | - 565 * th3 ** (5 / 3) 131 | ) 132 | ** 2 133 | ) 134 | ) 135 | + 15293365 136 | * 5 ** (1 / 3) 137 | * th3 ** (4 / 3) 138 | / (10838016 * 2 ** (2 / 3) * pi ** (4 / 3) * th0 ** (1 / 3)) 139 | + 617 * pi ** 2 * th0 / (384 * th3 ** 2) 140 | + 5429 / 5376 * (25 * pi * th0 / (2 * th3)) ** (1 / 3) 141 | ) 142 | psi_5 = ( 143 | ( 144 | 140311625 145 | * pi 146 | * th3 ** (2 / 3) 147 | * th3S 148 | / ( 149 | 180348 150 | * ( 151 | 565 * th3 ** (5 / 3) 152 | - 152 * 10 ** (1 / 3) * pi ** (5 / 3) * th0 ** (2 / 3) 153 | ) 154 | ) 155 | ) 156 | + ( 157 | 38645 158 | * (5 / pi) ** (2 / 3) 159 | * th3 ** (5 / 3) 160 | / (64512 * 2 ** (1 / 3) * th0 ** (2 / 3)) 161 | ) 162 | - ( 163 | 732985 164 | * 5 ** (2 / 3) 165 | * th3S 166 | * (th3 / th0) ** (2 / 3) 167 | / (455616 * 2 ** (1 / 3) * pi ** (2 / 3)) 168 | ) 169 | - 85 * pi * th3S / (152 * th3) 170 | - 65 * pi / 384 171 | + phi_0 172 | ) 173 | psi_6 = ( 174 | ( 175 | 15211 176 | * 5 ** (2 / 3) 177 | * pi ** (4 / 3) 178 | * th0 ** (1 / 3) 179 | / (73728 * 2 ** (1 / 3) * th3 ** (4 / 3)) 180 | ) 181 | - 25565 * pi ** 3 * th0 / (27648 * th3 ** 3) 182 | - 535 * jnp.euler_gamma * th3 ** 2 / (112 * pi ** 2 * th0) 183 | + (11583231236531 / (320458457088 * pi ** 2) - 25 / 8) * th3 ** 2 / th0 184 | - 535 * th3 ** 2 / (336 * pi ** 2 * th0) * jnp.log(10 * th3 / (pi * th0)) 185 | + ( 186 | 2255 * 5 ** (1 / 3) * pi ** (5 / 3) / (1024 * 2 ** (2 / 3)) 187 | - 15737765635 * (5 / pi) ** (1 / 3) / (260112384 * 2 ** (2 / 3)) 188 | ) 189 | * (th3 / th0) ** (1 / 3) 190 | ) 191 | psi_7 = ( 192 | 385483375 193 | * 5 ** (1 / 3) 194 | * th3 ** (7 / 3) 195 | / (173408256 * 2 ** (2 / 3) * pi ** (4 / 3) * th0 ** (4 / 3)) 196 | + 378515 * 5 ** (2 / 3) * (pi / 2) ** (1 / 3) / 516096 * (th3 / th0) ** (2 / 3) 197 | - 74045 * pi ** 2 / (129024 * th3) 198 | ) 199 | psi_8 = 2 * pi * f_0 * t_0 200 | 201 | psi_5_L = psi_5 - phi_0 # TODO: does this impact phi_0 maximization??? 202 | psi_6_L = -535 * th3 ** 2 / (336 * pi ** 2 * th0) 203 | 204 | powers = (jnp.arange(0, 9) - 5) / 3 205 | 206 | return ( 207 | psi_0 * (f / f_0) ** powers[0] 208 | + psi_1 * (f / f_0) ** powers[1] 209 | + psi_2 * (f / f_0) ** powers[2] 210 | + psi_3 * (f / f_0) ** powers[3] 211 | + psi_4 * (f / f_0) ** powers[4] 212 | + (psi_5 + psi_5_L * jnp.log(f / f_0)) * (f / f_0) ** powers[5] 213 | + (psi_6 + psi_6_L * jnp.log(f / f_0)) * (f / f_0) ** powers[6] 214 | + psi_7 * (f / f_0) ** powers[7] 215 | + psi_8 * (f / f_0) ** powers[8] 216 | ) 217 | -------------------------------------------------------------------------------- /src/diffbank/waveforms/threePN_simple.py: -------------------------------------------------------------------------------- 1 | from math import pi 2 | 3 | import jax.numpy as jnp 4 | 5 | """ 6 | 3PN metric parametrized by the black hole masses (m1, m2). 7 | """ 8 | 9 | 10 | def Psi(f: jnp.ndarray, theta: jnp.ndarray) -> jnp.ndarray: 11 | m1, m2 = theta 12 | # Mt, eta = theta 13 | gt = 4.92549094830932e-6 # GN*Msun/c**3 in seconds 14 | Mt = m1 + m2 15 | eta = m1 * m2 / (m1 + m2) ** 2 16 | EulerGamma = 0.57721566490153286060 17 | vlso = 1.0 / jnp.sqrt(6.0) 18 | 19 | v = (pi * Mt * (f + 1e-100) * gt) ** (1.0 / 3.0) 20 | v2 = v * v 21 | v3 = v2 * v 22 | v4 = v2 * v2 23 | v5 = v4 * v 24 | v6 = v3 * v3 25 | v7 = v3 * v4 26 | eta2 = eta ** 2 27 | eta3 = eta ** 3 28 | 29 | # # ------------------------- Non spinning part of the waveform 30 | psi_NS_0PN = 1.0 31 | psi_NS_1PN = (3715.0 / 756.0 + 55.0 * eta / 9.0) * v2 32 | psi_NS_15PN = -16.0 * pi * v3 33 | psi_NS_2PN = ( 34 | 15293365.0 / 508032.0 + 27145.0 * eta / 504.0 + 3085.0 * eta2 / 72.0 35 | ) * v4 36 | psi_NS_25PN = ( 37 | pi * (38645.0 / 756.0 - 65.0 * eta / 9.0) * (1 + 3.0 * jnp.log(v / vlso)) * v5 38 | ) 39 | psi_NS_3PN = ( 40 | ( 41 | 11583231236531.0 / 4694215680.0 42 | - 640.0 * pi ** 2 / 3.0 43 | - 6848.0 * EulerGamma / 21.0 44 | ) 45 | + (2255.0 * pi ** 2 / 12.0 - 15737765635.0 / 3048192.0) * eta 46 | + 76055.0 * eta2 / 1728.0 47 | - 127825.0 * eta3 / 1296.0 48 | - 6848.0 * jnp.log(4.0 * v) / 21.0 49 | ) * v6 50 | psi_NS_35PN = ( 51 | pi 52 | * (77096675.0 / 254016.0 + 378515.0 * eta / 1512.0 - 74045.0 * eta2 / 756.0) 53 | * v7 54 | ) 55 | 56 | psi_NS = ( 57 | psi_NS_0PN 58 | + psi_NS_1PN 59 | + psi_NS_15PN 60 | + psi_NS_2PN 61 | + psi_NS_25PN 62 | + psi_NS_3PN 63 | + psi_NS_35PN 64 | ) 65 | 66 | return 3.0 / 128.0 / eta / v5 * (psi_NS) 67 | 68 | 69 | def amp(f: jnp.ndarray, theta: jnp.ndarray) -> jnp.ndarray: 70 | m1, m2 = theta 71 | return f ** (-7 / 6) 72 | -------------------------------------------------------------------------------- /src/diffbank/waveforms/twoPN_chirptimes.py: -------------------------------------------------------------------------------- 1 | from math import pi 2 | 3 | import jax.numpy as jnp 4 | import numpy as np 5 | 6 | from ..constants import MSUN, C, G 7 | from ..utils import ms_to_Mc_eta 8 | 9 | """ 10 | 2PN metric parametrized by the chirp times from https://arxiv.org/pdf/gr-qc/0604037.pdf. 11 | """ 12 | f0 = 20 # Hz 13 | 14 | 15 | def Psi(f: jnp.ndarray, theta: jnp.ndarray) -> jnp.ndarray: 16 | theta1, theta2 = theta 17 | x = f / f0 18 | 19 | a01 = 3 / 5 20 | a21 = 11 * pi / 12 21 | a22 = (743 / 2016) * (25 / 2 / (pi ** 2)) ** (1 / 3) 22 | a31 = -3 / 2 23 | a41 = 617 * (pi ** 2) / 384 24 | a42 = (5429 / 5376) * (25 * pi / 2) ** (1 / 3) 25 | a43 = (15293365 / 10838016) * (5 / 4 / (pi ** 4)) ** (1 / 3) 26 | 27 | psi1 = (a01 * theta1) * (x ** (-5 / 3)) 28 | psi2 = (a21 * theta1 / theta2 + a22 * (theta1 * (theta2 ** 2)) ** (1 / 3)) * ( 29 | x ** (-1) 30 | ) 31 | psi3 = (a31 * theta2) * (x ** (-2 / 3)) 32 | psi4 = ( 33 | a41 * theta1 / (theta2 ** 2) 34 | + a42 * (theta1 / theta2) ** (1 / 3) 35 | + a43 * ((theta2 ** 4) / theta1) ** (1 / 3) 36 | ) * (x ** (-1 / 3)) 37 | 38 | psi = psi1 + psi2 + psi3 + psi4 39 | 40 | return psi 41 | 42 | 43 | def Amp(f: jnp.ndarray, theta: jnp.ndarray) -> jnp.ndarray: 44 | return f ** (-7 / 6) 45 | 46 | 47 | def phys_to_th(phys: jnp.ndarray, f_L) -> jnp.ndarray: 48 | """ 49 | Converts (M, eta) -> (theta1, theta2). 50 | """ 51 | _M_chirp, eta = phys 52 | M_chirp = _M_chirp / (C ** 3 / G) 53 | th0 = 5 / (128 * (f_L * M_chirp * pi) ** (5 / 3)) 54 | th3 = pi ** (1 / 3) / (4 * f_L ** (2 / 3) * M_chirp ** (2 / 3) * eta ** (3 / 5)) 55 | return jnp.stack([th0, th3]) 56 | 57 | 58 | def get_th_boundary_interps(m_min, m_max, f_L, n=200): 59 | """ 60 | Given a range of BH masses, returns corresponding range of th0 and 61 | interpolators for the minimum and maximum corresponding values of th3. 62 | """ 63 | 64 | # Lower boundary 65 | ms = jnp.linspace(m_min, m_max, n) 66 | M_chirps = ms_to_Mc_eta(jnp.stack([ms, ms]))[0] 67 | etas = jnp.full_like(M_chirps, 0.25) 68 | th0_lows, th3_lows = phys_to_th(jnp.stack([M_chirps, etas]), f_L)[:2, ::-1] 69 | interp_low = lambda th0: jnp.interp( 70 | th0, th0_lows, th3_lows, left=jnp.nan, right=jnp.nan 71 | ) 72 | 73 | # Upper boundary 74 | M_chirps, etas = ms_to_Mc_eta(jnp.stack([ms, jnp.full_like(ms, m_max)])) 75 | th0_highs_1, th3_highs_1 = phys_to_th(jnp.stack([M_chirps, etas]), f_L)[:2, ::-1] 76 | 77 | M_chirps, etas = ms_to_Mc_eta(jnp.stack([ms, jnp.full_like(ms, m_min)])) 78 | th0_highs_2, th3_highs_2 = phys_to_th(jnp.stack([M_chirps, etas]), f_L)[:2, ::-1] 79 | 80 | th0_highs = jnp.concatenate([th0_highs_1, th0_highs_2]) 81 | th3_highs = jnp.concatenate([th3_highs_1, th3_highs_2]) 82 | interp_high = lambda th0: jnp.interp( 83 | th0, th0_highs, th3_highs, left=jnp.nan, right=jnp.nan 84 | ) 85 | 86 | th0_min = min(th0_lows.min(), th0_highs_1.min()) 87 | th0_max = max(th0_lows.max(), th0_highs_2.max()) 88 | 89 | return (th0_min, th0_max), interp_low, interp_high 90 | 91 | 92 | ############################################################## 93 | ####### Functions for the analytic metric 94 | ############################################################## 95 | 96 | 97 | def calc_Iq(q, fs, Sn): 98 | Sns = Sn(fs) 99 | xs = fs / f0 100 | Sf0 = Sn(f0) 101 | Iq = np.trapz(Sf0 * (xs ** (-q / 3) / Sns), xs) 102 | return Iq 103 | 104 | 105 | def calc_Jq(q, fs, Sn): 106 | I7 = calc_Iq(7, fs, Sn) 107 | Iq = calc_Iq(q, fs, Sn) 108 | return Iq / I7 109 | 110 | 111 | def analytic_metric(f, theta, Sn): 112 | theta1, theta2 = theta 113 | 114 | a01 = 3 / 5 115 | a21 = 11 * pi / 12 116 | a22 = (743 / 2016) * (25 / 2 / (pi ** 2)) ** (1 / 3) 117 | a31 = -3 / 2 118 | a41 = 617 * (pi ** 2) / 384 119 | a42 = (5429 / 5376) * (25 * pi / 2) ** (1 / 3) 120 | a43 = (15293365 / 10838016) * (5 / 4 / (pi ** 4)) ** (1 / 3) 121 | 122 | Psi = np.array( 123 | [ 124 | [ 125 | a01, 126 | 0, 127 | a21 / theta2 + (a22 / 3) * (theta2 / theta1) ** (2 / 3), 128 | 0, 129 | (a41 / (theta2 ** 2)) 130 | + a42 / (3 * ((theta1 ** 2) * theta2) ** (1 / 3)) 131 | - (a43 / 3) * (theta2 / theta1) ** (4 / 3), 132 | ], 133 | [ 134 | 0, 135 | 0, 136 | -a21 * theta1 / (theta2 ** 2) 137 | + (2 * a22 / 3) * ((theta1 / theta2) ** (1 / 3)), 138 | a31, 139 | -2 * a41 * theta1 / (theta2 ** 3) 140 | - (a42 / 3) * (theta1 / (theta2 ** 4)) ** (1 / 3) 141 | + (4 * a43 / 3) * (theta2 / theta1) ** (1 / 3), 142 | ], 143 | ] 144 | ) 145 | 146 | def J_combined(k, j): 147 | J_comb = calc_Jq(17 - k - j, f, Sn) 148 | J_comb -= calc_Jq(12 - k, f, Sn) * calc_Jq(12 - j, f, Sn) 149 | J_comb -= ( 150 | (calc_Jq(9 - k, f, Sn) - calc_Jq(4, f, Sn) * calc_Jq(12 - k, f, Sn)) 151 | * (calc_Jq(9 - j, f, Sn) - calc_Jq(4, f, Sn) * calc_Jq(12 - j, f, Sn)) 152 | ) / (calc_Jq(1, f, Sn) - (calc_Jq(4, f, Sn) ** 2)) 153 | return J_comb 154 | 155 | def fill_g(m, l): 156 | gml = 0.0 157 | for k in range(5): 158 | for j in range(5): 159 | gml += Psi[m, k] * Psi[l, j] * J_combined(k, j) 160 | return gml * (1 / 2) 161 | 162 | g = np.zeros((2, 2)) 163 | g[0, 0] = fill_g(0, 0) 164 | g[0, 1] = fill_g(0, 1) 165 | g[1, 0] = fill_g(1, 0) 166 | g[1, 1] = fill_g(1, 1) 167 | return g 168 | -------------------------------------------------------------------------------- /src/diffbank/waveforms/twoPN_simple.py: -------------------------------------------------------------------------------- 1 | from math import pi 2 | 3 | import jax.numpy as jnp 4 | 5 | from ..constants import MSUN, C, G 6 | 7 | """ 8 | 2PN metric parametrized by the black hole masses (Mt, eta). 9 | """ 10 | 11 | 12 | def Psi(f: jnp.ndarray, theta: jnp.ndarray) -> jnp.ndarray: 13 | M, eta = theta 14 | M *= MSUN 15 | x = pi * G * M * f / C ** 3 # dimensionless 16 | return ( 17 | 3 18 | / (128 * eta) 19 | * x ** (-5 / 3) 20 | * ( 21 | 1 22 | + 20 / 9 * (743 / 336 + 11 / 4 * eta) * x ** (2 / 3) 23 | - 16 * pi * x 24 | + 10 25 | * (3058673 / 1016064 + 5429 / 1008 * eta + 617 / 144 * eta ** 2) 26 | * x ** (4 / 3) 27 | ) 28 | ) 29 | 30 | 31 | def ms_to_Meta(ms: jnp.ndarray) -> jnp.ndarray: 32 | m1, m2 = ms 33 | mu = m1 * m2 / (m1 + m2) # reduced mass 34 | M = m1 + m2 # total mass 35 | eta = mu / M 36 | return jnp.array([M, eta]) 37 | 38 | 39 | # def Psi(f: jnp.ndarray, theta: jnp.ndarray) -> jnp.ndarray: 40 | # return Psi_Meta(f, ms_to_Meta(theta)) 41 | 42 | 43 | def Amp(f: jnp.ndarray, theta: jnp.ndarray) -> jnp.ndarray: 44 | return f ** (-7 / 6) 45 | -------------------------------------------------------------------------------- /tests/README.md: -------------------------------------------------------------------------------- 1 | # Tests 2 | 3 | Run with `pytest`. 4 | -------------------------------------------------------------------------------- /tests/test_bank.py: -------------------------------------------------------------------------------- 1 | import jax.numpy as jnp 2 | from jax import random 3 | 4 | from diffbank.bank import Bank 5 | from diffbank.noise import Sn_LIGOI 6 | from diffbank.utils import get_m1_m2_sampler 7 | from diffbank.waveforms.threePN_simple import Psi, amp 8 | 9 | """ 10 | Make sure the bank's saving/loading, generation and effectualness calculation 11 | functionality works. 12 | """ 13 | 14 | 15 | def test_save_load(): 16 | """ 17 | Test saving and loading. 18 | """ 19 | amp = lambda f, t: None 20 | Psi = lambda f, t: None 21 | Sn = lambda f: None 22 | sample_base = lambda k, n: jnp.array([0.0, 0.9]) 23 | bank = Bank( 24 | amp, 25 | Psi, 26 | jnp.array([0.0, 1.0, 2.0]), 27 | Sn, 28 | 0.05, 29 | 0.99, 30 | sample_base, 31 | name="save_test", 32 | ) 33 | bank.ratio_max = jnp.array(100.0) 34 | bank.n_templates = jnp.array(2, dtype=int) 35 | bank.templates = jnp.array([[5.0, 4.0], [7.0, 79.0]]) 36 | bank.effectualnesses = jnp.array([1.0, 0.2, 3.0, 6.0, 1.0, 45.0, 9.0]) 37 | bank.effectualness_points = jnp.array([[0.0, 0.0], [1.0, 1.0], [2.0, 3.0]]) 38 | bank.eta_est = 0.92 39 | bank.eta_est_err = 0.001 40 | bank.save() 41 | bank_str = str(bank) 42 | print(bank_str) 43 | 44 | loaded_bank = Bank.load("save_test.npz", amp, Psi, Sn, sample_base) 45 | loaded_bank_str = str(loaded_bank) 46 | print(loaded_bank_str) 47 | 48 | # Computed variables 49 | assert bank.ratio_max == loaded_bank.ratio_max 50 | assert bank.n_templates == loaded_bank.n_templates 51 | assert jnp.all(bank.templates == loaded_bank.templates) 52 | assert jnp.all(bank.effectualness_points == loaded_bank.effectualness_points) 53 | assert jnp.all(bank.effectualnesses == loaded_bank.effectualnesses) 54 | assert bank.eta_est == bank.eta_est 55 | assert bank.eta_est_err == bank.eta_est_err 56 | assert bank.dim == loaded_bank.dim 57 | # Provided variables 58 | assert jnp.all(bank.fs == loaded_bank.fs) 59 | assert bank.m_star == loaded_bank.m_star 60 | assert bank.eta == loaded_bank.eta 61 | assert bank.name == loaded_bank.name 62 | # Functions 63 | assert bank.amp is loaded_bank.amp 64 | assert bank.Psi is loaded_bank.Psi 65 | assert bank.Sn is loaded_bank.Sn 66 | assert bank.sample_base is loaded_bank.sample_base 67 | 68 | 69 | def test_gen(): 70 | """ 71 | Make sure template bank generation works. 72 | """ 73 | key = random.PRNGKey(84) 74 | 75 | fs = jnp.linspace(20.0, 2000.0, 300) 76 | m_range = (2.8, 3.0) 77 | sampler = get_m1_m2_sampler(m_range, m_range) 78 | bank = Bank( 79 | amp, 80 | Psi, 81 | fs, 82 | Sn_LIGOI, 83 | m_star=1 - 0.8, 84 | eta=0.8, 85 | sample_base=sampler, 86 | name="3PN", 87 | ) 88 | 89 | bank.ratio_max = bank.density_fun(jnp.array([m_range[1], m_range[0]])) 90 | 91 | for kind in ["random", "stochastic"]: 92 | print(f"Testing {kind} bank") 93 | key, subkey = random.split(key) 94 | bank.fill_bank(subkey, kind) 95 | 96 | # Make sure templates are in bounds 97 | for m1, m2 in bank.templates: 98 | assert m1 >= m_range[0] and m1 <= m_range[1] 99 | assert m2 >= m_range[0] and m2 <= m_range[1] 100 | 101 | key, subkey = random.split(key) 102 | bank.calc_bank_effectualness(subkey, 10) 103 | print(f"eta = {bank.eta_est:.3f} +/- {bank.eta_est_err:.3f}\n") 104 | 105 | 106 | if __name__ == "__main__": 107 | test_save_load() 108 | test_gen() 109 | -------------------------------------------------------------------------------- /tests/test_checkpointing.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import jax 4 | import jax.numpy as jnp 5 | from jax import random 6 | from scipy.optimize import minimize_scalar 7 | 8 | from diffbank.bank import Bank 9 | from diffbank.utils import get_m1_m2_sampler 10 | from diffbank.waveforms.threePN_simple import Psi, amp 11 | 12 | """ 13 | Generates a small 3.5PN-2D bank to comprehensively test ``diffbank``. 14 | """ 15 | 16 | ##### Frequency settings 17 | f_u = 512.0 # Hz 18 | f_l = 24.0 # Hz 19 | N_fbins = 4880 20 | seed = 10 21 | ##### 22 | 23 | 24 | def test_checkpointing( 25 | kind="random", 26 | n_eta=10, 27 | mm=0.95, 28 | eta_star=0.9, 29 | n_eff=10, 30 | savedir="./", 31 | device="cpu", 32 | noise="analytic", 33 | ): 34 | jax.config.update("jax_platform_name", device) 35 | 36 | key = random.PRNGKey(seed) 37 | m_star = 1 - mm 38 | 39 | fs = jnp.linspace(f_l, f_u, N_fbins) 40 | m_range = (2.0, 3.0) 41 | sampler = get_m1_m2_sampler(m_range, m_range) 42 | if noise == "interpolated": 43 | from diffbank.noise import Sn_O3a as Sn 44 | elif noise == "analytic": 45 | from diffbank.noise import Sn_LIGOI as Sn 46 | else: 47 | raise ValueError("invalid 'noise' argument") 48 | 49 | bank = Bank( 50 | amp, 51 | Psi, 52 | fs, 53 | Sn, 54 | m_star, 55 | eta_star, 56 | sampler, 57 | name=f"3pn-{kind}-{seed}-mm={mm}-eta_star={eta_star}-n_eff={n_eff}", 58 | ) 59 | 60 | # Get max density 61 | fun = lambda m1: -bank.density_fun(jnp.array([m1, m_range[0]])) 62 | res = minimize_scalar(fun, bracket=m_range, bounds=m_range, method="bounded") 63 | assert res.success 64 | theta_dmax = jnp.array([res.x, m_range[0]]) 65 | ratio_max = bank.density_fun(theta_dmax) 66 | bank.ratio_max = ratio_max 67 | 68 | # Fill bank 69 | key, subkey = random.split(key) 70 | bank.fill_bank(subkey, kind, n_eff, save_interval=5) 71 | bank.save(savedir) 72 | print(f"Saved bank to {os.path.join(savedir, bank.name + '.npz')}") 73 | 74 | file = jnp.load( 75 | f"3pn-{kind}-{seed}-mm={mm}-eta_star={eta_star}-n_eff={n_eff}-checkpoint.npz", 76 | allow_pickle=True, 77 | ) 78 | temp = file["templates"] 79 | eff_points = file["eff_pts"] 80 | effs = file["eff"] 81 | key = file["key"] 82 | 83 | bank_checkpoint = Bank( 84 | amp, 85 | Psi, 86 | fs, 87 | Sn, 88 | m_star, 89 | eta_star, 90 | sampler, 91 | name=f"3pn-{kind}-{seed}-mm={mm}-eta_star={eta_star}-n_eff={n_eff}", 92 | ) 93 | bank_checkpoint.templates = temp 94 | bank_checkpoint.n_templates = len(bank.templates) 95 | bank_checkpoint.effectualness_points = eff_points 96 | 97 | # Get max density 98 | fun = lambda m1: -bank_checkpoint.density_fun(jnp.array([m1, m_range[0]])) 99 | res = minimize_scalar(fun, bracket=m_range, bounds=m_range, method="bounded") 100 | assert res.success 101 | theta_dmax = jnp.array([res.x, m_range[0]]) 102 | ratio_max = bank_checkpoint.density_fun(theta_dmax) 103 | bank_checkpoint.ratio_max = ratio_max 104 | 105 | # Fill bank from intermediate point 106 | bank_checkpoint.fill_bank(key, kind, n_eff, effs=effs) 107 | print("Both banks have the same size!") 108 | print(bank.templates.shape, bank_checkpoint.templates.shape) 109 | 110 | 111 | if __name__ == "__main__": 112 | test_checkpointing() 113 | -------------------------------------------------------------------------------- /tests/test_eff.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | import jax.numpy as jnp 4 | from jax import jit, vmap 5 | 6 | from diffbank.noise import Sn_LIGOI as Sn 7 | from diffbank.utils import get_eff_pads, get_match, get_phase_maximized_inner_product 8 | from diffbank.waveforms.threePN_simple import Psi, amp 9 | 10 | """ 11 | Runs various tests on the effectualness calculation to make sure it works and 12 | that the FFT properly maximizes over del_t. 13 | """ 14 | 15 | 16 | def test_match(): 17 | """ 18 | Checks that match function works and times it. 19 | """ 20 | theta_1 = jnp.array([3.0, 0.8]) 21 | theta_2 = jnp.array([3.0, 0.8]) 22 | fs = jnp.linspace(10, 500, 10000) 23 | pad_low, pad_high = get_eff_pads(fs) 24 | 25 | match_fn = jit(get_match, static_argnums=(2, 3, 4, 5, 7)) 26 | 27 | t0 = time.time() 28 | for _ in range(1000): 29 | match_fn( 30 | theta_1, theta_2, amp, Psi, amp, Psi, fs, Sn, pad_low, pad_high 31 | ).block_until_ready() 32 | 33 | t1 = time.time() 34 | 35 | print((t1 - t0) / 1000) 36 | 37 | 38 | def test_matches(): 39 | """ 40 | Makes sure the match function with an explicit del_t argument agrees with 41 | the FFT one which maximizes over del_t. 42 | """ 43 | theta_1 = jnp.array([2.20013935, 1.13180361]) 44 | theta_2 = jnp.array([2.27775711, 1.30239947]) 45 | del_ts = jnp.linspace(-0.5, 0.1, 2000) 46 | fs = jnp.linspace(10, 500, 10000) 47 | pad_low, pad_high = get_eff_pads(fs) 48 | 49 | matches = vmap( 50 | get_phase_maximized_inner_product, 51 | in_axes=(None, None, 0, None, None, None, None, None, None), 52 | )(theta_1, theta_2, del_ts, amp, Psi, amp, Psi, fs, Sn).max() 53 | match_fft = get_match( 54 | theta_1, theta_2, amp, Psi, amp, Psi, fs, Sn, pad_low, pad_high 55 | ) 56 | 57 | assert jnp.allclose(match_fft, matches.max(), rtol=5e-3) 58 | 59 | 60 | if __name__ == "__main__": 61 | test_match() 62 | test_matches() 63 | -------------------------------------------------------------------------------- /tests/test_est_ratio_max.py: -------------------------------------------------------------------------------- 1 | import jax 2 | import jax.numpy as jnp 3 | from jax import random 4 | from scipy.optimize import minimize_scalar 5 | 6 | from diffbank.bank import Bank 7 | from diffbank.noise import Sn_LIGOI 8 | from diffbank.utils import get_m1_m2_sampler 9 | from diffbank.waveforms.threePN_simple import Psi, amp 10 | 11 | """ 12 | Tests empirical supremum rejection sampling estimate of maximum target vs base 13 | density ratio. 14 | """ 15 | 16 | 17 | def test(): 18 | key = random.PRNGKey(90) 19 | 20 | # Setup 21 | minimum_match = 0.95 22 | m_star = 1 - minimum_match 23 | eta_star = 0.99 24 | fs = jnp.linspace(20.0, 2000.0, 1000) 25 | m_range = (1.4, 5.0) 26 | sampler = get_m1_m2_sampler(m_range, m_range) 27 | bank = Bank(amp, Psi, fs, Sn_LIGOI, m_star, eta_star, sampler) 28 | 29 | # Get max density 30 | fun = lambda m1: -bank.density_fun(jnp.array([m1, m_range[0]])) 31 | res = minimize_scalar(fun, bracket=m_range, bounds=m_range) 32 | assert res.success 33 | theta_dmax = jnp.array([res.x, m_range[0]]) 34 | ratio_max = bank.density_fun(theta_dmax) 35 | bank.ratio_max = ratio_max 36 | 37 | key, subkey = random.split(key) 38 | ratio_max_est = bank.est_ratio_max(subkey)[0] 39 | assert abs(bank.ratio_max - ratio_max_est) / bank.ratio_max < 0.1 40 | 41 | 42 | if __name__ == "__main__": 43 | test() 44 | -------------------------------------------------------------------------------- /tests/test_genbank_2D_threePN.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import jax 4 | import jax.numpy as jnp 5 | from jax import random 6 | from scipy.optimize import minimize_scalar 7 | 8 | from diffbank.bank import Bank 9 | from diffbank.utils import get_m1_m2_sampler 10 | from diffbank.waveforms.threePN_simple import Psi, amp 11 | 12 | """ 13 | Generates a small 3.5PN-2D bank to comprehensively test ``diffbank``. 14 | """ 15 | 16 | ##### Frequency settings 17 | f_u = 512.0 # Hz 18 | f_l = 24.0 # Hz 19 | N_fbins = 4880 20 | ##### 21 | 22 | 23 | def test_gen_2D_threePN( 24 | seed=1, 25 | kind="random", 26 | n_eta=10, 27 | mm=0.95, 28 | eta_star=0.9, 29 | n_eff=10, 30 | savedir="./", 31 | device="cpu", 32 | noise="analytic", 33 | ): 34 | jax.config.update("jax_platform_name", device) 35 | 36 | key = random.PRNGKey(seed) 37 | m_star = 1 - mm 38 | 39 | fs = jnp.linspace(f_l, f_u, N_fbins) 40 | m_range = (2.0, 3.0) 41 | sampler = get_m1_m2_sampler(m_range, m_range) 42 | if noise == "interpolated": 43 | from diffbank.noise import Sn_O3a as Sn 44 | elif noise == "analytic": 45 | from diffbank.noise import Sn_LIGOI as Sn 46 | else: 47 | raise ValueError("invalid 'noise' argument") 48 | 49 | bank = Bank( 50 | amp, 51 | Psi, 52 | fs, 53 | Sn, 54 | m_star, 55 | eta_star, 56 | sampler, 57 | name=f"3pn-{kind}-{seed}-mm={mm}-eta_star={eta_star}-n_eff={n_eff}", 58 | ) 59 | 60 | # Get max density 61 | fun = lambda m1: -bank.density_fun(jnp.array([m1, m_range[0]])) 62 | res = minimize_scalar(fun, bracket=m_range, bounds=m_range, method="bounded") 63 | assert res.success 64 | theta_dmax = jnp.array([res.x, m_range[0]]) 65 | ratio_max = bank.density_fun(theta_dmax) 66 | bank.ratio_max = ratio_max 67 | 68 | # Fill bank 69 | key, subkey = random.split(key) 70 | bank.fill_bank(subkey, kind, n_eff) 71 | bank.save(savedir) 72 | print(f"Saved bank to {os.path.join(savedir, bank.name + '.npz')}") 73 | 74 | # Get effectualnesses 75 | if n_eta > 0: 76 | key, subkey = random.split(key) 77 | bank.calc_bank_effectualness(subkey, n_eta) 78 | bank.save(savedir) 79 | else: 80 | print("Skipping effectualnesses calculation") 81 | 82 | 83 | if __name__ == "__main__": 84 | test_gen_2D_threePN() 85 | -------------------------------------------------------------------------------- /tests/test_metric_speed.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | import jax.numpy as jnp 4 | from jax import random 5 | 6 | from diffbank.metric import get_g 7 | from diffbank.noise import Sn_LIGOI 8 | from diffbank.utils import get_phase_maximized_inner_product 9 | from diffbank.waveforms.threePN_simple import Psi, amp 10 | 11 | 12 | def test_metric_speed(): 13 | """ 14 | Tests that the metric calculation runs, and times it. 15 | """ 16 | f = jnp.linspace(30.0, 512.0, 1000) 17 | 18 | print( 19 | "Match:", 20 | get_phase_maximized_inner_product( 21 | jnp.array([3, 1.0]), 22 | jnp.array([3.0, 1.01]), 23 | jnp.array(0.0), 24 | amp, 25 | Psi, 26 | amp, 27 | Psi, 28 | f, 29 | Sn_LIGOI, 30 | ) 31 | ) 32 | 33 | print( 34 | "Metric:", 35 | get_g(jnp.array([3, 1.0]), amp, Psi, f, Sn_LIGOI), 36 | ) 37 | 38 | N = 1000 39 | key = random.PRNGKey(89) 40 | m1_key, m2_key = random.split(key) 41 | m1 = random.uniform(m1_key, minval=2.0, maxval=3.0, shape=(N,)) 42 | m2 = random.uniform(m2_key, minval=1.0, maxval=2.0, shape=(N,)) 43 | 44 | t0 = time.time() 45 | for i in range(N): 46 | pars = jnp.array([m1[i], m2[i]]) 47 | get_g(pars, amp, Psi, f, Sn_LIGOI).block_until_ready() 48 | t1 = time.time() 49 | 50 | total = t1 - t0 51 | print("Time per metric evaluation:", total / N) 52 | 53 | 54 | if __name__ == "__main__": 55 | test_metric_speed() 56 | -------------------------------------------------------------------------------- /tests/test_taylorf2_metric.py: -------------------------------------------------------------------------------- 1 | import jax.numpy as jnp 2 | from jax import jit 3 | 4 | from diffbank.metric import get_g, get_gam 5 | from diffbank.noise import Sn_LIGOI 6 | from diffbank.waveforms.taylorF2 import Amp as _amp 7 | from diffbank.waveforms.taylorF2 import Psi as _Psi 8 | 9 | """ 10 | Makes sure the TaylorF2 metric calculation doesn't give infinities or nans. 11 | """ 12 | 13 | 14 | def test_taylorf2_dets(return_vals=False): 15 | amp = lambda f, theta: jit(_amp)(f, theta) 16 | Psi = lambda f, theta: jit(_Psi)(f, theta) 17 | 18 | m1 = jnp.array(3.0) 19 | m2 = jnp.array(1.5) 20 | 21 | f_min, f_max = jnp.array(10.0), jnp.array(512.0) 22 | fs = jnp.linspace(f_min, f_max, 10000) 23 | 24 | gam = get_gam(jnp.array([m1, m2, 0.3, 0.5]), amp, Psi, fs, Sn_LIGOI) 25 | g = get_g(jnp.array([m1, m2, 0.3, 0.5]), amp, Psi, fs, Sn_LIGOI) 26 | 27 | det_g = jnp.linalg.det(g) 28 | 29 | assert not (jnp.isinf(det_g) or jnp.isnan(det_g) or det_g <= 0) 30 | print(g, jnp.linalg.det(gam), det_g) 31 | 32 | if return_vals: 33 | return jnp.linalg.det(gam), det_g 34 | 35 | 36 | if __name__ == "__main__": 37 | test_taylorf2_dets(return_vals=True) 38 | --------------------------------------------------------------------------------