├── example ├── 179.pi.nh ├── 179.pi.z ├── 179.pi ├── cdfs4Ms_179.arf ├── cdfs4Ms_179.rmf ├── cdfs4Ms_179_bkg.pi └── cdfs4Ms_179_bkg.pi_model.fits ├── HISTORY.rst ├── pexmon.pdf ├── pip-requirements.txt ├── requirements_dev.txt ├── fastxsf ├── __init__.py ├── flux.py ├── data.py ├── model.py └── response.py ├── setup.cfg ├── MANIFEST.in ├── plotresponse.py ├── pyproject.toml ├── setup.py ├── LICENSE ├── .gitignore ├── NHdist.py ├── tests ├── test_lum.py └── test_table.py ├── Makefile ├── .github └── workflows │ └── tests.yml ├── CONTRIBUTING.rst ├── simple.py ├── README.rst ├── NHfast.py ├── simplev.py ├── simpleopt.py ├── reflopt.py ├── scripts └── xagnfitter.py └── multispecopt.py /example/179.pi.nh: -------------------------------------------------------------------------------- 1 | 8.8e19 2 | -------------------------------------------------------------------------------- /example/179.pi.z: -------------------------------------------------------------------------------- 1 | 0.605 2 | -------------------------------------------------------------------------------- /HISTORY.rst: -------------------------------------------------------------------------------- 1 | 1.0.0 2 | ----- 3 | First version 4 | 5 | -------------------------------------------------------------------------------- /pexmon.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JohannesBuchner/fastXSF/main/pexmon.pdf -------------------------------------------------------------------------------- /example/179.pi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JohannesBuchner/fastXSF/main/example/179.pi -------------------------------------------------------------------------------- /example/cdfs4Ms_179.arf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JohannesBuchner/fastXSF/main/example/cdfs4Ms_179.arf -------------------------------------------------------------------------------- /example/cdfs4Ms_179.rmf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JohannesBuchner/fastXSF/main/example/cdfs4Ms_179.rmf -------------------------------------------------------------------------------- /example/cdfs4Ms_179_bkg.pi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JohannesBuchner/fastXSF/main/example/cdfs4Ms_179_bkg.pi -------------------------------------------------------------------------------- /pip-requirements.txt: -------------------------------------------------------------------------------- 1 | numpy 2 | scipy 3 | matplotlib 4 | optns 5 | tqdm 6 | jax 7 | ultranest 8 | h5py 9 | astropy 10 | -------------------------------------------------------------------------------- /example/cdfs4Ms_179_bkg.pi_model.fits: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JohannesBuchner/fastXSF/main/example/cdfs4Ms_179_bkg.pi_model.fits -------------------------------------------------------------------------------- /requirements_dev.txt: -------------------------------------------------------------------------------- 1 | pip 2 | wheel 3 | flake8 4 | coverage 5 | Sphinx 6 | twine 7 | h5py 8 | numpy 9 | scipy 10 | matplotlib 11 | flake8 12 | pycodestyle 13 | pydocstyle 14 | coverage 15 | coverage-lcov 16 | pytest 17 | pytest-runner 18 | pytest-html 19 | pytest-xdist 20 | nbstripout 21 | nbsphinx 22 | sphinx_rtd_theme 23 | -------------------------------------------------------------------------------- /fastxsf/__init__.py: -------------------------------------------------------------------------------- 1 | """Fast X-ray spectral fitting.""" 2 | 3 | __author__ = """Johannes Buchner""" 4 | __email__ = 'johannes.buchner.acad@gmx.com' 5 | __version__ = '1.1.0' 6 | 7 | from .data import load_pha 8 | from .model import (FixedTable, Table, logPoissonPDF, logPoissonPDF_vectorized, 9 | x, xvec) 10 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [flake8] 2 | style = numpy 3 | check-return-types = False 4 | exclude = docs 5 | extend-ignore = E203,W503,E501,F401,E128,E231,E124,SIM114,DOC105,DOC106,DOC107,DOC301,DOC501,DOC503,DOC203,B006,SIM102,SIM113,DOC202,DOC403,DOC404 6 | max-line-length = 160 7 | 8 | [aliases] 9 | test = pytest 10 | 11 | 12 | [tool:pytest] 13 | collect_ignore = ['setup.py'] 14 | addopts = --junitxml=test-reports/junit.xml --html=tests/reports/index.html 15 | 16 | 17 | [pycodestyle] 18 | count = False 19 | ignore = W191,W291,W293,E231 20 | max-line-length = 160 21 | statistics = False 22 | 23 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include CONTRIBUTING.rst 2 | include HISTORY.rst 3 | include LICENSE 4 | include README.rst 5 | include requirements_dev.txt 6 | include pip-requirements.txt 7 | 8 | recursive-include * *.pyx 9 | recursive-include * *.pxd 10 | 11 | recursive-include tests * 12 | recursive-exclude * __pycache__ 13 | recursive-exclude * *.py[co] 14 | recursive-exclude * *.c 15 | recursive-exclude * *.orig 16 | recursive-exclude * *.pdf 17 | 18 | recursive-include docs *.rst conf.py Makefile make.bat *.jpg *.png *.gif 19 | 20 | # remove extraneous doc and test outputs 21 | prune tests/reports 22 | prune tests/.pytype 23 | recursive-exclude tests *.pdf 24 | -------------------------------------------------------------------------------- /plotresponse.py: -------------------------------------------------------------------------------- 1 | import sys 2 | 3 | import matplotlib.pyplot as plt 4 | import numpy as np 5 | import fastxsf 6 | 7 | filename = sys.argv[1] 8 | elo = float(sys.argv[2]) 9 | ehi = float(sys.argv[3]) 10 | 11 | data = fastxsf.load_pha(filename, elo, ehi) 12 | plt.imshow(data['RMF'], cmap='viridis') 13 | plt.colorbar() 14 | plt.xlabel('Energy') 15 | plt.ylabel('Energy channel') 16 | plt.savefig(sys.argv[1] + '_rmf.pdf', bbox_inches='tight') 17 | plt.close() 18 | 19 | plt.title('exposure: %d area: %f' % (data['src_expo'], data['src_expoarea'] / data['src_expo'])) 20 | plt.plot(data['ARF']) 21 | plt.ylabel('Sensitive area') 22 | plt.xlabel('Energy channel') 23 | plt.savefig(sys.argv[1] + '_arf.pdf', bbox_inches='tight') 24 | plt.close() 25 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | # Minimum requirements for the build system to execute. 3 | requires = ["setuptools", "wheel"] # PEP 508 specifications. 4 | build-backend = "setuptools.build_meta" 5 | 6 | [project] 7 | name = "fastxsf" 8 | version = "1.1.0" 9 | description = "Fast X-ray spectral fitting" 10 | readme = "README.rst" 11 | requires-python = ">=3.8" 12 | dependencies = ["numpy", "scipy", "scikit-learn", "tqdm", "jax", "ultranest", "corner", "matplotlib"] 13 | license = {file = "LICENSE"} 14 | authors = [ 15 | { name = "Johannes Buchner", email = "johannes.buchner.acad@gmx.com" }, 16 | ] 17 | keywords = ["astronomy", "X-ray", "spectroscopy", "spectral fitting", "X-ray spectroscopy", "Bayesian inference", "astrophysics"] 18 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | try: 5 | from setuptools import setup 6 | except: 7 | from distutils.core import setup 8 | 9 | import re 10 | 11 | with open('README.rst', encoding="utf-8") as readme_file: 12 | readme = readme_file.read() 13 | 14 | with open('HISTORY.rst', encoding="utf-8") as history_file: 15 | history = re.sub(r':py:class:`([^`]+)`', r'\1', 16 | history_file.read()) 17 | 18 | requirements = ['numpy', 'scipy', 'matplotlib', 'corner', 'optns', 'tqdm', 'ultranest', 'astropy'] 19 | 20 | setup_requirements = ['pytest-runner', ] 21 | 22 | test_requirements = ['pytest>=3', ] 23 | 24 | setup( 25 | python_requires='>=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*, !=3.5.*, !=3.6.*, !=3.7.*, !=3.8.*', 26 | install_requires=requirements, 27 | long_description=readme + '\n\n' + history, 28 | name='fastxsf', 29 | packages=['fastxsf'], 30 | setup_requires=setup_requirements, 31 | test_suite='tests', 32 | tests_require=test_requirements, 33 | url='https://github.com/JohannesBuchner/fastXSF', 34 | ) 35 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | GNU GENERAL PUBLIC LICENSE 2 | Version 3, 29 June 2007 3 | 4 | Fit and compare complex models reliably and rapidly. Advanced Nested Sampling. 5 | Copyright (C) 2019 Johannes Buchner 6 | 7 | This program is free software: you can redistribute it and/or modify 8 | it under the terms of the GNU General Public License as published by 9 | the Free Software Foundation, either version 3 of the License, or 10 | (at your option) any later version. 11 | 12 | This program is distributed in the hope that it will be useful, 13 | but WITHOUT ANY WARRANTY; without even the implied warranty of 14 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 | GNU General Public License for more details. 16 | 17 | You should have received a copy of the GNU General Public License 18 | along with this program. If not, see . 19 | 20 | Also add information on how to contact you by electronic and paper mail. 21 | 22 | You should also get your employer (if you work as a programmer) or school, 23 | if any, to sign a "copyright disclaimer" for the program, if necessary. 24 | For more information on this, and how to apply and follow the GNU GPL, see 25 | . 26 | 27 | The GNU General Public License does not permit incorporating your program 28 | into proprietary programs. If your program is a subroutine library, you 29 | may consider it more useful to permit linking proprietary applications with 30 | the library. If this is what you want to do, use the GNU Lesser General 31 | Public License instead of this License. But first, please read 32 | . 33 | 34 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | env/ 12 | build/ 13 | develop-eggs/ 14 | dist/ 15 | downloads/ 16 | eggs/ 17 | .eggs/ 18 | lib/ 19 | lib64/ 20 | parts/ 21 | sdist/ 22 | var/ 23 | wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | .pytest_cache/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | 58 | # Flask stuff: 59 | instance/ 60 | .webassets-cache 61 | 62 | # Scrapy stuff: 63 | .scrapy 64 | 65 | # Sphinx documentation 66 | docs/_build/ 67 | docs/joblib/ 68 | docs/example*.py 69 | 70 | # PyBuilder 71 | target/ 72 | 73 | # Jupyter Notebook 74 | .ipynb_checkpoints 75 | 76 | # pyenv 77 | .python-version 78 | 79 | # celery beat schedule file 80 | celerybeat-schedule 81 | 82 | # SageMath parsed files 83 | *.sage.py 84 | 85 | # dotenv 86 | .env 87 | 88 | # virtualenv 89 | .venv 90 | venv/ 91 | ENV/ 92 | 93 | # Spyder project settings 94 | .spyderproject 95 | .spyproject 96 | 97 | # Rope project settings 98 | .ropeproject 99 | 100 | # mkdocs documentation 101 | /site 102 | 103 | # mypy 104 | .mypy_cache/ 105 | 106 | *~ 107 | *.pyc 108 | *.pdf 109 | *.png 110 | *.log 111 | *.npz 112 | 113 | # cache 114 | joblib/ 115 | # profiling 116 | prof* 117 | # patches 118 | *.patch 119 | # old files 120 | *.orig 121 | -------------------------------------------------------------------------------- /NHdist.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import numpy as np 3 | import scipy.stats 4 | from scipy.special import logsumexp 5 | from matplotlib import pyplot as plt 6 | import tqdm 7 | 8 | filenames = sys.argv[1:] 9 | 10 | PhoIndex_gauss = scipy.stats.norm(1.95, 0.15) 11 | 12 | PhoIndex_grid = np.arange(1, 3.1, 0.1) 13 | logNH_grid = np.arange(20, 25.1, 0.1) 14 | 15 | PhoIndex_logprob = PhoIndex_gauss.logpdf(PhoIndex_grid) 16 | 17 | profile_likes = [np.loadtxt(filename) for filename in filenames] 18 | 19 | NHmean_grid = np.arange(20, 25, 0.2) 20 | NHstd_grid = np.arange(0.1, 4, 0.1) 21 | 22 | # for each NHmean, NHstd 23 | # for each spectrum 24 | # for each possible Gamma, NH 25 | # compute profile likelihood (of background and source) 26 | # store luminosity 27 | # compute marginal likelihood integrating over Gamma, NH with distributions 28 | # compute product of marginal likelihoods 29 | # compute surface of NHmean, NHstds 30 | # no sampling needed! 31 | 32 | marglike = np.zeros((len(NHmean_grid), len(NHstd_grid))) 33 | for i, NHmean in enumerate(tqdm.tqdm(NHmean_grid)): 34 | for j, NHstd in enumerate(NHstd_grid): 35 | NHlogprob = -0.5 * ((logNH_grid - NHmean) / NHstd)**2 36 | NHlogprob -= logsumexp(NHlogprob) 37 | 38 | # compute marginal likelihood integrating over Gamma, NH with distributions 39 | for profile_like in profile_likes: 40 | logprob = profile_like + PhoIndex_logprob.reshape((-1, 1)) + NHlogprob.reshape((1, -11)) 41 | marglike[i,j] += logsumexp(logprob) 42 | 43 | extent = [NHstd_grid.min(), NHstd_grid.max(), NHmean_grid.min(), NHmean_grid.max()] 44 | plt.imshow(-2 * (marglike - marglike.max()), vmin=0, vmax=10, 45 | extent=extent, origin='lower', cmap='Greys_r', aspect='auto') 46 | plt.ylabel(r'mean($\log N_\mathrm{H}$)') 47 | plt.xlabel(r'std($\log N_\mathrm{H}$)') 48 | plt.colorbar(orientation='horizontal') 49 | plt.contour(-2 * (marglike - marglike.max()), levels=[1, 2, 3], 50 | extent=extent, origin='lower', colors=['k'] * 3) 51 | plt.savefig(f'distNH_like.pdf') 52 | plt.close() 53 | 54 | dchi2 = -2 * (marglike - marglike.max()) 55 | 56 | for i, NHmean in enumerate(tqdm.tqdm(NHmean_grid)): 57 | for j, NHstd in enumerate(NHstd_grid): 58 | if dchi2[i,j] < 3: 59 | NHprob = np.exp(-0.5 * ((logNH_grid - NHmean) / NHstd)**2) 60 | NHprob /= NHprob.sum() 61 | color = 'r' if dchi2[i,j] < 1 else 'orange' if dchi2[i,j] < 2 else 'yellow' 62 | plt.plot(10**logNH_grid, NHprob, color=color, alpha=0.25) 63 | plt.xscale('log') 64 | plt.xlabel(r'Column density $N_\mathrm{H}$ [#/cm$^2$]') 65 | plt.savefig('distNH_curves.pdf') 66 | 67 | -------------------------------------------------------------------------------- /tests/test_lum.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import fastxsf 3 | from astropy import units as u 4 | from fastxsf.flux import energy_flux, photon_flux, luminosity, frac_overlap_interval 5 | import os 6 | from astropy.cosmology import LambdaCDM 7 | from numpy.testing import assert_allclose 8 | 9 | cosmo = LambdaCDM(H0=70, Om0=0.3, Ode0=0.730) 10 | 11 | # fastxsf.x.chatter(0) 12 | fastxsf.x.abundance('wilm') 13 | fastxsf.x.cross_section('vern') 14 | 15 | # load the spectrum, where we will consider data from 0.5 to 8 keV 16 | data = fastxsf.load_pha(os.path.join(os.path.dirname(__file__), '../example/179.pi'), 0.5, 8) 17 | energies = np.append(data['e_lo'], data['e_hi'][-1]) 18 | chan_e = np.append(data['chan_e_min'], data['chan_e_max'][-1]) 19 | 20 | print(data['src_expoarea']) 21 | 22 | def test_frac_overlap_interval(): 23 | assert_allclose(frac_overlap_interval(np.asarray([0, 1, 2, 3]), 1, 2), [0, 1, 0]) 24 | assert_allclose(frac_overlap_interval(np.asarray([0, 1, 2, 3]), 1.5, 2.5), [0, 0.5, 0.5]) 25 | assert_allclose(frac_overlap_interval(np.asarray([0.1, 1.1, 2.1, 3.1]), 1.5, 2.5), [0, 0.6, 0.4]) 26 | assert_allclose(frac_overlap_interval(np.asarray([0.1, 0.2, 0.3, 0.4]), 0.1, 0.12), [0.2, 0, 0]) 27 | assert_allclose(frac_overlap_interval(np.asarray([0.1, 0.2, 0.3, 0.4]), 0.28, 0.3), [0, 0.2, 0]) 28 | assert_allclose(frac_overlap_interval(np.asarray([0.1, 0.2, 0.3, 0.4]), 0.38, 0.4), [0, 0, 0.2]) 29 | 30 | def test_powerlaw_fluxes(): 31 | pl = fastxsf.x.zpowerlw(energies=energies, pars=[1, 0]) 32 | # print(photon_flux(pl, energies, 2.001, 2.002), 1.6014e-12 * u.erg/u.cm**2/u.s) 33 | assert np.isclose(photon_flux(pl, energies, 2.001, 2.002), 0.00049938 / u.cm**2/u.s, atol=1e-14) 34 | assert np.isclose(energy_flux(pl, energies, 2.001, 2.002), 1.6014e-12 * u.erg/u.cm**2/u.s, atol=1e-14) 35 | assert np.isclose(energy_flux(pl, energies, 2, 8), 9.6132e-09 * u.erg/u.cm**2/u.s, atol=1e-13) 36 | 37 | zpl0 = fastxsf.x.zpowerlw(energies=energies, pars=[1, 0.01]) 38 | assert np.isclose(energy_flux(zpl0, energies, 2, 8), 9.518e-09 * u.erg/u.cm**2/u.s, atol=1e-13) 39 | assert np.isclose(luminosity(zpl0, energies, 2, 8, 0.01, cosmo=cosmo), 2.0971e+45 * u.erg/u.s, rtol=0.04) 40 | 41 | zpl = fastxsf.x.zpowerlw(energies=energies, pars=[1, 0.5]) 42 | assert np.isclose(energy_flux(zpl, energies, 2, 8), 6.4088e-09 * u.erg/u.cm**2/u.s, atol=1e-11) 43 | assert np.isclose(luminosity(zpl, energies, 2, 10, 0.5, cosmo=cosmo), 5.5941e+48 * u.erg/u.s, rtol=0.04) 44 | 45 | zpl3 = fastxsf.x.zpowerlw(energies=energies, pars=[3, 0.5]) 46 | assert np.isclose(energy_flux(zpl3, energies, 2, 8), 1.7802e-10 * u.erg/u.cm**2/u.s, atol=1e-11) 47 | assert np.isclose(luminosity(zpl3, energies, 2, 10, 0.5, cosmo=cosmo), 2.7971e+47 * u.erg/u.s, rtol=0.04) 48 | assert np.isclose(photon_flux(zpl3, energies, 2, 8), 0.034722 / u.cm**2/u.s, atol=1e-14) 49 | 50 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | .PHONY: clean clean-test clean-pyc clean-build docs servedocs help install release release-test dist 2 | .DEFAULT_GOAL := help 3 | 4 | define BROWSER_PYSCRIPT 5 | import os, webbrowser, sys 6 | 7 | try: 8 | from urllib import pathname2url 9 | except: 10 | from urllib.request import pathname2url 11 | 12 | webbrowser.open("file://" + pathname2url(os.path.abspath(sys.argv[1]))) 13 | endef 14 | export BROWSER_PYSCRIPT 15 | 16 | define PRINT_HELP_PYSCRIPT 17 | import re, sys 18 | 19 | for line in sys.stdin: 20 | match = re.match(r'^([a-zA-Z_-]+):.*?## (.*)$$', line) 21 | if match: 22 | target, help = match.groups() 23 | print("%-20s %s" % (target, help)) 24 | endef 25 | export PRINT_HELP_PYSCRIPT 26 | 27 | PYTHON := python3 28 | 29 | BROWSER := $(PYTHON) -c "$$BROWSER_PYSCRIPT" 30 | 31 | help: 32 | @$(PYTHON) -c "$$PRINT_HELP_PYSCRIPT" < $(MAKEFILE_LIST) 33 | 34 | clean: clean-build clean-pyc clean-test clean-doc ## remove all build, test, coverage and Python artifacts 35 | 36 | clean-build: ## remove build artifacts 37 | rm -fr build/ 38 | rm -fr dist/ 39 | rm -fr .eggs/ 40 | find . -name '*.egg-info' -exec rm -fr {} + 41 | find . -name '*.egg' -exec rm -f {} + 42 | 43 | clean-pyc: ## remove Python file artifacts 44 | find . -name '*.pyc' -exec rm -f {} + 45 | find . -name '*.pyo' -exec rm -f {} + 46 | find . -name '*.pyx.py' -exec rm -f {} + 47 | find . -name '*~' -exec rm -f {} + 48 | find . -name '__pycache__' -exec rm -fr {} + 49 | find . -name '*.so' -exec rm -f {} + 50 | find fastxsf -name '*.c' -exec rm -f {} + 51 | 52 | clean-test: ## remove test and coverage artifacts 53 | rm -fr .tox/ 54 | rm -f .coverage 55 | rm -fr htmlcov/ 56 | rm -fr .pytest_cache 57 | 58 | clean-doc: 59 | rm -rf docs/build 60 | 61 | SOURCES := $(shell ls fastxsf/*.py) 62 | 63 | lint: ${SOURCES} ## check style 64 | flake8 ${SOURCES} 65 | pycodestyle ${SOURCES} 66 | pydocstyle ${SOURCES} 67 | 68 | test: ## run tests quickly with the default Python 69 | PYTHONPATH=. pytest 70 | 71 | test-all: ## run tests on every Python version with tox 72 | tox 73 | 74 | coverage: ## check code coverage quickly with the default Python 75 | PYTHONPATH=. coverage run --source fastxsf -m pytest 76 | coverage report -m 77 | coverage html 78 | $(BROWSER) htmlcov/index.html 79 | 80 | docs: ## generate Sphinx HTML documentation, including API docs 81 | rm -f docs/fastxsf.rst 82 | rm -f docs/modules.rst 83 | rm -f docs/API.rst 84 | python3 setup.py build_ext --inplace 85 | sphinx-apidoc -H API -o docs/ fastxsf 86 | cd docs; python3 modoverview.py 87 | $(MAKE) -C docs clean 88 | $(MAKE) -C docs html O=-jauto 89 | $(BROWSER) docs/build/html/index.html 90 | 91 | servedocs: docs ## compile the docs watching for changes 92 | watchmedo shell-command -p '*.rst' -c '$(MAKE) -C docs html' -R -D . 93 | 94 | release: dist ## package and upload a release 95 | twine upload --verbose dist/*.tar.gz 96 | 97 | dist: clean ## builds source and wheel package 98 | uv build 99 | ls -l dist 100 | 101 | install: clean ## install the package to the active Python's site-packages 102 | $(PYTHON) setup.py install --user 103 | -------------------------------------------------------------------------------- /.github/workflows/tests.yml: -------------------------------------------------------------------------------- 1 | name: build 2 | 3 | on: 4 | push: 5 | 6 | jobs: 7 | build: 8 | runs-on: ubuntu-latest 9 | timeout-minutes: 30 10 | strategy: 11 | fail-fast: false 12 | matrix: 13 | python-version: [3.12] 14 | 15 | defaults: 16 | run: 17 | # this is needed, because otherwise conda env is not available 18 | shell: bash -l {0} 19 | 20 | steps: 21 | - uses: actions/checkout@v4 22 | 23 | - name: Install system dependencies 24 | run: sudo apt-get update && sudo apt-get -y --no-install-recommends install -y ghostscript 25 | 26 | - uses: mamba-org/setup-micromamba@v1 27 | with: 28 | environment-name: test 29 | cache-environment: false 30 | cache-downloads: true 31 | 32 | - name: Set directory names 33 | run: | 34 | echo "MODELDIR=$HOME/Downloads/models" >> $GITHUB_ENV 35 | echo "MYCACHE=$HOME/Downloads/xspec" >> $GITHUB_ENV 36 | echo "PATH=$PATH:$HOME/.local/bin/" >> $GITHUB_ENV 37 | 38 | - name: Cache models 39 | uses: pat-s/always-upload-cache@v3.0.11 40 | id: cache-downloads 41 | with: 42 | path: ${{ env.MODELDIR }} 43 | key: cache-downloads 44 | 45 | - name: Download models (if necessary) 46 | run: | 47 | mkdir -p $MODELDIR 48 | pushd $MODELDIR 49 | wget -q -nc https://zenodo.org/record/1169181/files/uxclumpy-cutoff.fits https://zenodo.org/records/2224651/files/wedge.fits https://zenodo.org/records/2224472/files/diskreflect.fits 50 | popd 51 | 52 | - name: Install python dependencies 53 | run: | 54 | micromamba install --override-channels -c https://cxc.cfa.harvard.edu/conda/ciao -c conda-forge xspec-modelsonly "matplotlib>=3.5" ultranest "coverage<7.0.0" coveralls==3.3.1 scipy jax h5py astropy requests cython tqdm coverage toml flake8 pycodestyle pydocstyle pytest pytest-html pytest-xdist h5py joblib && 55 | echo "--- Environment dump start ---" && 56 | env && 57 | echo "--- Environment dump end ---" && 58 | pip install git+https://github.com/JohannesBuchner/coverage-lcov && 59 | sudo sed -i '/PDF/s/none/read|write/' /etc/ImageMagick-6/policy.xml && 60 | pip uninstall -y h5py && 61 | pip install --no-cache-dir -r pip-requirements.txt git+https://github.com/cxcsds/xspec-models-cxc 62 | 63 | - name: Conda info 64 | run: micromamba info 65 | - name: Conda list 66 | run: micromamba list 67 | - name: Conda paths 68 | run: | 69 | pwd 70 | echo $PATH 71 | ls $CONDA/bin/ 72 | which coverage 73 | 74 | - name: Lint with flake8 75 | run: flake8 fastxsf 76 | 77 | - name: Check code style 78 | run: pycodestyle fastxsf 79 | 80 | - name: Check doc style 81 | run: pydocstyle fastxsf 82 | 83 | - name: Check install 84 | run: | 85 | python -c "import numpy"; 86 | python -c "import xspec_models_cxc" 87 | 88 | - name: Prepare testing 89 | run: | 90 | echo "backend: Agg" > matplotlibrc 91 | 92 | - name: Test with pytest 93 | run: PYTHONPATH=. coverage run -p -m pytest 94 | 95 | - name: Run simple example 96 | run: PYTHONPATH=. coverage run -p simple.py 97 | 98 | - name: Run vectorized example 99 | run: PYTHONPATH=. coverage run -p simplev.py 100 | 101 | #- name: Run optimized nested sampling example 102 | # run: PYTHONPATH=. coverage run -p optsimple.py 103 | 104 | - name: Install package 105 | run: pip install . 106 | 107 | - name: Coverage report 108 | run: | 109 | coverage combine 110 | coverage report 111 | coverage-lcov 112 | # make paths relative 113 | sed -i s,$PWD/,,g lcov.info 114 | 115 | - name: Coveralls Finished 116 | uses: coverallsapp/github-action@master 117 | with: 118 | path-to-lcov: lcov.info 119 | github-token: ${{ secrets.github_token }} 120 | -------------------------------------------------------------------------------- /CONTRIBUTING.rst: -------------------------------------------------------------------------------- 1 | .. highlight:: shell 2 | 3 | ============ 4 | Contributing 5 | ============ 6 | 7 | Contributions are welcome, and they are greatly appreciated! Every little bit 8 | helps, and credit will always be given. 9 | 10 | You can contribute in many ways: 11 | 12 | Types of Contributions 13 | ---------------------- 14 | 15 | Report Bugs 16 | ~~~~~~~~~~~ 17 | 18 | Report bugs at https://github.com/JohannesBuchner/fastxsf/issues. 19 | 20 | If you are reporting a bug, please include: 21 | 22 | * Your operating system name and version. 23 | * Any details about your local setup that might be helpful in troubleshooting. 24 | * Detailed steps to reproduce the bug. 25 | 26 | Fix Bugs 27 | ~~~~~~~~ 28 | 29 | Look through the GitHub issues for bugs. Anything tagged with "bug" and "help 30 | wanted" is open to whoever wants to implement it. 31 | 32 | Implement Features 33 | ~~~~~~~~~~~~~~~~~~ 34 | 35 | Look through the GitHub issues for features. Anything tagged with "enhancement" 36 | and "help wanted" is open to whoever wants to implement it. 37 | 38 | Write Documentation 39 | ~~~~~~~~~~~~~~~~~~~ 40 | 41 | fastxsf could always use more documentation, whether as part of the 42 | official fastxsf docs, in docstrings, or even on the web in blog posts, 43 | articles, and such. 44 | 45 | Notebooks demonstrating how to use fastxsf are also appreciated. 46 | 47 | Submit Feedback 48 | ~~~~~~~~~~~~~~~ 49 | 50 | The best way to send feedback is to file an issue at https://github.com/JohannesBuchner/fastxsf/issues. 51 | 52 | If you are proposing a feature: 53 | 54 | * Explain in detail how it would work. 55 | * Keep the scope as narrow as possible, to make it easier to implement. 56 | * Remember that this is a volunteer-driven project, and that contributions 57 | are welcome :) 58 | 59 | Get Started! 60 | ------------ 61 | 62 | Ready to contribute? Here's how to set up `fastxsf` for local development. 63 | 64 | 1. Fork the `fastxsf` repo on GitHub. 65 | 2. Clone your fork locally:: 66 | 67 | $ git clone git@github.com:JohannesBuchner/fastxsf.git 68 | 69 | 3. Install your local copy into a virtualenv. Assuming you have virtualenvwrapper installed, this is how you set up your fork for local development:: 70 | 71 | $ mkvirtualenv fastxsf 72 | $ cd fastxsf/ 73 | $ python setup.py develop 74 | 75 | 4. Create a branch for local development:: 76 | 77 | $ git checkout -b name-of-your-bugfix-or-feature 78 | 79 | Now you can make your changes locally. 80 | 81 | 5. When you're done making changes, check that your changes pass flake8 and the 82 | tests, including testing other Python versions with tox:: 83 | 84 | $ flake8 fastxsf tests 85 | $ PYTHONPATH=. pytest 86 | $ tox 87 | 88 | To get flake8 and tox, just pip install them into your virtualenv. 89 | 90 | 6. Commit your changes and push your branch to GitHub:: 91 | 92 | $ git add . 93 | $ git commit -m "Your detailed description of your changes." 94 | $ git push origin name-of-your-bugfix-or-feature 95 | 96 | 7. Submit a pull request through the GitHub website. 97 | 98 | Pull Request Guidelines 99 | ----------------------- 100 | 101 | Before you submit a pull request, check that it meets these guidelines: 102 | 103 | 1. The pull request should include tests. 104 | 2. If the pull request adds functionality, the docs should be updated. Put 105 | your new functionality into a function with a docstring, and add the 106 | feature to the list in README.rst. 107 | 3. The pull request should work for Python 2.7, 3.5, 3.6 and 3.7, and for PyPy. Check 108 | https://github.com/JohannesBuchner/fastxsf/actions/ 109 | and make sure that the tests pass for all supported Python versions. 110 | 111 | Tips 112 | ---- 113 | 114 | To run a subset of tests:: 115 | 116 | $ PYTHONPATH=. pytest tests.test_utils 117 | 118 | 119 | Deploying 120 | --------- 121 | 122 | A reminder for the maintainers on how to deploy. 123 | Make sure all your changes are committed (including an entry in HISTORY.rst). 124 | Then run:: 125 | 126 | $ bump2version patch # possible: major / minor / patch 127 | $ git push 128 | $ git push --tags 129 | 130 | Travis will then deploy to PyPI if tests pass. 131 | -------------------------------------------------------------------------------- /simple.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import numpy as np 4 | import scipy.stats 5 | from matplotlib import pyplot as plt 6 | 7 | import fastxsf 8 | 9 | # fastxsf.x.chatter(0) 10 | fastxsf.x.abundance('wilm') 11 | fastxsf.x.cross_section('vern') 12 | 13 | # load the spectrum, where we will consider data from 0.5 to 8 keV 14 | data = fastxsf.load_pha('example/179.pi', 0.5, 8) 15 | 16 | # fetch some basic information about our spectrum 17 | e_lo = data['e_lo'] 18 | e_hi = data['e_hi'] 19 | e_mid = (data['e_hi'] + data['e_lo']) / 2. 20 | e_width = data['e_hi'] - data['e_lo'] 21 | energies = np.append(e_lo, e_hi[-1]) 22 | RMF_src = data['RMF_src'] 23 | 24 | chan_e = (data['chan_e_min'] + data['chan_e_max']) / 2. 25 | 26 | # load a Table model 27 | absAGN = fastxsf.Table(os.path.join(os.environ.get('MODELDIR', '.'), 'uxclumpy-cutoff.fits')) 28 | 29 | # pre-compute the absorption factors -- no need to call this again and again if the parameters do not change! 30 | galabso = fastxsf.x.TBabs(energies=energies, pars=[data['galnh']]) 31 | 32 | z = data['redshift'] 33 | # define the model parameters: 34 | bkg_norm = 1.0 35 | norm = 3e-7 36 | scat_norm = norm * 0.08 37 | PhoIndex = 2.0 38 | TORsigma = 28.0 39 | CTKcover = 0.1 40 | Incl = 45.0 41 | NH22 = 1.0 42 | Ecut = 400 43 | 44 | 45 | # define a likelihood 46 | def loglikelihood(params, plot=False): 47 | norm, NH22, rel_scat_norm, PhoIndex, TORsigma, CTKcover, Incl, Ecut, bkg_norm = params 48 | 49 | scat_norm = norm * rel_scat_norm 50 | # here we are taking z from the global context -- so it is fixed! 51 | abs_component = absAGN(energies=energies, pars=[NH22, PhoIndex, Ecut, TORsigma, CTKcover, Incl, z]) 52 | 53 | scat_component = fastxsf.x.zpowerlw(energies=energies, pars=[norm, PhoIndex]) 54 | 55 | pred_spec = abs_component * norm + scat_component * scat_norm 56 | 57 | pred_counts_src_srcreg = RMF_src.apply_rmf(data['ARF'] * (galabso * pred_spec))[data['chan_mask']] * data['src_expoarea'] 58 | pred_counts_bkg_srcreg = data['bkg_model_src_region'] * bkg_norm * data['src_expoarea'] 59 | pred_counts_srcreg = pred_counts_src_srcreg + pred_counts_bkg_srcreg 60 | pred_counts_bkg_bkgreg = data['bkg_model_bkg_region'] * bkg_norm * data['bkg_expoarea'] 61 | 62 | if plot: 63 | plt.figure() 64 | plt.legend() 65 | plt.plot(data['chan_e_min'], data['src_region_counts'] / (data['chan_e_max'] - data['chan_e_min']), 'o', label='data', mfc='none') 66 | plt.plot(data['chan_e_min'], pred_counts_srcreg / (data['chan_e_max'] - data['chan_e_min']), label='src+bkg') 67 | plt.plot(data['chan_e_min'], pred_counts_src_srcreg / (data['chan_e_max'] - data['chan_e_min']), label='src') 68 | plt.plot(data['chan_e_min'], pred_counts_bkg_srcreg / (data['chan_e_max'] - data['chan_e_min']), label='bkg') 69 | plt.xlabel('Channel Energy [keV]') 70 | plt.ylabel('Counts / keV') 71 | plt.legend() 72 | plt.savefig('src_region_counts.pdf') 73 | plt.close() 74 | 75 | plt.figure() 76 | plt.plot(data['chan_e_min'], data['bkg_region_counts'] / (data['chan_e_max'] - data['chan_e_min']), 'o', label='data', mfc='none') 77 | plt.plot(data['chan_e_min'], pred_counts_bkg_bkgreg / (data['chan_e_max'] - data['chan_e_min']), label='bkg') 78 | plt.xlabel('Channel Energy [keV]') 79 | plt.ylabel('Counts / keV') 80 | plt.legend() 81 | plt.savefig('bkg_region_counts.pdf') 82 | plt.close() 83 | 84 | # compute log Poisson probability 85 | like_srcreg = fastxsf.logPoissonPDF(pred_counts_srcreg, data['src_region_counts']) 86 | like_bkgreg = fastxsf.logPoissonPDF(pred_counts_bkg_bkgreg, data['bkg_region_counts']) 87 | return like_srcreg + like_bkgreg 88 | 89 | 90 | # lets define a prior 91 | PhoIndex_gauss = scipy.stats.norm(1.95, 0.15) 92 | 93 | 94 | # define the prior transform function 95 | def prior_transform(cube): 96 | params = cube.copy() 97 | # uniform from 1e-10 to 1 98 | params[0] = 10**(cube[0] * -10) 99 | # uniform from 1e-2 to 1e2 100 | params[1] = 10**(cube[1] * (2 - -2) + -2) 101 | params[2] = 10**(cube[2] * (-1 - -5) + -5) 102 | # Gaussian prior 103 | params[3] = PhoIndex_gauss.ppf(cube[3]) 104 | # uniform priors 105 | params[4] = cube[4] * (80 - 7) + 7 106 | params[5] = cube[5] * (0.4) + 0 107 | params[6] = cube[6] * 90 108 | params[7] = cube[7] * (400 - 300) + 300 109 | # log-uniform prior on the background normalisation between 0.1 and 10 110 | params[8] = 10**(cube[8] * (1 - -1) + -1) 111 | return params 112 | 113 | 114 | # compute a likelihood: 115 | print(loglikelihood((norm, NH22, 0.08, PhoIndex, TORsigma, CTKcover, Incl, Ecut, bkg_norm), plot=True)) 116 | -------------------------------------------------------------------------------- /README.rst: -------------------------------------------------------------------------------- 1 | FastXSF 2 | ------- 3 | 4 | Fast X-ray spectral fitting. 5 | 6 | .. image:: https://img.shields.io/pypi/v/fastxsf.svg 7 | :target: https://pypi.python.org/pypi/fastxsf 8 | 9 | .. image:: https://github.com/JohannesBuchner/fastxsf/actions/workflows/tests.yml/badge.svg 10 | :target: https://github.com/JohannesBuchner/fastxsf/actions/workflows/tests.yml 11 | 12 | .. image:: https://coveralls.io/repos/github/JohannesBuchner/fastxsf/badge.svg?branch=main 13 | :target: https://coveralls.io/github/JohannesBuchner/fastxsf?branch=main 14 | 15 | .. image:: https://img.shields.io/badge/GitHub-JohannesBuchner%2Ffastxsf-blue.svg?style=flat 16 | :target: https://github.com/JohannesBuchner/fastxsf/ 17 | :alt: Github repository 18 | 19 | Background 20 | ---------- 21 | 22 | Currently, there are the following issues in modern X-ray spectral fitting software: 23 | 24 | 1. Response matrices have become huge (e.g. XRISM, NuSTAR), making models slow to evaluate. 25 | 2. XSPEC is not developed openly and its quirks make it difficult to build upon it and extend it. 26 | 3. Yet models are maintained by the community in XSPEC. 27 | 4. Maintaining additional software packages requires substantial institutional efforts (CXC: sherpa, SRON: spex). 28 | 5. Not all models are differentiable. Reimplementing them in a differentiable language one by one is a significant effort with little recognition. 29 | Nevertheless, it has been and is being tried. Yet such reimplementations tend to fade out (see also 3ML astromodels). 30 | 6. Inference parameter spaces are complicated, with multiple modes and other complicated degeneracies being common in X-ray spectral fitting. 31 | 7. `Bayesian model comparison `_ is powerful and we want it. 32 | 8. The X-ray community is walled off from other communities by vendor lock-in and its own odd terminology. But `X-ray spectral fitting is neither complicated nor a special case! `_ 33 | 34 | Therefore, we want: 35 | 36 | 1) A performant software package 37 | 2) All community packages from XSPEC 38 | 3) Nested sampling for model comparison and robust parameter estimation 39 | 4) Minimum package maintainance effort 40 | 41 | FastXSF does that. 42 | 43 | xspex&jaxspec do 1, xspec/sherpa+BXA does 2+3. 44 | 45 | FastXSF is a few hundred lines of code. 46 | 47 | Approach 48 | -------- 49 | 50 | 1) Vectorization. 51 | Folding the spectrum through the RMF is vectorized. 52 | Handling many proposed spectra at once keeps memory low and efficiency high: 53 | Each chunk of the response matrix is applied to all spectra. 54 | Modern Bayesian samplers such as UltraNest can handle vectorized functions. 55 | 56 | 2) Building upon the CXC (Doug Burke's) wrapper for Xspec models. https://github.com/cxcsds/xspec-models-cxc/ 57 | All XSPEC models are available for use! 58 | 59 | 3) Some further niceties (all optional) include handling of backgrounds, redshifts and galactic NH: 60 | 61 | * Use BXA's autobackground folder to create a background spectral model from your background region. 62 | * Use BXA's galnh.py to fetch the galactic NH for the position of your observation and store it in my.pha.nh as a string (e.g. 1.2e20). 63 | * Store the redshift in my.pha.z as a string. 64 | 65 | 4) We treat X-ray spectral fitting as a normal inference problem like any other! 66 | 67 | Define a likelihood, prior and call a sampler. No need to carry around 68 | legacy awkwardness such as chi-square, C-stat, 69 | background-subtraction, identify matrix folding, multi-source to data mappings. 70 | 71 | Installation 72 | ------------ 73 | 74 | Prerequisites: 75 | 76 | * install and load xspec/heasoft 77 | * install https://github.com/cxcsds/xspec-models-cxc/ 78 | * install ultranest (pip install ultranest) 79 | * download this repository and enter it from the command line. 80 | 81 | Test with:: 82 | 83 | `python -c 'import xspec_models_cxc; import ultranest'` 84 | 85 | Getting started 86 | --------------- 87 | 88 | To start, have a look at simple.py, which demonstrates: 89 | 90 | * loading a spectrum 91 | * loading a ATable (download the table from `the xars models page `_) 92 | * setting up a XSPEC model 93 | * passing the model through the ARF and RMF 94 | * adding a background model 95 | * plotting the spectrum of the source and background model on top of the data 96 | * computing the likelihood and print it 97 | 98 | Next, the vectorization is in simplev.py, which demonstrates the same as above plus: 99 | 100 | * vectorized handling of proposals 101 | * launching UltraNest for sampling the posterior, make corner plots. 102 | 103 | Next, there is simpleopt.py, which demonstrates optimized nested sampling (optNS). 104 | This is much faster, especially when there are many components with free normalisations. 105 | 106 | Take it for a spin and adapt it! 107 | 108 | Todo 109 | ---- 110 | 111 | * ✓ Profile where the code is still slow. 112 | * ✓ There is a python loop in fastxsf/model.py::Table.__call__ which should be replaced with something smarter 113 | * ✓ Compute fluxes and luminosities. 114 | * ✓ Create some unit tests for loading and evaluating atables/mtables, poisson probability, plotting ARF/RMF. 115 | * ✓ Make a atable that is precomputed for a given rest-frame energy grid at fixed redshift. -> FixedTable 116 | 117 | Credits 118 | -------- 119 | 120 | This builds upon work by the Chandra X-ray Center (in particular Doug Burke's wrapper), 121 | and Daniela Huppenkothen's RMF/ARF reader (based in turn on sherpa code, IIRC). 122 | 123 | License: GPL v3 124 | 125 | Contributions are welcome. 126 | -------------------------------------------------------------------------------- /NHfast.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | from astropy.cosmology import LambdaCDM 4 | from astropy import units as u 5 | import numpy as np 6 | import tqdm 7 | from matplotlib import pyplot as plt 8 | from optns.profilelike import ComponentModel 9 | import fastxsf 10 | from fastxsf.flux import luminosity 11 | from fastxsf.model import FixedTable 12 | 13 | cosmo = LambdaCDM(H0=70, Om0=0.3, Ode0=0.730) 14 | 15 | 16 | # fastxsf.x.chatter(0) 17 | fastxsf.x.abundance('wilm') 18 | fastxsf.x.cross_section('vern') 19 | 20 | PhoIndex_grid = np.arange(1, 3.1, 0.1) 21 | PhoIndex_grid[-1] = 3.0 22 | logNH_grid = np.arange(20, 25.1, 0.1) 23 | 24 | filename = sys.argv[1] 25 | 26 | elo = 0.3 27 | ehi = 8 28 | data = fastxsf.load_pha(filename, elo, ehi) 29 | # fetch some basic information about our spectrum 30 | e_lo = data['e_lo'] 31 | e_hi = data['e_hi'] 32 | #e_mid = (data['e_hi'] + data['e_lo']) / 2. 33 | #e_width = data['e_hi'] - data['e_lo'] 34 | energies = np.append(e_lo, e_hi[-1]) 35 | RMF_src = data['RMF_src'] 36 | #chan_e = (data['chan_e_min'] + data['chan_e_max']) / 2. 37 | 38 | # pre-compute the absorption factors -- no need to call this again and again if the parameters do not change! 39 | galabso = fastxsf.x.TBabs(energies=energies, pars=[data['galnh']]) 40 | 41 | z = data['redshift'] 42 | absAGN = FixedTable( 43 | os.path.join(os.environ.get('MODELDIR', '.'), 'uxclumpy-cutoff.fits'), 44 | energies=energies, redshift=z) 45 | 46 | Nsrc_chan = len(data['src_region_counts']) 47 | Nbkg_chan = len(data['bkg_region_counts']) 48 | counts_flat = np.hstack((data['src_region_counts'], data['bkg_region_counts'])) 49 | nonlinear_param_names = ['logNH', 'PhoIndex'] 50 | linear_param_names = ['norm_src', 'norm_bkg'] 51 | # create OptNS object, and give it all of these ingredients, 52 | # as well as our data 53 | statmodel = ComponentModel(2, counts_flat) 54 | 55 | TORsigma = 28.0 56 | CTKcover = 0.1 57 | Incl = 45.0 58 | Ecut = 400 59 | 60 | def compute_model_components(params): 61 | logNH, PhoIndex = params 62 | # first component: a absorbed power law 63 | plabso = absAGN(energies=energies, pars=[10**(logNH - 22), PhoIndex, Ecut, TORsigma, CTKcover, Incl]) 64 | 65 | # now we need to project all of our components through the response. 66 | src_components = data['ARF'] * galabso * plabso 67 | pred_counts_src_srcreg = RMF_src.apply_rmf(src_components)[data['chan_mask']] * data['src_expoarea'] 68 | # add non-folded background to source region components 69 | pred_counts = np.zeros((2, Nsrc_chan + Nbkg_chan)) 70 | # the three folded source components in the source region 71 | pred_counts[0, :Nsrc_chan] = pred_counts_src_srcreg 72 | # the unfolded background components in the source region 73 | pred_counts[1, :Nsrc_chan] = data['bkg_model_src_region'] * data['src_expoarea'] 74 | # the unfolded background components in the background region 75 | pred_counts[1, Nsrc_chan:] = data['bkg_model_bkg_region'] * data['bkg_expoarea'] 76 | # notice how the source does not affect the background: 77 | # pred_counts[0, Nsrc_chan:] = 0 # they remain zero 78 | return pred_counts.T 79 | 80 | 81 | profile_like = np.zeros((len(PhoIndex_grid), len(logNH_grid))) 82 | Lint = np.zeros((len(PhoIndex_grid), len(logNH_grid))) 83 | extent = [logNH_grid.min(), logNH_grid.max(), PhoIndex_grid.min(), PhoIndex_grid.max()] 84 | 85 | for i, PhoIndex in enumerate(tqdm.tqdm(PhoIndex_grid)): 86 | for j, logNH in enumerate(logNH_grid): 87 | X = compute_model_components([logNH, PhoIndex]) 88 | res = statmodel.loglike_poisson_optimize(X) 89 | norms = res.x 90 | profile_like[i,j] = -res.fun 91 | srcnorm = np.exp(norms[0]) 92 | Lint[i,j] = np.log10(luminosity( 93 | srcnorm * fastxsf.x.zpowerlw(energies=energies, pars=[PhoIndex, z]), 94 | energies, 2, 10, z, cosmo 95 | ) / (u.erg/u.s)) 96 | # plot likelihoods and 97 | plt.imshow(Lint, vmin=42, vmax=47, extent=extent, cmap='rainbow') 98 | plt.colorbar(orientation='horizontal') 99 | plt.savefig(f'{filename}_L.pdf') 100 | plt.close() 101 | 102 | plt.imshow(-2 * (profile_like - profile_like.max()), vmin=0, vmax=10, extent=extent, origin='upper', cmap='Greys_r') 103 | plt.colorbar(orientation='horizontal') 104 | plt.contour(-2 * (profile_like - profile_like.max()), levels=[1, 2, 3], extent=extent, origin='upper', colors=['k'] * 3) 105 | plt.savefig(f'{filename}_like.pdf') 106 | plt.close() 107 | 108 | # compute posterior distribution of logNH, L 109 | prob = np.exp(profile_like - profile_like.max()) 110 | logNH_probs = np.mean(prob, axis=0) 111 | logNH_probs /= logNH_probs.sum() 112 | fig, axs = plt.subplots(1, 2, sharey=True, figsize=(7, 2)) 113 | axs[0].plot(10**logNH_grid, logNH_probs / logNH_probs.max()) 114 | Lgrid = np.arange(max(42, Lint.min() - 0.1), Lint.max() + 0.1, 0.05) 115 | Lgrid_probs = 0 * Lgrid 116 | for i, PhoIndex in enumerate(tqdm.tqdm(PhoIndex_grid)): 117 | for j, logNH in enumerate(logNH_grid): 118 | k = np.argmin(np.abs(Lint[i,j] - Lgrid)) 119 | Lgrid_probs[k] += prob[i,j] 120 | axs[1].plot(10**Lgrid, Lgrid_probs / Lgrid_probs.max()) 121 | axs[0].set_xlabel(r'Column density $N_\mathrm{H}$ [#/cm$^2$]') 122 | axs[1].set_xlabel(r'Luminosity (2-10keV, intr.) [erg/s]') 123 | axs[0].set_xscale('log') 124 | axs[1].set_xscale('log') 125 | #axs[0].set_xticks(10**np.arange(20, 26)) 126 | #axs[1].set_xticks(10**np.arange(int(Lgrid.min()), int(Lgrid.max()))) 127 | axs[0].set_yticks([0,1]) 128 | plt.savefig(f'{filename}_prob.pdf') 129 | plt.savefig(f'{filename}_prob.png') 130 | plt.close() 131 | 132 | np.savetxt(f'{filename}_like.txt', profile_like, fmt='%.3f') 133 | np.savetxt(f'{filename}_L.txt', [Lgrid, Lgrid_probs], fmt='%.6f') 134 | np.savetxt(f'{filename}_NH.txt', [logNH_grid, logNH_probs], fmt='%.6f') 135 | -------------------------------------------------------------------------------- /tests/test_table.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import fastxsf 3 | from fastxsf.model import Table, FixedTable 4 | import os 5 | import requests 6 | import matplotlib.pyplot as plt 7 | from numpy.testing import assert_allclose 8 | 9 | def get_fullfilename(filename, modeldir = os.environ.get('MODELDIR', '.')): 10 | return os.path.join(modeldir, filename) 11 | 12 | def download(url, filename): 13 | # download file if it does not exist 14 | fullfilename = get_fullfilename(filename) 15 | if not os.path.exists(fullfilename): 16 | print("downloading", url, "-->", fullfilename) 17 | response = requests.get(url) 18 | assert response.status_code == 200 19 | with open(filename, 'wb') as fout: 20 | fout.write(response.content) 21 | 22 | #download('https://zenodo.org/records/1169181/files/uxclumpy-cutoff.fits?download=1', 'uxclumpy-cutoff.fits') 23 | #download('https://zenodo.org/records/2235505/files/wada-cutoff.fits?download=1', 'wada-cutoff.fits') 24 | #download('https://zenodo.org/records/2235457/files/blob_uniform.fits?download=1', 'blob_uniform.fits') 25 | download('https://zenodo.org/records/2224651/files/wedge.fits?download=1', 'wedge.fits') 26 | download('https://zenodo.org/records/2224472/files/diskreflect.fits?download=1', 'diskreflect.fits') 27 | 28 | 29 | def test_disk_table(): 30 | energies = np.logspace(-0.5, 2, 1000) 31 | e_lo = energies[:-1] 32 | e_hi = energies[1:] 33 | e_mid = (e_lo + e_hi) / 2.0 34 | deltae = e_hi - e_lo 35 | 36 | fastxsf.x.abundance("angr") 37 | fastxsf.x.cross_section("vern") 38 | # compare diskreflect to pexmon 39 | atable = Table(get_fullfilename("diskreflect.fits")) 40 | Ecut = 400 41 | Incl = 70 42 | PhoIndex = 2.0 43 | ZHe = 1 44 | ZFe = 1 45 | for z in 0.0, 1.0, 2.0, 4.0: 46 | print(f"Case: redshift={z}") 47 | ftable = FixedTable(get_fullfilename("diskreflect.fits"), energies, redshift=z) 48 | A = atable(energies, [PhoIndex, Ecut, Incl, z]) 49 | B = fastxsf.x.pexmon(energies=energies, pars=[PhoIndex, Ecut, -1, z, ZHe, ZFe, Incl]) / (1 + z)**2 / 2 50 | C = ftable(energies=energies, pars=[PhoIndex, Ecut, Incl]) 51 | l, = plt.plot(e_mid, A / deltae / (1 + z)**2, label="atable") 52 | plt.plot(e_mid, B / deltae / (1 + z)**2, label="pexmon", ls=':', color=l.get_color()) 53 | plt.xlabel("Energy [keV]") 54 | plt.ylabel("Spectrum [photons/cm$^2$/s]") 55 | plt.yscale("log") 56 | plt.xscale("log") 57 | plt.legend() 58 | plt.savefig("pexmon.pdf") 59 | #plt.close() 60 | mask = np.logical_and(energies[:-1] > 8 / (1 + z), energies[:-1] < 80 / (1 + z)) 61 | assert_allclose(A[mask], B[mask], rtol=0.1) 62 | assert_allclose(A, C) 63 | 64 | def test_pexpl_table(): 65 | energies = np.logspace(-0.5, 2, 1000) 66 | e_lo = energies[:-1] 67 | e_hi = energies[1:] 68 | e_mid = (e_lo + e_hi) / 2.0 69 | deltae = e_hi - e_lo 70 | Ecut = 1000 71 | Incl = 70 72 | ZHe = 1 73 | ZFe = 1 74 | for PhoIndex in 2.4, 2.0, 1.2: 75 | for z in 0, 1, 2: 76 | A = fastxsf.x.zpowerlw(energies=energies, pars=[PhoIndex, z]) 77 | B = fastxsf.x.pexmon(energies=energies, pars=[PhoIndex, Ecut, 0, z, ZHe, ZFe, Incl]) / (1 + z)**2 78 | l, = plt.plot(e_mid * (1 + z), A / deltae, label="atable") 79 | plt.plot(e_mid * (1 + z), B / deltae / (1 + z)**(PhoIndex - 2), label="pexmon", ls=':', color=l.get_color()) 80 | plt.xlabel("Energy [keV]") 81 | plt.ylabel("Spectrum [photons/cm$^2$/s]") 82 | plt.yscale("log") 83 | plt.xscale("log") 84 | plt.legend() 85 | plt.savefig("pexmonpl.pdf") 86 | assert_allclose(A, B / (1 + z)**(PhoIndex - 2), rtol=0.2, atol=1e-4) 87 | 88 | def test_absorber_table(): 89 | fastxsf.x.abundance("angr") 90 | fastxsf.x.cross_section("bcmc") 91 | # compare uxclumpy to ztbabs * zpowerlw 92 | atable = Table(get_fullfilename("wedge.fits")) 93 | PhoIndex = 1.0 94 | Incl = 80 95 | 96 | for z in 0, 1, 2: 97 | plt.figure(figsize=(20, 5)) 98 | print("Redshift:", z) 99 | for elo, NH22 in (0.2, 0.01), (0.3, 0.1), (0.6, 0.4): 100 | #for elo, NH22 in (0.3, 0.01),: 101 | energies = np.geomspace(elo / (1 + z), 10, 400) 102 | e_lo = energies[:-1] 103 | e_hi = energies[1:] 104 | e_mid = (e_lo + e_hi) / 2.0 105 | deltae = e_hi - e_lo 106 | 107 | A = atable(energies, [NH22, PhoIndex, 45.6, Incl, z]) 108 | B = fastxsf.x.zpowerlw(energies=energies, pars=[PhoIndex, z]) 109 | C = B * fastxsf.x.zphabs(energies=energies, pars=[NH22, z]) 110 | mask = np.logical_and(energies[:-1] > elo / (1 + z), energies[:-1] / (1 + z) < 80) 111 | mask[np.abs(energies[:-1] - 6.4 / (1 + z)) < 0.1] = False 112 | plt.plot(e_mid, A / deltae, label="atable", ls='--', color='k', lw=0.5) 113 | A[~mask] = np.nan 114 | B[~mask] = np.nan 115 | C[~mask] = np.nan 116 | plt.plot(e_mid, A / deltae, label="atable", color='k') 117 | plt.plot(e_mid, B / deltae, label="pl", ls="--", color='orange', lw=0.5) 118 | plt.plot(e_mid, C / deltae, label="pl*tbabs", color='r', lw=1) 119 | plt.xlabel("Energy [keV]") 120 | plt.ylabel("Spectrum [photons/cm$^2$/s/keV]") 121 | plt.ylim(0.04, 3) 122 | plt.yscale("log") 123 | plt.xscale("log") 124 | plt.legend() 125 | plt.savefig("abspl_z%d.pdf" % z) 126 | print(energies[np.argmax(np.abs(np.log10(A[mask] / C[mask])))]) 127 | assert_allclose(A[mask], C[mask], rtol=0.2) 128 | plt.close() 129 | -------------------------------------------------------------------------------- /simplev.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import numpy as np 4 | import scipy.stats 5 | import ultranest 6 | from matplotlib import pyplot as plt 7 | 8 | import fastxsf 9 | 10 | # fastxsf.x.chatter(0) 11 | fastxsf.x.abundance('wilm') 12 | fastxsf.x.cross_section('vern') 13 | 14 | # load the spectrum, where we will consider data from 0.5 to 8 keV 15 | data = fastxsf.load_pha('example/179.pi', 0.5, 8) 16 | 17 | # fetch some basic information about our spectrum 18 | e_lo = data['e_lo'] 19 | e_hi = data['e_hi'] 20 | e_mid = (data['e_hi'] + data['e_lo']) / 2. 21 | e_width = data['e_hi'] - data['e_lo'] 22 | energies = np.append(e_lo, e_hi[-1]) 23 | RMF_src = data['RMF_src'] 24 | 25 | chan_e = (data['chan_e_min'] + data['chan_e_max']) / 2. 26 | 27 | # load a Table model 28 | absAGN = fastxsf.Table(os.path.join(os.environ.get('MODELDIR', '.'), 'uxclumpy-cutoff.fits')) 29 | 30 | # pre-compute the absorption factors -- no need to call this again and again if the parameters do not change! 31 | galabso = fastxsf.x.TBabs(energies=energies, pars=[data['galnh']]) 32 | 33 | 34 | # define a likelihood 35 | def loglikelihood(params, plot=False): 36 | norm, NH22, rel_scat_norm, PhoIndex, TORsigma, CTKcover, z, bkg_norm = np.transpose(params) 37 | Incl = 45 + norm * 0 38 | Ecut = 400 + norm * 0 39 | scat_norm = norm * rel_scat_norm 40 | 41 | abs_component = absAGN(energies=energies, pars=np.transpose([NH22, PhoIndex, Ecut, TORsigma, CTKcover, Incl, z]), vectorized=True) 42 | 43 | scat_component = fastxsf.xvec(fastxsf.x.zpowerlw, energies=energies, pars=np.transpose([norm, PhoIndex])) 44 | 45 | pred_spec = np.einsum('ij,i->ij', abs_component, norm) + np.einsum('ij,i->ij', scat_component, scat_norm) 46 | 47 | pred_counts_src_srcreg = RMF_src.apply_rmf_vectorized(np.einsum('i,i,ji->ji', data['ARF'], galabso, pred_spec))[:,data['chan_mask']] * data['src_expoarea'] 48 | pred_counts_bkg_srcreg = np.einsum('j,i->ij', data['bkg_model_src_region'], bkg_norm) * data['src_expoarea'] 49 | pred_counts_srcreg = pred_counts_src_srcreg + pred_counts_bkg_srcreg 50 | pred_counts_bkg_bkgreg = np.einsum('j,i->ij', data['bkg_model_bkg_region'], bkg_norm) * data['bkg_expoarea'] 51 | 52 | if plot: 53 | print(params[0]) 54 | plt.figure() 55 | plt.plot(data['chan_e_min'], data['src_region_counts'] / (data['chan_e_max'] - data['chan_e_min']), 'o', label='data', mfc='none') 56 | plt.plot(data['chan_e_min'], pred_counts_srcreg[0] / (data['chan_e_max'] - data['chan_e_min']), label='src+bkg') 57 | plt.plot(data['chan_e_min'], pred_counts_src_srcreg[0] / (data['chan_e_max'] - data['chan_e_min']), label='src') 58 | plt.plot(data['chan_e_min'], pred_counts_bkg_srcreg[0] / (data['chan_e_max'] - data['chan_e_min']), label='bkg') 59 | plt.xlabel('Channel Energy [keV]') 60 | plt.ylabel('Counts / keV') 61 | plt.legend() 62 | plt.savefig('src_region_counts.pdf') 63 | plt.close() 64 | 65 | plt.figure() 66 | plt.plot(data['chan_e_min'], data['bkg_region_counts'] / (data['chan_e_max'] - data['chan_e_min']), 'o', label='data', mfc='none') 67 | plt.plot(data['chan_e_min'], pred_counts_bkg_bkgreg[0] / (data['chan_e_max'] - data['chan_e_min']), label='bkg') 68 | plt.xlabel('Channel Energy [keV]') 69 | plt.ylabel('Counts / keV') 70 | plt.legend() 71 | plt.savefig('bkg_region_counts.pdf') 72 | plt.close() 73 | 74 | # compute log Poisson probability 75 | like_srcreg = fastxsf.logPoissonPDF_vectorized(pred_counts_srcreg, data['src_region_counts']) 76 | like_bkgreg = fastxsf.logPoissonPDF_vectorized(pred_counts_bkg_bkgreg, data['bkg_region_counts']) 77 | # combined the probabilities. If fitting multiple spectra, you would add them up here as well 78 | return like_srcreg + like_bkgreg 79 | 80 | 81 | # compute a likelihood: 82 | z = np.array([data['redshift']] * 2) 83 | bkg_norm = np.array([1.0] * 2) 84 | norm = np.array([3e-7] * 2) 85 | scat_norm = norm * 0.08 86 | PhoIndex = np.array([1.9, 2.0]) 87 | TORsigma = np.array([28.0] * 2) 88 | CTKcover = np.array([0.1] * 2) 89 | NH22 = np.array([1.0] * 2) 90 | print(loglikelihood(np.transpose([norm, NH22, np.array([0.08] * 2), PhoIndex, TORsigma, CTKcover, z, bkg_norm]), plot=True)) 91 | 92 | # lets define a prior 93 | 94 | PhoIndex_gauss = scipy.stats.norm(1.95, 0.15) 95 | # we are cool, we can let redshift be a free parameter informed from photo-z 96 | z_gauss = scipy.stats.norm(data['redshift'], 0.05) 97 | 98 | 99 | # define the prior transform function 100 | def prior_transform(cube): 101 | params = cube.copy() 102 | # uniform from 1e-10 to 1 103 | params[:,0] = 10**(cube[:,0] * 10 + -10) 104 | # uniform from 1e-2 to 1e2 105 | params[:,1] = 10**(cube[:,1] * (2 - -2) + -2) 106 | params[:,2] = 10**(cube[:,2] * (-1 - -5) + -5) 107 | # Gaussian prior 108 | params[:,3] = PhoIndex_gauss.ppf(cube[:,3]) 109 | # uniform priors 110 | params[:,4] = cube[:,4] * (80 - 7) + 7 111 | params[:,5] = cube[:,5] * (0.4) + 0 112 | # informative Gaussian prior on the redshift 113 | params[:,6] = z_gauss.ppf(cube[:,6]) 114 | # log-uniform prior on the background normalisation between 0.1 and 10 115 | params[:,7] = 10**(cube[:,7] * (1 - -1) + -1) 116 | return params 117 | 118 | 119 | # define parameter names 120 | param_names = ['norm', 'logNH22', 'scatnorm', 'PhoIndex', 'TORsigma', 'CTKcover', 'redshift', 'bkg_norm'] 121 | 122 | 123 | # run sampler 124 | sampler = ultranest.ReactiveNestedSampler( 125 | param_names, loglikelihood, prior_transform, 126 | log_dir='simplev', resume=True, 127 | vectorized=True) 128 | 129 | # then to run: 130 | # results = sampler.run(max_num_improvement_loops=0, frac_remain=0.5) 131 | 132 | # and to plot a model: 133 | # loglikelihood(results['samples'][:10,:], plot=True) 134 | 135 | # and to plot the posterior corner plot: 136 | # sampler.plot() 137 | -------------------------------------------------------------------------------- /fastxsf/flux.py: -------------------------------------------------------------------------------- 1 | """Functions for flux and luminosity computations.""" 2 | 3 | import numpy as np 4 | from astropy import units as u 5 | 6 | 7 | def frac_overlap_interval(edges, lo, hi): 8 | """Compute overlap of bins. 9 | 10 | Parameters 11 | ---------- 12 | edges: array 13 | edges of bins. 14 | lo: float 15 | lower limit 16 | hi: float 17 | upper limit 18 | 19 | Returns 20 | ------- 21 | weights: array 22 | for each bin, what fraction of the bin lies between lo and hi. 23 | """ 24 | weights = np.zeros(len(edges) - 1) 25 | for i, (edge_lo, edge_hi) in enumerate(zip(edges[:-1], edges[1:])): 26 | if edge_lo > lo and edge_hi < hi: 27 | weight = 1 28 | elif edge_hi < lo: 29 | weight = 0 30 | elif edge_lo > hi: 31 | weight = 0 32 | elif hi > edge_hi and lo > edge_lo: 33 | weight = (edge_hi - lo) / (edge_hi - edge_lo) 34 | elif lo < edge_lo and hi < edge_hi: 35 | weight = (hi - edge_lo) / (edge_hi - edge_lo) 36 | else: 37 | weight = (min(hi, edge_hi) - max(lo, edge_lo)) / (edge_hi - edge_lo) 38 | weights[i] = weight 39 | return weights 40 | 41 | 42 | def bins_sum(values, edges, lo, hi): 43 | """Sum up bin values. 44 | 45 | Parameters 46 | ---------- 47 | values: array 48 | values in bins 49 | edges: array 50 | bin edges 51 | lo: float 52 | lower limit 53 | hi: float 54 | upper limit 55 | 56 | Returns 57 | ------- 58 | sum: float 59 | values summed from bins between lo and hi. 60 | """ 61 | widths = edges[1:] - edges[:-1] 62 | fracs = frac_overlap_interval(edges, lo, hi) 63 | return np.sum(widths * values * fracs) 64 | 65 | 66 | def bins_integrate1(values, edges, lo, hi, axis=None): 67 | """Integrate up bin values. 68 | 69 | Parameters 70 | ---------- 71 | values: array 72 | values in bins 73 | edges: array 74 | bin edges 75 | lo: float 76 | lower limit 77 | hi: float 78 | upper limit 79 | axis: int | None 80 | summing axis 81 | 82 | Returns 83 | ------- 84 | I: float 85 | integral from bins between lo and hi of values times bin center. 86 | """ 87 | mids = (edges[1:] + edges[:-1]) / 2.0 88 | fracs = frac_overlap_interval(edges, lo, hi) 89 | return np.sum(mids * fracs * values, axis=axis) 90 | 91 | 92 | def bins_integrate(values, edges, lo, hi, axis=None): 93 | """Integrate up bin values. 94 | 95 | Parameters 96 | ---------- 97 | values: array 98 | values in bins 99 | edges: array 100 | bin edges 101 | lo: float 102 | lower limit 103 | hi: float 104 | upper limit 105 | axis: int | None 106 | summing axis 107 | 108 | Returns 109 | ------- 110 | I: float 111 | values integrated from bins between lo and hi. 112 | """ 113 | return np.sum(values * frac_overlap_interval(edges, lo, hi), axis=axis) 114 | 115 | 116 | def photon_flux(unfolded_model_spectrum, energies, energy_lo, energy_hi, axis=None): 117 | """Compute photon flux. 118 | 119 | Parameters 120 | ---------- 121 | unfolded_model_spectrum: array 122 | Model spectral density 123 | energies: array 124 | energies 125 | energy_lo: float 126 | lower limit 127 | energy_hi: float 128 | upper limit 129 | axis: int | None 130 | summing axis 131 | 132 | Returns 133 | ------- 134 | photon_flux: float 135 | Photon flux in phot/cm^2/s 136 | """ 137 | Nchan = len(energies) - 1 138 | assert unfolded_model_spectrum.shape == (Nchan,) 139 | assert energies.shape == (Nchan + 1,) 140 | integral = bins_integrate(unfolded_model_spectrum, energies, energy_lo, energy_hi, axis=axis) 141 | return integral / u.cm**2 / u.s 142 | 143 | 144 | def energy_flux(unfolded_model_spectrum, energies, energy_lo, energy_hi, axis=None): 145 | """Compute energy flux. 146 | 147 | Parameters 148 | ---------- 149 | unfolded_model_spectrum: array 150 | Model spectral density 151 | energies: array 152 | energies 153 | energy_lo: float 154 | lower limit 155 | energy_hi: float 156 | upper limit 157 | axis: int | None 158 | summing axis 159 | 160 | Returns 161 | ------- 162 | energy_flux: float 163 | Energy flux in erg/cm^2/s 164 | """ 165 | Nchan = len(energies) - 1 166 | assert unfolded_model_spectrum.shape[-1] == Nchan 167 | assert energies.shape == (Nchan + 1,) 168 | integral1 = bins_integrate1(unfolded_model_spectrum, energies, energy_lo, energy_hi, axis=axis) 169 | return integral1 * ((1 * u.keV).to(u.erg)) / u.cm**2 / u.s 170 | 171 | 172 | def luminosity( 173 | unfolded_model_spectrum, energies, rest_energy_lo, rest_energy_hi, z, cosmo 174 | ): 175 | """Compute luminosity. 176 | 177 | Parameters 178 | ---------- 179 | unfolded_model_spectrum: array 180 | Model spectral density 181 | energies: array 182 | energies 183 | rest_energy_lo: float 184 | lower limit 185 | rest_energy_hi: float 186 | upper limit 187 | z: float 188 | redshift 189 | cosmo: object 190 | astropy cosmology object 191 | 192 | Returns 193 | ------- 194 | luminosity: float 195 | Isotropic luminosity in erg/s 196 | """ 197 | Nchan = len(energies) - 1 198 | assert unfolded_model_spectrum.shape == (Nchan,) 199 | assert energies.shape == (Nchan + 1,) 200 | rest_energies = energies * (1 + z) 201 | rest_flux = energy_flux( 202 | unfolded_model_spectrum, rest_energies, rest_energy_lo, rest_energy_hi 203 | ) / (1 + z) 204 | DL = cosmo.luminosity_distance(z) 205 | return (rest_flux * (4 * np.pi * DL**2)).to(u.erg / u.s) 206 | -------------------------------------------------------------------------------- /fastxsf/data.py: -------------------------------------------------------------------------------- 1 | """Functionality for loading data.""" 2 | import os 3 | from functools import cache 4 | 5 | import astropy.io.fits as pyfits 6 | import numpy as np 7 | from astropy import units as u 8 | from astropy.cosmology import Planck18 as cosmo 9 | 10 | from .response import ARF, RMF, MockARF 11 | 12 | 13 | @cache 14 | def get_ARF(arf_filename): 15 | """Read ancillary response file. 16 | 17 | Avoids building a new object for the same file with caching. 18 | 19 | Parameters 20 | ---------- 21 | arf_filename: str 22 | filename 23 | 24 | Returns 25 | ------- 26 | ARF: object 27 | ARF object 28 | """ 29 | return ARF(arf_filename) 30 | 31 | 32 | @cache 33 | def get_RMF(rmf_filename): 34 | """Read response matrix file. 35 | 36 | Avoids building a new object for the same file with caching. 37 | 38 | Parameters 39 | ---------- 40 | rmf_filename: str 41 | filename 42 | 43 | Returns 44 | ------- 45 | RMF: object 46 | RMF object 47 | """ 48 | return RMF(rmf_filename) 49 | 50 | 51 | def load_pha(filename, elo, ehi, load_absorption=True, z=None, require_background=True): 52 | """Load PHA file. 53 | 54 | Parameters 55 | ---------- 56 | filename: str 57 | file name of the PHA-style spectrum 58 | elo: float 59 | lowest energy channel to consider 60 | ehi: float 61 | highest energy channel to consider 62 | load_absorption: bool 63 | whether to try to load the .nh file 64 | z: float or None 65 | if given, set data['redshift'] to z. 66 | Otherwise try to load the .z file. 67 | require_background: bool 68 | whether to fail if no corresponding background file is associated 69 | with the data. 70 | 71 | Returns 72 | ------- 73 | data: dict 74 | All information about the observation. 75 | """ 76 | path = os.path.dirname(filename) 77 | a = pyfits.open(filename) 78 | header = a["SPECTRUM"].header 79 | exposure = header["EXPOSURE"] 80 | backscal = header["BACKSCAL"] 81 | areascal = header["AREASCAL"] 82 | backfile = os.path.join(path, header["BACKFILE"]) 83 | rmffile = os.path.join(path, header["RESPFILE"]) 84 | arffile = os.path.join(path, header["ANCRFILE"]) 85 | channels = a["SPECTRUM"].data["CHANNEL"] 86 | 87 | b = pyfits.open(backfile) 88 | bheader = b["SPECTRUM"].header 89 | bexposure = bheader["EXPOSURE"] 90 | bbackscal = bheader["BACKSCAL"] 91 | bareascal = bheader["AREASCAL"] 92 | assert ( 93 | "RESPFILE" not in bheader or bheader["RESPFILE"] == header["RESPFILE"] 94 | ), "background must have same RMF" 95 | assert ( 96 | "ANCRFILE" not in bheader or bheader["ANCRFILE"] == header["ANCRFILE"] 97 | ), "background must have same ARF" 98 | 99 | ebounds = pyfits.getdata(rmffile, "EBOUNDS") 100 | chan_e_min = ebounds["E_MIN"] 101 | chan_e_max = ebounds["E_MAX"] 102 | mask = np.logical_and(chan_e_min > elo, chan_e_max < ehi) 103 | armf = get_RMF(rmffile) 104 | m = armf.get_dense_matrix() 105 | Nflux, Nchan = m.shape 106 | 107 | assert (Nflux,) == armf.energ_lo.shape == armf.energ_hi.shape 108 | 109 | if header["ANCRFILE"].strip() == 'NONE': 110 | aarf = MockARF(armf) 111 | else: 112 | aarf = get_ARF(arffile) 113 | 114 | assert (Nflux,) == aarf.e_low.shape == aarf.e_high.shape 115 | assert len(channels) == Nchan, (len(channels), Nchan) 116 | 117 | strip_mask = armf.strip(mask) 118 | aarf.strip(strip_mask) 119 | assert armf.energ_lo.shape == armf.energ_hi.shape == aarf.e_low.shape == aarf.e_high.shape 120 | 121 | # assert np.allclose(channels, np.arange(Nchan)+1), (channels, Nchan) 122 | fcounts = a["SPECTRUM"].data["COUNTS"] 123 | assert (Nchan,) == fcounts.shape, (fcounts.shape, Nchan) 124 | counts = fcounts.astype(int) 125 | assert (counts == fcounts).all() 126 | 127 | bchannels = b["SPECTRUM"].data["CHANNEL"] 128 | # assert np.allclose(bchannels, np.arange(Nchan)+1), (bchannels, Nchan) 129 | assert len(bchannels) == Nchan, (len(bchannels), Nchan) 130 | bfcounts = b["SPECTRUM"].data["COUNTS"] 131 | assert (Nchan,) == bfcounts.shape, (bfcounts.shape, Nchan) 132 | bcounts = bfcounts.astype(int) 133 | assert (bcounts == bfcounts).all() 134 | 135 | data = dict( 136 | Nflux=Nflux, 137 | Nchan=mask.sum(), 138 | src_region_counts=counts[mask], 139 | bkg_region_counts=bcounts[mask], 140 | chan_mask=mask, 141 | RMF_src=armf, 142 | RMF=np.array(m[:, mask][strip_mask,:]), 143 | ARF=np.array(aarf.specresp, dtype=float), 144 | e_lo=np.array(aarf.e_low), 145 | e_hi=np.array(aarf.e_high), 146 | e_delta=np.array(aarf.e_high - aarf.e_low), 147 | energies=np.append(aarf.e_low, aarf.e_high[-1]), 148 | chan_e_min=chan_e_min[mask], 149 | chan_e_max=chan_e_max[mask], 150 | src_expo=exposure, 151 | bkg_expo=bexposure, 152 | src_expoarea=exposure * areascal, 153 | bkg_expoarea=bexposure * bareascal, 154 | src_to_bkg_ratio=areascal / bareascal * backscal / bbackscal * exposure / bexposure, 155 | ) 156 | data["chan_const_spec_weighting"] = data["RMF_src"].apply_rmf( 157 | data["ARF"] * data['e_lo']**-1.4)[data['chan_mask']] * data['src_expoarea'] 158 | 159 | # data["chan_const_spec_weighting"] = np.dot(data["ARF"], data["RMF"]) * data["src_expoarea"] 160 | 161 | if os.path.exists(backfile + "_model.fits"): 162 | bkg_model = pyfits.getdata(backfile + "_model.fits", "SPECTRA") 163 | data["bkg_model_bkg_region"] = np.array(bkg_model[0]["INTPSPEC"][mask]) 164 | data["bkg_model_src_region"] = np.array(bkg_model[1]["INTPSPEC"][mask]) 165 | if os.path.exists(filename + ".nh") and load_absorption: 166 | data["galnh"] = float(np.loadtxt(filename + ".nh") / 1e22) 167 | if z is not None: 168 | data["redshift"] = z 169 | data["e_mid_restframe"] = (data["e_hi"] + data["e_lo"]) / 2 * (1 + z) 170 | data["luminosity_distance"] = cosmo.luminosity_distance(z) 171 | elif os.path.exists(filename + ".z"): 172 | z = float(np.loadtxt(filename + ".z")) 173 | data["redshift"] = z 174 | data["e_mid_restframe"] = (data["e_hi"] + data["e_lo"]) / 2 * (1 + z) 175 | data["luminosity_distance"] = cosmo.luminosity_distance(z).to(u.cm) 176 | else: 177 | z = 0.0 178 | 179 | return data 180 | -------------------------------------------------------------------------------- /simpleopt.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | A new approach to X-ray spectral fitting with Xspec models and 4 | optimized nested sampling. 5 | 6 | This script explains how to set up your model. 7 | 8 | The idea is that you create a function which computes all model components. 9 | 10 | """ 11 | import os 12 | import corner 13 | import numpy as np 14 | import scipy.stats 15 | from matplotlib import pyplot as plt 16 | from optns.sampler import OptNS 17 | 18 | import fastxsf 19 | 20 | # fastxsf.x.chatter(0) 21 | fastxsf.x.abundance('wilm') 22 | fastxsf.x.cross_section('vern') 23 | 24 | # load the spectrum, where we will consider data from 0.5 to 8 keV 25 | data = fastxsf.load_pha('example/179.pi', 0.5, 8) 26 | 27 | # fetch some basic information about our spectrum 28 | e_lo = data['e_lo'] 29 | e_hi = data['e_hi'] 30 | e_mid = (data['e_hi'] + data['e_lo']) / 2. 31 | e_width = data['e_hi'] - data['e_lo'] 32 | energies = np.append(e_lo, e_hi[-1]) 33 | RMF_src = data['RMF_src'] 34 | 35 | chan_e = (data['chan_e_min'] + data['chan_e_max']) / 2. 36 | 37 | # load a Table model 38 | absAGN = fastxsf.Table(os.path.join(os.environ.get('MODELDIR', '.'), 'uxclumpy-cutoff.fits')) 39 | 40 | # pre-compute the absorption factors -- no need to call this again and again if the parameters do not change! 41 | galabso = fastxsf.x.TBabs(energies=energies, pars=[data['galnh']]) 42 | 43 | # z = data['redshift'] 44 | # define the model parameters: 45 | #bkg_norm = 1.0 46 | #norm = 3e-7 47 | #scat_norm = norm * 0.08 48 | #PhoIndex = 2.0 49 | #TORsigma = 28.0 50 | #CTKcover = 0.1 51 | Incl = 45.0 52 | #NH22 = 1.0 53 | Ecut = 400 54 | 55 | Nsrc_chan = len(data['src_region_counts']) 56 | Nbkg_chan = len(data['bkg_region_counts']) 57 | counts_flat = np.hstack((data['src_region_counts'], data['bkg_region_counts'])) 58 | 59 | # lets now start using optimized nested sampling. 60 | 61 | # set up function which computes the various model components: 62 | # the parameters are: 63 | nonlinear_param_names = ['logNH', 'PhoIndex', 'TORsigma', 'CTKcover', 'redshift'] 64 | 65 | 66 | def compute_model_components(params): 67 | logNH, PhoIndex, TORsigma, CTKcover, z = params 68 | 69 | # first component: a absorbed power law 70 | NH22 = 10**(logNH - 22) 71 | abs_component = absAGN(energies=energies, pars=[NH22, PhoIndex, Ecut, TORsigma, CTKcover, Incl, z]) 72 | 73 | # second component, a copy of the unabsorbed power law 74 | scat = fastxsf.x.zpowerlw(energies=energies, pars=[PhoIndex, z]) 75 | 76 | # lets compute the detected counts: 77 | pred_counts = np.zeros((3, Nsrc_chan + Nbkg_chan)) 78 | # the two folded source components in the source region: 79 | pred_counts[0, :Nsrc_chan] = RMF_src.apply_rmf(data['ARF'] * galabso * abs_component)[data['chan_mask']] * data['src_expoarea'] 80 | pred_counts[1, :Nsrc_chan] = RMF_src.apply_rmf(data['ARF'] * galabso * scat)[data['chan_mask']] * data['src_expoarea'] 81 | # the unfolded background components in the source region 82 | pred_counts[2, :Nsrc_chan] = data['bkg_model_src_region'] * data['src_expoarea'] 83 | # the unfolded background components in the background region 84 | pred_counts[2, Nsrc_chan:] = data['bkg_model_bkg_region'] * data['bkg_expoarea'] 85 | # notice how the first three components do not affect the background: 86 | # pred_counts[0:3, Nsrc_chan:] = 0 # they remain zero 87 | assert (pred_counts[0] > 0).any(), (params, abs_component) 88 | assert (pred_counts[1] > 0).any(), (params, scat) 89 | assert (pred_counts[2] > 0).any(), (params,) 90 | 91 | return pred_counts.T 92 | 93 | 94 | # set up a prior transform for these nonlinear parameters 95 | PhoIndex_gauss = scipy.stats.norm(1.95, 0.15) 96 | # we are cool, we can let redshift be a free parameter informed from photo-z 97 | z_gauss = scipy.stats.norm(data['redshift'], 0.05) 98 | 99 | 100 | def nonlinear_param_transform(cube): 101 | params = cube.copy() 102 | params[0] = cube[0] * 4 + 20 # logNH 103 | params[1] = PhoIndex_gauss.ppf(cube[1]) 104 | params[2] = cube[2] * (80 - 7) + 7 105 | params[3] = cube[3] * (0.4) + 0 106 | # informative Gaussian prior on the redshift 107 | params[4] = z_gauss.ppf(cube[4]) 108 | return params 109 | 110 | 111 | # now for the linear (normalisation) parameters: 112 | linear_param_names = ['Nsrc', 'Nscat', 'Nbkg'] 113 | 114 | # set up a prior log-probability density function for these linear parameters: 115 | 116 | 117 | def linear_param_logprior(params): 118 | assert np.all(params > 0) 119 | Nsrc, Nscat, Nbkg = params.transpose() 120 | # a log-uniform prior on the source luminosity 121 | logp = -np.log(Nsrc) 122 | # a log-uniform prior on the relative scattering normalisation. 123 | # logp = -np.log(Nscat / Nsrc) 124 | assert np.isfinite(logp).all(), logp 125 | # limits: 126 | logp[Nscat > 0.1 * Nsrc] = -np.inf 127 | return logp 128 | 129 | 130 | # create OptNS object, and give it all of these ingredients, 131 | # as well as our data 132 | statmodel = OptNS( 133 | linear_param_names, nonlinear_param_names, compute_model_components, 134 | nonlinear_param_transform, linear_param_logprior, 135 | counts_flat, positive=True) 136 | 137 | # prior predictive checks: 138 | fig = plt.figure(figsize=(15, 4)) 139 | statmodel.prior_predictive_check_plot(fig.gca()) 140 | plt.legend() 141 | plt.ylim(0.1, counts_flat.max() * 1.1) 142 | plt.yscale('log') 143 | plt.savefig('simpleopt-ppc.pdf') 144 | plt.close() 145 | 146 | # create a UltraNest sampler from this. You can pass additional arguments like here: 147 | optsampler = statmodel.ReactiveNestedSampler( 148 | log_dir='simpleopt', resume=True) 149 | # run the UltraNest optimized sampler on the nonlinear parameter space: 150 | optresults = optsampler.run(max_num_improvement_loops=0, frac_remain=0.5) 151 | optsampler.print_results() 152 | optsampler.plot() 153 | 154 | # now for postprocessing the results, we want to get the full posterior: 155 | # this samples up to 1000 normalisations for each nonlinear posterior sample: 156 | fullsamples, weights, y_preds = statmodel.get_weighted_samples(optresults['samples'][:400], 100) 157 | print(f'Obtained {len(fullsamples)} weighted posterior samples') 158 | 159 | print('weights:', weights, np.nanmin(weights), np.nanmax(weights), np.mean(weights)) 160 | # make a corner plot: 161 | mask = weights > 1e-6 * np.nanmax(weights) 162 | fullsamples_selected = fullsamples[mask,:] 163 | fullsamples_selected[:, :len(linear_param_names)] = np.log10(fullsamples_selected[:, :len(linear_param_names)]) 164 | 165 | print(f'Obtained {mask.sum()} with not minuscule weight.') 166 | fig = corner.corner( 167 | fullsamples_selected, weights=weights[mask], 168 | labels=linear_param_names + nonlinear_param_names, 169 | show_titles=True, quiet=True, 170 | plot_datapoints=False, plot_density=False, 171 | levels=[0.9973, 0.9545, 0.6827, 0.3934], quantiles=[0.15866, 0.5, 0.8413], 172 | contour_kwargs=dict(linestyles=['-','-.',':','--'], colors=['navy','navy','navy','purple']), 173 | color='purple' 174 | ) 175 | plt.savefig('simpleopt-corner.pdf') 176 | plt.close() 177 | 178 | # to obtain equally weighted samples, we resample 179 | # this respects the effective sample size. If you get too few samples here, 180 | # crank up the number just above. 181 | samples, y_pred_samples = statmodel.resample(fullsamples, weights, y_preds) 182 | print(f'Obtained {len(samples)} equally weighted posterior samples') 183 | 184 | 185 | # prior predictive checks: 186 | fig = plt.figure(figsize=(15, 10)) 187 | statmodel.posterior_predictive_check_plot(fig.gca(), samples[:100]) 188 | plt.legend() 189 | plt.ylim(0.1, counts_flat.max() * 1.1) 190 | plt.yscale('log') 191 | plt.savefig('simpleopt-postpc.pdf') 192 | plt.close() 193 | 194 | -------------------------------------------------------------------------------- /reflopt.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | A new approach to X-ray spectral fitting with Xspec models and 4 | optimized nested sampling. 5 | 6 | This script explains how to set up your model. 7 | 8 | The idea is that you create a function which computes all model components. 9 | 10 | """ 11 | 12 | import corner 13 | import numpy as np 14 | import scipy.stats 15 | from matplotlib import pyplot as plt 16 | from optns.sampler import OptNS 17 | 18 | import fastxsf 19 | 20 | # fastxsf.x.chatter(0) 21 | fastxsf.x.abundance('wilm') 22 | fastxsf.x.cross_section('vern') 23 | 24 | # load the spectrum, where we will consider data from 0.5 to 8 keV 25 | data = fastxsf.load_pha('example/179.pi', 0.5, 8) 26 | 27 | # fetch some basic information about our spectrum 28 | e_lo = data['e_lo'] 29 | e_hi = data['e_hi'] 30 | e_mid = (data['e_hi'] + data['e_lo']) / 2. 31 | e_width = data['e_hi'] - data['e_lo'] 32 | energies = np.append(e_lo, e_hi[-1]) 33 | RMF_src = data['RMF_src'] 34 | 35 | chan_e = (data['chan_e_min'] + data['chan_e_max']) / 2. 36 | 37 | # pre-compute the absorption factors -- no need to call this again and again if the parameters do not change! 38 | galabso = fastxsf.x.TBabs(energies=energies, pars=[data['galnh']]) 39 | 40 | z = data['redshift'] 41 | # define the model parameters: 42 | bkg_norm = 1.0 43 | norm = 3e-7 44 | scat_norm = norm * 0.08 45 | PhoIndex = 2.0 46 | TORsigma = 28.0 47 | CTKcover = 0.1 48 | Incl = 45.0 49 | NH22 = 1.0 50 | Ecut = 400 51 | 52 | Nsrc_chan = len(data['src_region_counts']) 53 | Nbkg_chan = len(data['bkg_region_counts']) 54 | counts_flat = np.hstack((data['src_region_counts'], data['bkg_region_counts'])) 55 | 56 | # lets now start using optimized nested sampling. 57 | 58 | # set up function which computes the various model components: 59 | # the parameters are: 60 | nonlinear_param_names = ['logNH', 'PhoIndex', 'emissPL', 'Rin', 'Rout', 'incl'] 61 | 62 | 63 | def compute_model_components(params): 64 | logNH, PhoIndex, emissivityPL, Rin, Rout, incl = params 65 | 66 | # first component: a absorbed power law 67 | 68 | pl = fastxsf.x.zpowerlw(energies=energies, pars=[PhoIndex, z]) 69 | abso = fastxsf.x.zTBabs(energies=energies, pars=[10**(logNH - 22), z]) 70 | plabso = pl * abso 71 | 72 | # second component, a disk reflection 73 | Eline = 6.4 # keV 74 | refl = fastxsf.x.diskline(energies=energies, pars=[Eline, emissivityPL, Rin, Rout, incl]) 75 | 76 | # third component, a copy of the unabsorbed power law 77 | scat = pl 78 | assert (pl >= 0).all() 79 | assert (plabso >= 0).all() 80 | assert (refl >= 0).all() 81 | 82 | # now we need to project all of our components through the response. 83 | src_components = data['ARF'] * galabso * np.array([plabso, refl, scat]) 84 | pred_counts_src_srcreg = RMF_src.apply_rmf_vectorized(src_components)[:,data['chan_mask']] * data['src_expoarea'] 85 | # add non-folded background to source region components 86 | pred_counts = np.zeros((4, Nsrc_chan + Nbkg_chan)) 87 | # the three folded source components in the source region 88 | pred_counts[0:3, :Nsrc_chan] = pred_counts_src_srcreg 89 | # the unfolded background components in the source region 90 | pred_counts[3, :Nsrc_chan] = data['bkg_model_src_region'] * data['src_expoarea'] 91 | # the unfolded background components in the background region 92 | pred_counts[3, Nsrc_chan:] = data['bkg_model_bkg_region'] * data['bkg_expoarea'] 93 | # notice how the first three components do not affect the background: 94 | # pred_counts[0:3, Nsrc_chan:] = 0 # they remain zero 95 | assert (pred_counts[0] > 0).any(), (params, pl, abso) 96 | assert (pred_counts[1] > 0).any(), (params, refl) 97 | assert (pred_counts[2] > 0).any(), (params, pl) 98 | assert (pred_counts[3] > 0).any(), (params,) 99 | 100 | return pred_counts.T 101 | 102 | 103 | # set up a prior transform for these nonlinear parameters 104 | PhoIndex_gauss = scipy.stats.norm(1.95, 0.15) 105 | 106 | 107 | def nonlinear_param_transform(cube): 108 | params = cube.copy() 109 | params[0] = cube[0] * 4 + 20 # logNH 110 | params[1] = PhoIndex_gauss.ppf(cube[1]) 111 | params[2] = -(cube[2] * 2 + 1) # emissivity index from -3 to -1 112 | params[3] = cube[3] * 14 + 6 # Rin 113 | params[4] = 10**(cube[4] * 1.7 + 1.7) # Rout from 50 to 3000 114 | params[5] = np.arccos(cube[5]) * 180 / np.pi # inclination from 0 to 90 degrees 115 | return params 116 | 117 | 118 | # now for the linear (normalisation) parameters: 119 | linear_param_names = ['Nsrc', 'Nrefl', 'Nscat', 'Nbkg'] 120 | # set up a prior log-probability density function for these linear parameters: 121 | 122 | 123 | def linear_param_logprior(params): 124 | assert np.all(params > 0) 125 | Nsrc, Nrefl, Nscat, Nbkg = params.transpose() 126 | # a log-uniform prior on the source luminosity 127 | logp = -np.log(Nsrc) 128 | # a log-uniform prior on the relative scattering normalisation. 129 | # logp += -np.log(Nscat / Nsrc) 130 | # a log-uniform prior on the relative reflection normalisation. 131 | # logp += -np.log(Nrefl / Nsrc) 132 | assert np.isfinite(logp).all(), logp 133 | # limits: 134 | logp[Nscat > 0.1 * Nsrc] = -np.inf 135 | logp[Nrefl > 50 * Nsrc] = -np.inf 136 | logp[Nrefl < Nsrc / 300] = -np.inf 137 | return logp 138 | 139 | 140 | # create OptNS object, and give it all of these ingredients, 141 | # as well as our data 142 | statmodel = OptNS( 143 | linear_param_names, nonlinear_param_names, compute_model_components, 144 | nonlinear_param_transform, linear_param_logprior, 145 | counts_flat, positive=True) 146 | 147 | # prior predictive checks: 148 | fig = plt.figure(figsize=(15, 4)) 149 | statmodel.prior_predictive_check_plot(fig.gca()) 150 | plt.legend() 151 | plt.ylim(0.1, counts_flat.max() * 1.1) 152 | plt.yscale('log') 153 | plt.savefig('optrefl-ppc.pdf') 154 | plt.close() 155 | 156 | # create a UltraNest sampler from this. You can pass additional arguments like here: 157 | optsampler = statmodel.ReactiveNestedSampler( 158 | log_dir='optrefl', resume=True) 159 | # run the UltraNest optimized sampler on the nonlinear parameter space: 160 | optresults = optsampler.run(max_num_improvement_loops=0, frac_remain=0.5) 161 | optsampler.print_results() 162 | optsampler.plot() 163 | 164 | # now for postprocessing the results, we want to get the full posterior: 165 | # this samples up to 1000 normalisations for each nonlinear posterior sample: 166 | fullsamples, weights, y_preds = statmodel.get_weighted_samples(optresults['samples'][:400], 100) 167 | print(f'Obtained {len(fullsamples)} weighted posterior samples') 168 | 169 | print('weights:', weights, np.nanmin(weights), np.nanmax(weights), np.mean(weights)) 170 | # make a corner plot: 171 | mask = weights > 1e-6 * np.nanmax(weights) 172 | fullsamples_selected = fullsamples[mask,:] 173 | fullsamples_selected[:, :len(linear_param_names)] = np.log10(fullsamples_selected[:, :len(linear_param_names)]) 174 | 175 | print(f'Obtained {mask.sum()} with not minuscule weight.') 176 | fig = corner.corner( 177 | fullsamples_selected, weights=weights[mask], 178 | labels=linear_param_names + nonlinear_param_names, 179 | show_titles=True, quiet=True, 180 | plot_datapoints=False, plot_density=False, 181 | levels=[0.9973, 0.9545, 0.6827, 0.3934], quantiles=[0.15866, 0.5, 0.8413], 182 | contour_kwargs=dict(linestyles=['-','-.',':','--'], colors=['navy','navy','navy','purple']), 183 | color='purple' 184 | ) 185 | plt.savefig('optrefl-corner.pdf') 186 | plt.close() 187 | 188 | # to obtain equally weighted samples, we resample 189 | # this respects the effective sample size. If you get too few samples here, 190 | # crank up the number just above. 191 | samples, y_pred_samples = statmodel.resample(fullsamples, weights, y_preds) 192 | print(f'Obtained {len(samples)} equally weighted posterior samples') 193 | 194 | 195 | # prior predictive checks: 196 | fig = plt.figure(figsize=(15, 10)) 197 | statmodel.posterior_predictive_check_plot(fig.gca(), samples[:100]) 198 | plt.legend() 199 | plt.ylim(0.1, counts_flat.max() * 1.1) 200 | plt.yscale('log') 201 | plt.savefig('optrefl-postpc.pdf') 202 | plt.close() 203 | 204 | -------------------------------------------------------------------------------- /scripts/xagnfitter.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import numpy as np 4 | import scipy.stats 5 | import ultranest 6 | from matplotlib import pyplot as plt 7 | 8 | import fastxsf 9 | import fastxsf.flux 10 | import sys 11 | import astropy.units as u 12 | 13 | # fastxsf.x.chatter(0) 14 | fastxsf.x.abundance('wilm') 15 | fastxsf.x.cross_section('vern') 16 | 17 | # load the spectrum, where we will consider data from 0.5 to 8 keV 18 | data = fastxsf.load_pha(sys.argv[1], float(os.environ['ELO']), float(os.environ['EHI'])) 19 | 20 | # fetch some basic information about our spectrum 21 | e_lo = data['e_lo'] 22 | e_hi = data['e_hi'] 23 | e_mid = (data['e_hi'] + data['e_lo']) / 2. 24 | e_width = data['e_hi'] - data['e_lo'] 25 | energies = np.append(e_lo, e_hi[-1]) 26 | RMF_src = data['RMF_src'] 27 | 28 | chan_e = (data['chan_e_min'] + data['chan_e_max']) / 2. 29 | 30 | # load a Table model 31 | absAGN = fastxsf.Table(os.path.join(os.environ.get('MODELDIR', '.'), 'uxclumpy-cutoff.fits')) 32 | 33 | # pre-compute the absorption factors -- no need to call this again and again if the parameters do not change! 34 | galabso = fastxsf.x.TBabs(energies=energies, pars=[data['galnh']]) 35 | 36 | def setup_model(params): 37 | lognorm, logNH22, logrel_scat_norm, PhoIndex, TORsigma, CTKcover, z, bkg_norm = np.transpose(params) 38 | norm, NH22, rel_scat_norm = 10**lognorm, 10**logNH22, 10**logrel_scat_norm 39 | Incl = 45 + norm * 0 40 | Ecut = 400 + norm * 0 41 | scat_norm = norm * rel_scat_norm 42 | 43 | abs_component = np.einsum('ij,i->ij', 44 | absAGN(energies=energies, pars=np.transpose([NH22, PhoIndex, Ecut, TORsigma, CTKcover, Incl, z]), vectorized=True), 45 | norm) 46 | 47 | scat_component = np.einsum('ij,i->ij', 48 | fastxsf.xvec(fastxsf.x.zpowerlw, energies=energies, pars=np.transpose([PhoIndex, z])), 49 | scat_norm) 50 | 51 | return abs_component, scat_component 52 | 53 | def get_flux(params, e_min, e_max): 54 | abs_component, scat_component = setup_model(params) 55 | return fastxsf.flux.energy_flux(galabso * (abs_component + scat_component), energies, e_min, e_max, axis=1) 56 | 57 | 58 | # define a likelihood 59 | def loglikelihood(params, plot=False, plot_prefix=''): 60 | norm, NH22, rel_scat_norm, PhoIndex, TORsigma, CTKcover, z, bkg_norm = np.transpose(params) 61 | abs_component, scat_component = setup_model(params) 62 | pred_spec = abs_component + scat_component 63 | 64 | pred_counts_src_srcreg = RMF_src.apply_rmf_vectorized(np.einsum('i,i,ji->ji', data['ARF'], galabso, pred_spec))[:,data['chan_mask']] * data['src_expoarea'] 65 | pred_counts_bkg_srcreg = np.einsum('j,i->ij', data['bkg_model_src_region'], bkg_norm) * data['src_expoarea'] 66 | pred_counts_srcreg = pred_counts_src_srcreg + pred_counts_bkg_srcreg 67 | pred_counts_bkg_bkgreg = np.einsum('j,i->ij', data['bkg_model_bkg_region'], bkg_norm) * data['bkg_expoarea'] 68 | 69 | if plot: 70 | plt.figure() 71 | l = plt.plot(e_mid, np.einsum('i,ji->ij', e_mid**2 * galabso, pred_spec), ':', color='darkblue', label='src*galabso', mfc='none', alpha=0.3)[0] 72 | l2 = plt.plot(e_mid, (e_mid**2 * pred_spec).T, '-', label='src', mfc='none', color='k', alpha=0.3)[0] 73 | l3 = plt.plot(e_mid, (e_mid**2 * abs_component).T, '-', label='torus', mfc='none', color='green', alpha=0.3)[0] 74 | l4 = plt.plot(e_mid, (e_mid**2 * scat_component).T, '--', label='scat', mfc='none', color='lightblue', alpha=0.3)[0] 75 | plt.yscale('log') 76 | plt.xscale('log') 77 | plt.ylim((e_mid**2 * pred_spec).min() / 5, (e_mid**2 * pred_spec).max() * 2) 78 | plt.xlabel('Energy [keV]') 79 | plt.ylabel('Energy Flux Density * keV$^2$ [keV$^2$ * erg/s/cm$^2$/keV]') 80 | plt.legend([l, l2, l3, l4], ['src*galabso', 'src', 'torus', 'scat']) 81 | plt.savefig(f'{plot_prefix}src_model_E2.pdf') 82 | plt.close() 83 | 84 | plt.figure() 85 | l = plt.plot(e_mid, np.einsum('i,ji->ij', galabso, pred_spec), ':', color='darkblue', label='src*galabso', mfc='none', alpha=0.3)[0] 86 | l2 = plt.plot(e_mid, pred_spec.T, '-', label='src', mfc='none', color='k', alpha=0.3)[0] 87 | l3 = plt.plot(e_mid, abs_component.T, '-', label='torus', mfc='none', color='green', alpha=0.3)[0] 88 | l4 = plt.plot(e_mid, scat_component.T, '--', label='scat', mfc='none', color='lightblue', alpha=0.3)[0] 89 | plt.yscale('log') 90 | plt.xscale('log') 91 | plt.ylim(pred_spec.min() / 5, pred_spec.max() * 2) 92 | plt.xlabel('Energy [keV]') 93 | plt.ylabel('Energy Flux Density * keV$^2$ [keV$^2$ * erg/s/cm$^2$/keV]') 94 | plt.legend([l, l2, l3, l4], ['src*galabso', 'src', 'torus', 'scat']) 95 | plt.savefig(f'{plot_prefix}src_model.pdf') 96 | plt.close() 97 | 98 | plt.figure() 99 | plt.plot(data['chan_e_min'], data['src_region_counts'] / (data['chan_e_max'] - data['chan_e_min']), 'o', label='data', mfc='none') 100 | plt.plot(data['chan_e_min'], pred_counts_srcreg[0] / (data['chan_e_max'] - data['chan_e_min']), label='src+bkg') 101 | plt.plot(data['chan_e_min'], pred_counts_src_srcreg[0] / (data['chan_e_max'] - data['chan_e_min']), label='src') 102 | plt.plot(data['chan_e_min'], pred_counts_bkg_srcreg[0] / (data['chan_e_max'] - data['chan_e_min']), label='bkg') 103 | plt.xlabel('Channel Energy [keV]') 104 | plt.ylabel('Counts / keV') 105 | plt.legend() 106 | plt.savefig(f'{plot_prefix}src_region_counts.pdf') 107 | plt.close() 108 | 109 | plt.figure() 110 | plt.plot(data['chan_e_min'], data['bkg_region_counts'] / (data['chan_e_max'] - data['chan_e_min']), 'o', label='data', mfc='none') 111 | plt.plot(data['chan_e_min'], pred_counts_bkg_bkgreg[0] / (data['chan_e_max'] - data['chan_e_min']), label='bkg') 112 | plt.xlabel('Channel Energy [keV]') 113 | plt.ylabel('Counts / keV') 114 | plt.legend() 115 | plt.savefig(f'{plot_prefix}bkg_region_counts.pdf') 116 | plt.close() 117 | 118 | # compute log Poisson probability 119 | like_srcreg = fastxsf.logPoissonPDF_vectorized(pred_counts_srcreg, data['src_region_counts']) 120 | like_bkgreg = fastxsf.logPoissonPDF_vectorized(pred_counts_bkg_bkgreg, data['bkg_region_counts']) 121 | # combined the probabilities. If fitting multiple spectra, you would add them up here as well 122 | return like_srcreg + like_bkgreg 123 | 124 | def main(): 125 | # compute a likelihood: 126 | z = np.array([data['redshift'], 0, 0.404]) 127 | bkg_norm = np.array([1.0] * 3) 128 | lognorm = np.log10([1] * 2 + [8.45e-5]) 129 | logscat_norm = np.log10([0.08] * 3) 130 | PhoIndex = np.array([1.9, 2.0, 2.0]) 131 | TORsigma = np.array([28.0] * 3) 132 | CTKcover = np.array([0.1] * 3) 133 | logNH22 = np.log10([0.01] * 3) 134 | print(loglikelihood(np.transpose([lognorm, logNH22, logscat_norm, PhoIndex, TORsigma, CTKcover, z, bkg_norm]), plot=True)) 135 | test_fluxes = get_flux(np.transpose([lognorm, logNH22, logscat_norm, PhoIndex, TORsigma, CTKcover, z, bkg_norm]), 0.5, 2) 136 | print('flux:', test_fluxes) 137 | print('ARF', np.median(data['ARF'])) 138 | print('src_expoarea:', data['src_expoarea']) 139 | print('bkg_expoarea:', data['bkg_expoarea']) 140 | 141 | ratio_from_models = (data["bkg_model_src_region"].sum() * data["src_expoarea"]) / \ 142 | (data["bkg_model_bkg_region"].sum() * data["bkg_expoarea"]) 143 | print("bkg source/background count ratio from model:", ratio_from_models) 144 | print("expected ratio (header-based):", data["src_to_bkg_ratio"]) 145 | assert np.log10(test_fluxes[1] / (2.2211e-09 * u.erg/u.s/u.cm**2)) < 0.2 146 | assert np.log10(test_fluxes[2] / (9.5175e-14 * u.erg/u.s/u.cm**2)) < 0.2 147 | 148 | # lets define a prior 149 | 150 | PhoIndex_gauss = scipy.stats.norm(1.95, 0.15) 151 | # we are cool, we can let redshift be a free parameter informed from photo-z 152 | if 'redshift' in data: 153 | z_gauss = scipy.stats.norm(data['redshift'], 0.001) 154 | else: 155 | z_gauss = scipy.stats.uniform(0, 6) 156 | 157 | 158 | # define the prior transform function 159 | def prior_transform(cube): 160 | params = cube.copy() 161 | # uniform from 1e-10 to 1 162 | params[:,0] = (cube[:,0] * 10 + -10) 163 | # uniform from 1e-2 to 1e2 164 | params[:,1] = (cube[:,1] * (4 - -2) + -2) 165 | params[:,2] = (cube[:,2] * (-1 - -5) + -5) 166 | # Gaussian prior 167 | params[:,3] = PhoIndex_gauss.ppf(cube[:,3]) 168 | # uniform priors 169 | params[:,4] = cube[:,4] * (80 - 7) + 7 170 | params[:,5] = cube[:,5] * (0.4) + 0 171 | # informative Gaussian prior on the redshift 172 | params[:,6] = z_gauss.ppf(cube[:,6]) 173 | # log-uniform prior on the background normalisation between 0.1 and 10 174 | params[:,7] = 10**(cube[:,7] * (1 - -1) + -1) 175 | return params 176 | 177 | 178 | # define parameter names 179 | param_names = ['lognorm', 'logNH22', 'scatnorm', 'PhoIndex', 'TORsigma', 'CTKcover', 'redshift', 'bkg_norm'] 180 | 181 | outprefix = sys.argv[1] + '_out_clumpy/' 182 | # run sampler 183 | try: 184 | sampler = ultranest.ReactiveNestedSampler( 185 | param_names, loglikelihood, prior_transform, 186 | log_dir=outprefix, resume=True, 187 | vectorized=True) 188 | except: 189 | sampler = ultranest.ReactiveNestedSampler( 190 | param_names, loglikelihood, prior_transform, 191 | log_dir=outprefix, resume='overwrite', 192 | vectorized=True) 193 | # then to run: 194 | results = sampler.run(max_num_improvement_loops=0, frac_remain=0.9, Lepsilon=0.1) 195 | sampler.print_results() 196 | #sampler.plot() 197 | samples = results['samples'].copy() 198 | abs_component, scat_component = setup_model(samples) 199 | soft_flux_obs = fastxsf.flux.energy_flux(galabso * (abs_component + scat_component), energies, 0.5, 2, axis=1) 200 | hard_flux_obs = fastxsf.flux.energy_flux(galabso * (abs_component + scat_component), energies, 2, 10, axis=1) 201 | print('soft flux obs:', soft_flux_obs.min(), soft_flux_obs.max(), soft_flux_obs.mean(), soft_flux_obs.std()) 202 | print('hard flux obs:', hard_flux_obs.min(), hard_flux_obs.max(), hard_flux_obs.mean(), hard_flux_obs.std()) 203 | print("setting NH to zero") 204 | samples[:,1] = -2 205 | samples[:,2] = -5 206 | abs_component, scat_component = setup_model(samples) 207 | soft_flux = fastxsf.flux.energy_flux(abs_component, energies, 0.5, 2, axis=1) 208 | hard_flux = fastxsf.flux.energy_flux(abs_component, energies, 2, 10, axis=1) 209 | hard_flux_restintr = np.empty(len(samples)) 210 | for i, row in enumerate(samples): 211 | z = row[6] 212 | hard_flux_restintr[i] = fastxsf.flux.energy_flux(abs_component[i, :], energies, 2 / (1 + z), 10 / (1 + z)).to_value(u.erg / u.cm**2 / u.s) 213 | 214 | print('soft flux:', soft_flux.min(), soft_flux.max(), soft_flux.mean(), soft_flux.std()) 215 | print('hard flux:', hard_flux.min(), hard_flux.max(), hard_flux.mean(), hard_flux.std()) 216 | # z, Fint, Fabs, lognorm, Gamma, logNH, fscat 217 | samples = results['samples'].copy() 218 | lognorm = samples[:, param_names.index('lognorm')] 219 | PhoIndex = samples[:, param_names.index('PhoIndex')] 220 | logNH = samples[:, param_names.index('logNH22')] + 22 221 | fscat = samples[:, param_names.index('scatnorm')] 222 | np.savetxt(f'{outprefix}/chains/intrinsic_photon_flux.txt.gz', np.transpose([ 223 | samples[:,6], hard_flux_restintr, soft_flux_obs, 224 | lognorm, PhoIndex, logNH, fscat, 225 | ])) 226 | 227 | # and to plot a model: 228 | loglikelihood(results['samples'][:20,:], plot=True, plot_prefix=f'{sys.argv[1]}_out_clumpy/plots/') 229 | 230 | # and to plot the posterior corner plot: 231 | sampler.plot() 232 | 233 | if __name__ == '__main__': 234 | main() 235 | -------------------------------------------------------------------------------- /fastxsf/model.py: -------------------------------------------------------------------------------- 1 | """Statistical and astrophysical models.""" 2 | import hashlib 3 | import itertools 4 | 5 | import astropy.io.fits as pyfits 6 | import numpy as np 7 | import tqdm 8 | import xspec_models_cxc as x 9 | from scipy.interpolate import RegularGridInterpolator 10 | 11 | from joblib import Memory 12 | 13 | mem = Memory('.', verbose=False) 14 | 15 | 16 | def logPoissonPDF_vectorized(models, counts): 17 | """Compute poisson probability. 18 | 19 | Parameters 20 | ---------- 21 | models: array 22 | expected number of counts. shape is (num_models, len(counts)). 23 | counts: array 24 | observed counts (non-negative integer) 25 | 26 | Returns 27 | ------- 28 | loglikelihood: array 29 | ln of the Poisson likelihood, neglecting the factorial(counts) factor, 30 | shape=(num_models,). 31 | """ 32 | log_models = np.log(np.clip(models, 1e-100, None)) 33 | return np.sum(log_models * counts.reshape((1, -1)), axis=1) - models.sum(axis=1) 34 | 35 | 36 | def logPoissonPDF(model, counts): 37 | """Compute poisson probability. 38 | 39 | Parameters 40 | ---------- 41 | model: array 42 | expected number of counts 43 | counts: array 44 | observed counts (non-negative integer) 45 | 46 | Returns 47 | ------- 48 | loglikelihood: float 49 | ln of the Poisson likelihood, neglecting the factorial(counts) factor. 50 | """ 51 | log_model = np.log(np.clip(model, 1e-100, None)) 52 | return np.sum(log_model * counts) - model.sum() 53 | 54 | 55 | def xvec(model, energies, pars): 56 | """Evaluate a model in a vectorized way. 57 | 58 | Parameters 59 | ---------- 60 | model: object 61 | xspec model (from fastxsf.x module, which is xspec_models_cxc) 62 | energies: array 63 | energies in keV where to evaluate model 64 | pars: array 65 | list of parameter vectors 66 | 67 | Returns 68 | ------- 69 | results: array 70 | for each parameter vector in pars, evaluates the model 71 | at the given energies. Has shape (pars.shape[0], energies.shape[0]) 72 | """ 73 | results = np.empty((len(pars), len(energies) - 1)) 74 | for i, pars_i in enumerate(pars): 75 | results[i, :] = model(energies=energies, pars=pars_i) 76 | return results 77 | 78 | 79 | @mem.cache 80 | def check_if_sorted(param_vals, parameter_grid): 81 | """Check if parameters are stored in a sorted way. 82 | 83 | Parameters 84 | ---------- 85 | param_vals: array 86 | list of parameter values stored 87 | parameter_grid: array 88 | list of possible values for each parameter 89 | 90 | Returns 91 | ------- 92 | sorted: bool 93 | True if param_vals==itertools.product(*parameter_grid) 94 | """ 95 | for i, params in enumerate(itertools.product(*parameter_grid)): 96 | if not np.all(param_vals[i] == params): 97 | return False 98 | return True 99 | 100 | 101 | def hashfile(filename): 102 | """Compute a hash for the content of a file.""" 103 | with open(filename, 'rb', buffering=0) as f: 104 | return hashlib.file_digest(f, 'sha256').hexdigest() 105 | 106 | 107 | @mem.cache 108 | def _load_table(filename, filehash=None): 109 | """Load data from table file. 110 | 111 | Parameters 112 | ---------- 113 | filename: str 114 | filename of a OGIP FITS file. 115 | filehash: str 116 | hash of the file 117 | 118 | Returns 119 | ------- 120 | parameter_grid: list 121 | list of values for each parameter 122 | data: array 123 | array of shape `(len(g) for g in parameter_grid)` 124 | containing the spectra for each parameter grid point. 125 | info: dict 126 | information about the table, including parameter_names, name, 127 | e_model_lo, e_model_hi, e_model_mid, deltae, parameter_grid 128 | """ 129 | f = pyfits.open(filename) 130 | assert f[0].header["MODLUNIT"] in ("photons/cm^2/s", "ergs/cm**2/s") 131 | assert f[0].header["HDUCLASS"] == "OGIP" 132 | self = {} 133 | self['parameter_names'] = f["PARAMETERS"].data["NAME"] 134 | self['name'] = f[0].header["MODLNAME"] 135 | parameter_grid = [ 136 | row["VALUE"][: row["NUMBVALS"]] for row in f["PARAMETERS"].data 137 | ] 138 | self['e_model_lo'] = f["ENERGIES"].data["ENERG_LO"] 139 | self['e_model_hi'] = f["ENERGIES"].data["ENERG_HI"] 140 | self['e_model_mid'] = (self['e_model_lo'] + self['e_model_hi']) / 2.0 141 | self['deltae'] = self['e_model_hi'] - self['e_model_lo'] 142 | specdata = f["SPECTRA"].data 143 | is_sorted = check_if_sorted(specdata["PARAMVAL"], parameter_grid) 144 | shape = tuple([len(g) for g in parameter_grid] + [len(specdata["INTPSPEC"][0])]) 145 | 146 | if is_sorted: 147 | data = specdata["INTPSPEC"].reshape(shape) 148 | else: 149 | data = np.nan * np.zeros( 150 | [len(g) for g in parameter_grid] + [len(specdata["INTPSPEC"][0])] 151 | ) 152 | for index, params, row in zip( 153 | tqdm.tqdm( 154 | list(itertools.product(*[range(len(g)) for g in parameter_grid])) 155 | ), 156 | itertools.product(*parameter_grid), 157 | sorted(specdata, key=lambda row: tuple(row["PARAMVAL"])), 158 | ): 159 | np.testing.assert_allclose(params, row["PARAMVAL"]) 160 | data[index] = row["INTPSPEC"] 161 | assert np.isfinite(data).all(), data 162 | return parameter_grid, data, self 163 | 164 | 165 | class Table: 166 | """Additive or multiplicative table model.""" 167 | 168 | def __init__(self, filename, method="linear", verbose=True): 169 | """Initialise. 170 | 171 | Parameters 172 | ---------- 173 | filename: str 174 | filename of a OGIP FITS file. 175 | method: str 176 | interpolation kind, passed to RegularGridInterpolator 177 | """ 178 | parameter_grid, data, info = _load_table(filename, hashfile(filename)) 179 | self.__dict__.update(info) 180 | if verbose: 181 | print(f'ATABLE "{self.name}"') 182 | for param_name, param_values in zip(self.parameter_names, parameter_grid): 183 | print(f" {param_name}: {param_values.tolist()}") 184 | self.interpolator = RegularGridInterpolator(parameter_grid, data, method=method) 185 | 186 | def __call__(self, energies, pars, vectorized=False): 187 | """Evaluate spectrum. 188 | 189 | Parameters 190 | ---------- 191 | energies: array 192 | energies in keV where spectrum should be computed 193 | pars: list 194 | parameter values. 195 | vectorized: bool 196 | if true, pars is a list of parameter vectors, 197 | and the function returns a list of spectra. 198 | 199 | Returns 200 | ------- 201 | spectrum: array 202 | photon spectrum, corresponding to the parameter values, 203 | one entry for each value in energies in phot/cm^2/s. 204 | """ 205 | if vectorized: 206 | z = pars[:, -1] 207 | e_lo = energies[:-1] 208 | e_hi = energies[1:] 209 | e_mid = (e_lo + e_hi) / 2.0 210 | delta_e = e_hi - e_lo 211 | model_int_spectrum = self.interpolator(pars[:, :-1]) 212 | results = np.empty((len(z), len(e_mid))) 213 | for i, zi in enumerate(z): 214 | # this model spectrum contains for each bin [e_lo...e_hi] the integral of energy 215 | # now we have a new energy, energies 216 | results[i, :] = ( 217 | np.interp( 218 | # look up in rest-frame, which is at higher energies at higher redshifts 219 | x=e_mid * (1 + zi), 220 | # in the model spectral grid 221 | xp=self.e_model_mid, 222 | # use spectral density, which is stretched out if redshifted. 223 | fp=model_int_spectrum[i, :] / self.deltae * (1 + zi), 224 | ) * delta_e / (1 + zi) 225 | ) 226 | return results 227 | else: 228 | z = pars[-1] 229 | e_lo = energies[:-1] 230 | e_hi = energies[1:] 231 | e_mid = (e_lo + e_hi) / 2.0 232 | delta_e = e_hi - e_lo 233 | (model_int_spectrum,) = self.interpolator([pars[:-1]]) 234 | # this model spectrum contains for each bin [e_lo...e_hi] the integral of energy 235 | # now we have a new energy, energies 236 | return ( 237 | np.interp( 238 | # look up in rest-frame, which is at higher energies at higher redshifts 239 | x=e_mid * (1 + z), 240 | # in the model spectral grid 241 | xp=self.e_model_mid, 242 | # use spectral density, which is stretched out if redshifted. 243 | fp=model_int_spectrum / self.deltae * (1 + z), 244 | ) * delta_e / (1 + z) 245 | ) 246 | 247 | 248 | @mem.cache 249 | def _load_redshift_interpolated_table(filename, filehash, energies, redshift, fix={}): 250 | """Load data from table file, precomputed for a energy grid. 251 | 252 | Parameters 253 | ---------- 254 | filename: str 255 | filename of a OGIP FITS file. 256 | filehash: str 257 | hash of the file 258 | energies: array 259 | Energies at which to compute model. 260 | redshift: float 261 | redshift to use for precomputing 262 | fix: dict 263 | dictionary of parameter names and their values to fix 264 | for faster data loading. 265 | 266 | Returns 267 | ------- 268 | newshape: tuple 269 | multidimensional shape of data according to info['parameter_grid'] 270 | data: array 271 | data table 272 | info: dict 273 | information about the table, including parameter_names, name, e_model_lo, e_model_hi, e_model_mid, deltae, parameter_grid 274 | """ 275 | parameter_grid, data, info = _load_table(filename, filehash) 276 | # interpolate data from original energy grid onto new energy grid 277 | e_lo = energies[:-1] 278 | e_hi = energies[1:] 279 | e_mid = (e_lo + e_hi) / 2.0 280 | delta_e = e_hi - e_lo 281 | # look up in rest-frame, which is at higher energies at higher redshifts 282 | e_mid_rest = e_mid * (1 + redshift) 283 | deltae_rest = delta_e / (1 + redshift) 284 | e_model_mid = info['e_model_mid'] 285 | info['energies'] = energies 286 | model_deltae_rest = info['deltae'] / (1 + redshift) 287 | 288 | # param_shapes = [len(p) for p in parameter_grid] 289 | # ndim = len(param_shapes) 290 | 291 | # Flatten the parameter grid into indices 292 | data_reshaped = data.reshape((-1, data.shape[-1])) 293 | n_points = data_reshaped.shape[0] 294 | mask = np.ones(n_points, dtype=bool) 295 | 296 | # Precompute grids 297 | param_grids = np.meshgrid(*parameter_grid, indexing='ij') 298 | # same shape as data without last dim 299 | param_grids_flat = [g.flatten() for g in param_grids] 300 | # each entry is flattened to match reshaped data 301 | parameter_names = [str(pname) for pname in info['parameter_names']] 302 | 303 | # Now apply fix conditions 304 | for pname, val in fix.items(): 305 | assert pname in parameter_names, (pname, parameter_names) 306 | param_idx = parameter_names.index(pname) 307 | mask &= (param_grids_flat[param_idx] == val) 308 | assert mask.any(), (f'You can only fix parameter {pname} to one of:', parameter_grid[param_idx]) 309 | 310 | # Mask valid rows 311 | valid_data = data_reshaped[mask] 312 | # Build new parameter grid (only for unfixed parameters) 313 | newparameter_grid = [] 314 | for p, pname in zip(parameter_grid, info['parameter_names']): 315 | if pname not in fix: 316 | newparameter_grid.append(p) 317 | 318 | # Interpolate 319 | newdata = np.zeros((valid_data.shape[0], len(e_mid_rest))) 320 | for i, row in enumerate(valid_data): 321 | newdata[i, :] = np.interp( 322 | x=e_mid_rest, 323 | xp=e_model_mid, 324 | fp=row / model_deltae_rest, 325 | ) * deltae_rest 326 | 327 | info['parameter_grid'] = newparameter_grid 328 | newshape = tuple([len(g) for g in newparameter_grid] + [len(e_mid_rest)]) 329 | return newshape, newdata, info 330 | 331 | 332 | @mem.cache 333 | def _load_redshift_interpolated_table_folded(filename, filehash, energies, redshift, ARF, RMF, fix={}): 334 | """Load data from table file, and fold it through the response. 335 | 336 | Parameters 337 | ---------- 338 | filename: str 339 | filename of a OGIP FITS file. 340 | filehash: str 341 | hash of the file 342 | energies: array 343 | Energies at which to compute model. 344 | redshift: float 345 | redshift to use for precomputing 346 | ARF: ARF 347 | area response function 348 | RMF: RMF 349 | response matrix 350 | fix: dict 351 | dictionary of parameter names and their values to fix 352 | for faster data loading. 353 | 354 | Returns 355 | ------- 356 | newshape: tuple 357 | multidimensional shape of data according to info['parameter_grid'] 358 | data: array 359 | data table 360 | info: dict 361 | information about the table, including parameter_names, name, e_model_lo, e_model_hi, e_model_mid, deltae, parameter_grid 362 | """ 363 | oldshape, olddata, info = _load_redshift_interpolated_table( 364 | filename, filehash, energies, redshift=redshift, fix=fix) 365 | newshape = list(oldshape) 366 | newshape[-1] = RMF.detchans 367 | newdata = RMF.apply_rmf_vectorized(olddata * ARF) 368 | assert newdata.shape == (len(olddata), RMF.detchans), (newdata.shape, olddata.shape, len(olddata), newshape[-1]) 369 | return newshape, newdata.reshape(newshape), info 370 | 371 | 372 | class FixedTable(Table): 373 | """Additive or multiplicative table model with fixed energy grid.""" 374 | 375 | def __init__(self, filename, energies, redshift=0, method="linear", fix={}, verbose=True): 376 | """Initialise. 377 | 378 | Parameters 379 | ---------- 380 | filename: str 381 | filename of a OGIP FITS file. 382 | energies: array 383 | energies in keV where spectrum should be computed 384 | redshift: float 385 | Redshift 386 | method: str 387 | interpolation kind, passed to RegularGridInterpolator 388 | fix: dict 389 | dictionary of parameter names and their values to fix 390 | for faster data loading. 391 | """ 392 | shape, data, info = _load_redshift_interpolated_table( 393 | filename, hashfile(filename), energies, redshift=redshift, fix=fix) 394 | self.__dict__.update(info) 395 | if verbose: 396 | print(f'ATABLE "{self.name}" (redshift={redshift})') 397 | for param_name, param_values in zip(self.parameter_names, self.parameter_grid): 398 | print(f" {param_name}: {param_values.tolist()}") 399 | self.interpolator = RegularGridInterpolator( 400 | self.parameter_grid, data.reshape(shape), 401 | method=method) 402 | 403 | def __call__(self, pars, vectorized=False, energies=None): 404 | """Evaluate spectrum. 405 | 406 | Parameters 407 | ---------- 408 | pars: list 409 | parameter values. 410 | vectorized: bool 411 | if true, pars is a list of parameter vectors, 412 | and the function returns a list of spectra. 413 | energies: array 414 | energies in keV where spectrum should be computed (ignored) 415 | 416 | Returns 417 | ------- 418 | spectrum: array 419 | photon spectrum, corresponding to the parameter values, 420 | one entry for each value in energies in phot/cm^2/s. 421 | """ 422 | assert np.size(vectorized) == 1 423 | assert np.ndim(vectorized) == 0 424 | if energies is not None: 425 | assert np.all(self.energies == energies) 426 | if vectorized: 427 | assert np.ndim(pars) == 2 428 | try: 429 | return self.interpolator(pars) 430 | except ValueError as e: 431 | raise ValueError(f"Interpolator with parameters {self.parameter_names} called with {pars}") from e 432 | else: 433 | assert np.ndim(pars) == 1 434 | try: 435 | return self.interpolator([pars])[0] 436 | except ValueError as e: 437 | pars_assigned = ' '.join([f'{k}={v}' for k, v in zip(self.parameter_names, pars)]) 438 | raise ValueError(f"Interpolator called with {pars_assigned}") from e 439 | 440 | 441 | class FixedFoldedTable(FixedTable): 442 | """Additive or multiplicative table model folded through response.""" 443 | 444 | def __init__(self, filename, energies, RMF, ARF, redshift=0, method="linear", fix={}, verbose=True): 445 | """Initialise. 446 | 447 | Parameters 448 | ---------- 449 | filename: str 450 | filename of a OGIP FITS file. 451 | energies: array 452 | energies in keV where spectrum should be computed 453 | redshift: float 454 | Redshift 455 | method: str 456 | interpolation kind, passed to RegularGridInterpolator 457 | fix: dict 458 | dictionary of parameter names and their values to fix 459 | for faster data loading. 460 | """ 461 | shape, data, info = _load_redshift_interpolated_table_folded( 462 | filename, hashfile(filename), energies, redshift=redshift, fix=fix, 463 | RMF=RMF, ARF=ARF) 464 | self.__dict__.update(info) 465 | if verbose: 466 | print(f'ATABLE "{self.name}" (redshift={redshift})') 467 | for param_name, param_values in zip(self.parameter_names, self.parameter_grid): 468 | print(f" {param_name}: {param_values.tolist()}") 469 | self.interpolator = RegularGridInterpolator( 470 | self.parameter_grid, data.reshape(shape), 471 | method=method) 472 | 473 | 474 | def prepare_folded_model0d(model, energies, pars, ARF, RMF, nonnegative=True): 475 | """Prepare a folded spectrum. 476 | 477 | Parameters 478 | ---------- 479 | model: object 480 | xspec model (from fastxsf.x module, which is xspec_models_cxc) 481 | energies: array 482 | energies in keV where spectrum should be computed (ignored) 483 | pars: list 484 | parameter values. 485 | ARF: array 486 | vector for multiplication before applying the RMF 487 | RMF: RMF 488 | RMF object for folding 489 | nonnegative: bool 490 | 491 | 492 | Returns 493 | ------- 494 | folded_spectrum: array 495 | folded spectrum after applying RMF & ARF 496 | """ 497 | if nonnegative: 498 | return RMF.apply_rmf(np.clip(model(energies=energies, pars=pars), 0, None) * ARF) 499 | else: 500 | return RMF.apply_rmf(model(energies=energies, pars=pars) * ARF) 501 | 502 | 503 | @mem.cache(ignore=['model']) 504 | def _prepare_folded_model1d(model, modelname, energies, pars, ARF, RMF, nonnegative=True): 505 | """Prepare a function that returns the folded model. 506 | 507 | Parameters 508 | ---------- 509 | model: object 510 | xspec model (from fastxsf.x module, which is xspec_models_cxc) 511 | modelname: str 512 | name of xspec model 513 | energies: array 514 | energies in keV where spectrum should be computed (ignored) 515 | pars: list 516 | parameter values; one of the entries can be an array, 517 | which will be the interpolation range. 518 | ARF: array 519 | vector for multiplication before applying the RMF 520 | RMF: RMF 521 | RMF object for folding 522 | nonnegative: bool 523 | 524 | 525 | Returns 526 | ------- 527 | freeparam_grid: tuple 528 | the pars element which is an array. 529 | folded_spectrum: array 530 | folded spectrum after applying RMF & ARF 531 | """ 532 | mask_fixed = np.array([np.size(p) == 1 for p in pars]) 533 | assert (~mask_fixed).sum() == 1, mask_fixed 534 | i_variable = np.where(~(mask_fixed))[0][0] 535 | assert pars[i_variable].ndim == 1 536 | 537 | data = np.zeros((len(pars[i_variable]), len(ARF))) 538 | for i, variable in enumerate(pars[i_variable]): 539 | pars_row = list(pars) 540 | pars_row[i_variable] = variable 541 | data[i] = model(energies=energies, pars=pars_row) 542 | foldeddata = RMF.apply_rmf_vectorized(data * ARF) 543 | if nonnegative: 544 | foldeddata = np.clip(foldeddata, 0, None) 545 | return pars[i_variable], foldeddata 546 | 547 | 548 | def prepare_folded_model1d(model, energies, pars, ARF, RMF, nonnegative=True, method='linear'): 549 | """Prepare a function that returns the folded model. 550 | 551 | Parameters 552 | ---------- 553 | model: object 554 | xspec model (from fastxsf.x module, which is xspec_models_cxc) 555 | energies: array 556 | energies in keV where spectrum should be computed (ignored) 557 | pars: list 558 | parameter values; one of the entries can be an array, 559 | which will be the interpolation range. 560 | ARF: array 561 | vector for multiplication before applying the RMF 562 | RMF: RMF 563 | RMF object for folding 564 | nonnegative: bool 565 | 566 | method: str 567 | interpolation kind, passed to RegularGridInterpolator 568 | 569 | Returns 570 | ------- 571 | simple_interpolator: func 572 | function that given the free parameter value returns a folded spectrum. 573 | """ 574 | grid, foldeddata = _prepare_folded_model1d( 575 | model=model, modelname=model.__name__, pars=pars, 576 | energies=energies, ARF=ARF, RMF=RMF, nonnegative=True) 577 | interp = RegularGridInterpolator((grid,), foldeddata, method=method) 578 | 579 | def simple_interpolator(par): 580 | """Interpolator for a single parameter. 581 | 582 | Parameters 583 | ---------- 584 | par: float 585 | The value for the one free model parameter 586 | 587 | Returns 588 | ------- 589 | spectrum: array 590 | photon spectrum 591 | """ 592 | try: 593 | return interp([par])[0] 594 | except ValueError as e: 595 | raise ValueError(f'invalid parameter value passed: {par}') from e 596 | 597 | return simple_interpolator 598 | 599 | # build a folded model with mean and std of a parameter being a distribution? 600 | -------------------------------------------------------------------------------- /fastxsf/response.py: -------------------------------------------------------------------------------- 1 | """Functionality for linear instrument response.""" 2 | # from https://github.com/dhuppenkothen/clarsach/blob/master/clarsach/respond.py 3 | # GPL licenced code from the Clàrsach project 4 | from functools import partial 5 | 6 | import astropy.io.fits as fits 7 | import jax 8 | import numpy as np 9 | 10 | __all__ = ["RMF", "ARF", "MockARF"] 11 | 12 | 13 | def _apply_rmf(spec, in_indices, out_indices, weights, detchans): 14 | """Apply RMF. 15 | 16 | Parameters 17 | ---------- 18 | spec: array 19 | Input spectrum. 20 | in_indices: array 21 | List of indices for spec. 22 | out_indices: array 23 | list of indices for where to add into output array 24 | weights: array 25 | list of weights for multiplying *spec[outindex]* when adding to output array 26 | detchans: int 27 | length of output array 28 | 29 | Returns 30 | ------- 31 | out: array 32 | Summed entries. 33 | """ 34 | contribs = spec[in_indices] * weights 35 | out = jax.numpy.zeros(detchans) 36 | out = out.at[out_indices].add(contribs) 37 | return out 38 | 39 | 40 | def _apply_rmf_vectorized(specs, in_indices, out_indices, weights, detchans): 41 | """Apply RMF to many spectra. 42 | 43 | Parameters 44 | ---------- 45 | specs: array 46 | List of input spectra. 47 | in_indices: array 48 | List of indices for spec. 49 | out_indices: array 50 | list of indices for where to add into output array 51 | weights: array 52 | list of weights for multiplying *spec[outindex]* when adding to output array 53 | detchans: int 54 | length of output array 55 | 56 | Returns 57 | ------- 58 | out: array 59 | Summed entries. Shape=(len(specs), detchans) 60 | """ 61 | out = jax.numpy.zeros((len(specs), detchans)) 62 | 63 | def body_fun(j, out): 64 | """Sum up one spectrum. 65 | 66 | Parameters 67 | ---------- 68 | j: int 69 | index of spectrum. 70 | out: array 71 | will store into out[j,:] 72 | 73 | Returns 74 | ------- 75 | out_row: array 76 | returns out[j] 77 | """ 78 | spec = specs[j] 79 | contribs = spec[in_indices] * weights 80 | return out.at[j, out_indices].add(contribs) 81 | 82 | out = jax.lax.fori_loop(0, len(specs), body_fun, out) 83 | return out 84 | 85 | 86 | class RMF(object): 87 | """Response matrix file.""" 88 | 89 | def __init__(self, filename): 90 | """ 91 | Initialise. 92 | 93 | Parameters 94 | ---------- 95 | filename : str 96 | The file name with the RMF FITS file 97 | """ 98 | self._load_rmf(filename) 99 | 100 | def _load_rmf(self, filename): 101 | """ 102 | Load an RMF from a FITS file. 103 | 104 | Parameters 105 | ---------- 106 | filename : str 107 | The file name with the RMF file 108 | 109 | Attributes 110 | ---------- 111 | n_grp : numpy.ndarray 112 | the Array with the number of channels in each 113 | channel set 114 | 115 | f_chan : numpy.ndarray 116 | The starting channel for each channel group; 117 | If an element i in n_grp > 1, then the resulting 118 | row entry in f_chan will be a list of length n_grp[i]; 119 | otherwise it will be a single number 120 | 121 | n_chan : numpy.ndarray 122 | The number of channels in each channel group. The same 123 | logic as for f_chan applies 124 | 125 | matrix : numpy.ndarray 126 | The redistribution matrix as a flattened 1D vector 127 | 128 | energ_lo : numpy.ndarray 129 | The lower edges of the energy bins 130 | 131 | energ_hi : numpy.ndarray 132 | The upper edges of the energy bins 133 | 134 | detchans : int 135 | The number of channels in the detector 136 | """ 137 | # open the FITS file and extract the MATRIX extension 138 | # which contains the redistribution matrix and 139 | # anxillary information 140 | hdulist = fits.open(filename) 141 | 142 | # get all the extension names 143 | extnames = np.array([h.name for h in hdulist]) 144 | 145 | # figure out the right extension to use 146 | if "MATRIX" in extnames: 147 | h = hdulist["MATRIX"] 148 | 149 | elif "SPECRESP MATRIX" in extnames: 150 | h = hdulist["SPECRESP MATRIX"] 151 | elif "SPECRESP" in extnames: 152 | h = hdulist["SPECRESP"] 153 | else: 154 | raise AssertionError(f"{extnames} does not contain MATRIX or SPECRESP") 155 | 156 | data = h.data 157 | hdr = dict(h.header) 158 | hdulist.close() 159 | 160 | # extract + store the attributes described in the docstring 161 | n_grp = np.array(data.field("N_GRP")) 162 | f_chan = np.array(data.field("F_CHAN")) 163 | n_chan = np.array(data.field("N_CHAN")) 164 | matrix = np.array(data.field("MATRIX")) 165 | 166 | self.energ_lo = np.array(data.field("ENERG_LO")) 167 | self.energ_hi = np.array(data.field("ENERG_HI")) 168 | self.energ_unit = data.columns["ENERG_LO"].unit 169 | self.detchans = int(hdr["DETCHANS"]) 170 | self.offset = self.__get_tlmin(h) 171 | 172 | # flatten the variable-length arrays 173 | results = self._flatten_arrays(n_grp, f_chan, n_chan, matrix) 174 | self.n_grp, self.f_chan, self.n_chan, self.matrix = results 175 | self.dense_info = None 176 | 177 | def __get_tlmin(self, h): 178 | """ 179 | Get the tlmin keyword for `F_CHAN`. 180 | 181 | Parameters 182 | ---------- 183 | h : an astropy.io.fits.hdu.table.BinTableHDU object 184 | The extension containing the `F_CHAN` column 185 | 186 | Returns 187 | ------- 188 | tlmin : int 189 | The tlmin keyword 190 | """ 191 | # get the header 192 | hdr = h.header 193 | # get the keys of all 194 | keys = np.array(list(hdr.keys())) 195 | 196 | # find the place where the tlmin keyword is defined 197 | t = np.array(["TLMIN" in k for k in keys]) 198 | 199 | # get the index of the TLMIN keyword 200 | tlmin_idx = np.hstack(np.where(t))[0] 201 | 202 | # get the corresponding value 203 | tlmin = int(list(hdr.items())[tlmin_idx][1]) 204 | 205 | return tlmin 206 | 207 | def _flatten_arrays(self, n_grp, f_chan, n_chan, matrix): 208 | """Flatten array. 209 | 210 | Parameters 211 | ---------- 212 | n_grp: array 213 | number of groups 214 | f_chan: array 215 | output start indices 216 | n_chan: array 217 | number of output indices 218 | matrix: array 219 | weights 220 | 221 | Returns 222 | ------- 223 | n_grp: array 224 | number of groups 225 | f_chan: array 226 | output start indices 227 | n_chan: array 228 | number of output indices 229 | matrix: array 230 | weights 231 | """ 232 | if not len(n_grp) == len(f_chan) == len(n_chan) == len(matrix): 233 | raise ValueError("Arrays must be of same length!") 234 | 235 | # find all non-zero groups 236 | nz_idx = n_grp > 0 237 | 238 | # stack all non-zero rows in the matrix 239 | matrix_flat = np.hstack(matrix[nz_idx], dtype=float) 240 | 241 | # stack all nonzero rows in n_chan and f_chan 242 | # n_chan_flat = np.hstack(n_chan[nz_idx]) 243 | # f_chan_flat = np.hstack(f_chan[nz_idx]) 244 | 245 | # some matrices actually have more elements 246 | # than groups in `n_grp`, so we'll only pick out 247 | # those values that have a correspondence in 248 | # n_grp 249 | f_chan_new = [] 250 | n_chan_new = [] 251 | for i, t in enumerate(nz_idx): 252 | if t: 253 | n = n_grp[i] 254 | f = f_chan[i] 255 | nc = n_chan[i] 256 | if np.size(f) == 1: 257 | f_chan_new.append(f.astype(np.int64) - self.offset) 258 | n_chan_new.append(nc.astype(np.int64)) 259 | else: 260 | f_chan_new.append(f[:n].astype(np.int64) - self.offset) 261 | n_chan_new.append(nc[:n].astype(np.int64)) 262 | 263 | n_chan_flat = np.hstack(n_chan_new) 264 | f_chan_flat = np.hstack(f_chan_new) 265 | 266 | return n_grp, f_chan_flat, n_chan_flat, matrix_flat 267 | 268 | def strip(self, channel_mask): 269 | """Strip response matrix of entries outside the channel mask. 270 | 271 | Parameters 272 | ---------- 273 | channel_mask : array 274 | Boolean array indicating which detector channel to keep. 275 | 276 | Returns 277 | ------- 278 | energy_mask : array 279 | Boolean array indicating which energy channels were kept. 280 | """ 281 | n_grp_new = np.zeros_like(self.n_grp) 282 | n_chan_new = [] 283 | f_chan_new = [] 284 | matrix_new = [] 285 | energ_lo_new = [] 286 | energ_hi_new = [] 287 | 288 | in_indices = [] 289 | i_new = 0 290 | out_indices = [] 291 | weights = [] 292 | k = 0 293 | resp_idx = 0 294 | # loop over all channels 295 | for i in range(len(self.energ_lo)): 296 | # get the current number of groups 297 | current_num_groups = self.n_grp[i] 298 | 299 | # loop over the current number of groups 300 | for current_num_chans, counts_idx in zip( 301 | self.n_chan[k:k + current_num_groups], 302 | self.f_chan[k:k + current_num_groups], 303 | ): 304 | # add the flux to the subarray of the counts array that starts with 305 | # counts_idx and runs over current_num_chans channels 306 | # outslice = slice(counts_idx, counts_idx + current_num_chans) 307 | inslice = slice(resp_idx, resp_idx + current_num_chans) 308 | mask_valid = channel_mask[counts_idx:counts_idx + current_num_chans] 309 | if current_num_chans > 0 and mask_valid.any(): 310 | # add block 311 | n_grp_new[i_new] += 1 312 | # length 313 | n_chan_new.append(current_num_chans) 314 | # location in matrix 315 | f_chan_new.append(counts_idx) 316 | matrix_new.append(self.matrix[inslice]) 317 | 318 | in_indices.append((i_new + np.zeros(current_num_chans, dtype=int))[mask_valid]) 319 | out_indices.append(np.arange(counts_idx, counts_idx + current_num_chans)[mask_valid]) 320 | weights.append(self.matrix[inslice][mask_valid]) 321 | resp_idx += current_num_chans 322 | 323 | k += current_num_groups 324 | if n_grp_new[i_new] > 0: 325 | energ_lo_new.append(self.energ_lo[i]) 326 | energ_hi_new.append(self.energ_hi[i]) 327 | i_new += 1 328 | 329 | out_indices = np.hstack(out_indices) 330 | in_indices = np.hstack(in_indices) 331 | weights = np.hstack(weights) 332 | self.n_chan = np.array(n_chan_new) 333 | self.f_chan = np.array(f_chan_new) 334 | 335 | # cut down input array as well 336 | strip_mask = n_grp_new > 0 337 | self.n_grp = n_grp_new[strip_mask] 338 | self.energ_lo = np.array(energ_lo_new) 339 | self.energ_hi = np.array(energ_hi_new) 340 | 341 | self.matrix = np.hstack(matrix_new) 342 | i = np.argsort(out_indices) 343 | # self.dense_info = in_indices[i], out_indices[i], weights[i] 344 | self.dense_info = in_indices, out_indices, weights 345 | self._compile() 346 | return strip_mask 347 | 348 | def _compile(self): 349 | """Prepare internal functions.""" 350 | if self.dense_info is None: 351 | return 352 | in_indices, out_indices, weights = self.dense_info 353 | self._apply_rmf = jax.jit( 354 | partial( 355 | _apply_rmf, 356 | in_indices=in_indices, 357 | out_indices=out_indices, 358 | weights=weights, 359 | detchans=self.detchans, 360 | ) 361 | ) 362 | self._apply_rmf_vectorized = jax.jit( 363 | partial( 364 | _apply_rmf_vectorized, 365 | in_indices=in_indices, 366 | out_indices=out_indices, 367 | weights=weights, 368 | detchans=self.detchans, 369 | ) 370 | ) 371 | 372 | def __getstate__(self): 373 | """Get state for pickling.""" 374 | state = self.__dict__.copy() 375 | # Remove non-pickleable functions 376 | if "_apply_rmf" in state: 377 | del state["_apply_rmf"] 378 | if "_apply_rmf_vectorized" in state: 379 | del state["_apply_rmf_vectorized"] 380 | return state 381 | 382 | def __setstate__(self, state): 383 | """Restore state from pickling.""" 384 | self.__dict__.update(state) 385 | self._compile() 386 | 387 | def apply_rmf(self, spec): 388 | """ 389 | Fold the spectrum through the redistribution matrix. 390 | 391 | The redistribution matrix is saved as a flattened 1-dimensional 392 | vector to save space. In reality, for each entry in the flux 393 | vector, there exists one or more sets of channels that this 394 | flux is redistributed into. The additional arrays `n_grp`, 395 | `f_chan` and `n_chan` store this information: 396 | * `n_group` stores the number of channel groups for each 397 | energy bin 398 | * `f_chan` stores the *first channel* that each channel 399 | for each channel set 400 | * `n_chan` stores the number of channels in each channel 401 | set 402 | 403 | As a result, for a given energy bin i, we need to look up the 404 | number of channel sets in `n_grp` for that energy bin. We 405 | then need to loop over the number of channel sets. For each 406 | channel set, we look up the first channel into which flux 407 | will be distributed as well as the number of channels in the 408 | group. We then need to also loop over the these channels and 409 | actually use the corresponding elements in the redistribution 410 | matrix to redistribute the photon flux into channels. 411 | 412 | All of this is basically a big bookkeeping exercise in making 413 | sure to get the indices right. 414 | 415 | Parameters 416 | ---------- 417 | spec : numpy.ndarray 418 | The (model) spectrum to be folded 419 | 420 | Returns 421 | ------- 422 | counts : numpy.ndarray 423 | The (model) spectrum after folding, in 424 | counts/s/channel 425 | """ 426 | if self.dense_info is not None: 427 | # in_indices, out_indices, weights = self.dense_info 428 | # 0.001658s/call; 0.096s/likelihood eval 429 | # out = np.zeros(self.detchans) 430 | # np.add.at(out, out_indices, spec[in_indices] * weights) 431 | # 0.004929s/call; 0.106s/likelihood eval 432 | # out = np.bincount(out_indices, weights=spec[in_indices] * weights, minlength=self.detchans) 433 | # 0.001963s/call; 0.077s/likelihood eval 434 | out = self._apply_rmf( 435 | spec 436 | ) # , in_indices, out_indices, weights, self.detchans) 437 | return out 438 | 439 | # get the number of channels in the data 440 | nchannels = spec.shape[0] 441 | 442 | # an empty array for the output counts 443 | counts = np.zeros(nchannels) 444 | 445 | # index for n_chan and f_chan incrementation 446 | k = 0 447 | 448 | # index for the response matrix incrementation 449 | resp_idx = 0 450 | 451 | # loop over all channels 452 | for i in range(nchannels): 453 | # this is the current bin in the flux spectrum to 454 | # be folded 455 | source_bin_i = spec[i] 456 | 457 | # get the current number of groups 458 | current_num_groups = self.n_grp[i] 459 | 460 | # loop over the current number of groups 461 | for current_num_chans, counts_idx in zip( 462 | self.n_chan[k:k + current_num_groups], 463 | self.f_chan[k:k + current_num_groups] 464 | ): 465 | # add the flux to the subarray of the counts array that starts with 466 | # counts_idx and runs over current_num_chans channels 467 | outslice = slice(counts_idx, counts_idx + current_num_chans) 468 | inslice = slice(resp_idx, resp_idx + current_num_chans) 469 | counts[outslice] += self.matrix[inslice] * source_bin_i 470 | # iterate the response index for next round 471 | resp_idx += current_num_chans 472 | k += current_num_groups 473 | 474 | return counts[:self.detchans] 475 | 476 | def apply_rmf_vectorized(self, specs): 477 | """ 478 | Fold the spectrum through the redistribution matrix. 479 | 480 | The redistribution matrix is saved as a flattened 1-dimensional 481 | vector to save space. In reality, for each entry in the flux 482 | vector, there exists one or more sets of channels that this 483 | flux is redistributed into. The additional arrays `n_grp`, 484 | `f_chan` and `n_chan` store this information: 485 | * `n_group` stores the number of channel groups for each 486 | energy bin 487 | * `f_chan` stores the *first channel* that each channel 488 | for each channel set 489 | * `n_chan` stores the number of channels in each channel 490 | set 491 | 492 | As a result, for a given energy bin i, we need to look up the 493 | number of channel sets in `n_grp` for that energy bin. We 494 | then need to loop over the number of channel sets. For each 495 | channel set, we look up the first channel into which flux 496 | will be distributed as well as the number of channels in the 497 | group. We then need to also loop over the these channels and 498 | actually use the corresponding elements in the redistribution 499 | matrix to redistribute the photon flux into channels. 500 | 501 | All of this is basically a big bookkeeping exercise in making 502 | sure to get the indices right. 503 | 504 | Parameters 505 | ---------- 506 | specs : numpy.ndarray 507 | The (model) spectra to be folded 508 | 509 | Returns 510 | ------- 511 | counts : numpy.ndarray 512 | The (model) spectrum after folding, in counts/s/channel 513 | """ 514 | # get the number of channels in the data 515 | nspecs, nchannels = specs.shape 516 | if self.dense_info is not None: # and nspecs < 40: 517 | in_indices, out_indices, weights = self.dense_info 518 | out = np.zeros((nspecs, self.detchans)) 519 | for i in range(nspecs): 520 | # out[i] = np.bincount(out_indices, weights=specs[i,in_indices] * weights, minlength=self.detchans) 521 | # out[i] = self._apply_rmf(specs[i]) 522 | out[i] = jax.numpy.zeros(self.detchans).at[out_indices].add(specs[i,in_indices] * weights) 523 | # out = self._apply_rmf_vectorized(specs) 524 | return out 525 | 526 | # an empty array for the output counts 527 | counts = np.zeros((nspecs, nchannels)) 528 | 529 | # index for n_chan and f_chan incrementation 530 | k = 0 531 | 532 | # index for the response matrix incrementation 533 | resp_idx = 0 534 | 535 | # loop over all channels 536 | for i in range(nchannels): 537 | # this is the current bin in the flux spectrum to 538 | # be folded 539 | source_bin_i = specs[:,i] 540 | 541 | # get the current number of groups 542 | current_num_groups = self.n_grp[i] 543 | 544 | # loop over the current number of groups 545 | for current_num_chans, counts_idx in zip( 546 | self.n_chan[k:k + current_num_groups], 547 | self.f_chan[k:k + current_num_groups] 548 | ): 549 | # add the flux to the subarray of the counts array that starts with 550 | # counts_idx and runs over current_num_chans channels 551 | to_add = np.outer( 552 | source_bin_i, self.matrix[resp_idx:resp_idx + current_num_chans] 553 | ) 554 | counts[:, counts_idx:counts_idx + current_num_chans] += to_add 555 | 556 | # iterate the response index for next round 557 | resp_idx += current_num_chans 558 | k += current_num_groups 559 | 560 | return counts[:, : self.detchans] 561 | 562 | def get_dense_matrix(self): 563 | """Extract the redistribution matrix as a dense numpy matrix. 564 | 565 | The redistribution matrix is saved as a 1-dimensional 566 | vector to save space (see apply_rmf for more information). 567 | This function converts it into a dense array. 568 | 569 | Returns 570 | ------- 571 | dense_matrix : numpy.ndarray 572 | The RMF as a dense 2d matrix. 573 | """ 574 | # get the number of channels in the data 575 | nchannels = len(self.energ_lo) 576 | nenergies = self.detchans 577 | 578 | # an empty array for the output counts 579 | dense_matrix = np.zeros((nchannels, nenergies)) 580 | 581 | # index for n_chan and f_chan incrementation 582 | k = 0 583 | 584 | # index for the response matrix incrementation 585 | resp_idx = 0 586 | 587 | # loop over all channels 588 | for i in range(nchannels): 589 | # get the current number of groups 590 | current_num_groups = self.n_grp[i] 591 | 592 | # loop over the current number of groups 593 | for _ in range(current_num_groups): 594 | current_num_chans = int(self.n_chan[k]) 595 | # get the right index for the start of the counts array 596 | # to put the data into 597 | counts_idx = self.f_chan[k] 598 | # this is the current number of channels to use 599 | 600 | k += 1 601 | 602 | # assign the subarray of the counts array that starts with 603 | # counts_idx and runs over current_num_chans channels 604 | 605 | outslice = slice(counts_idx, counts_idx + current_num_chans) 606 | inslice = slice(resp_idx, resp_idx + current_num_chans) 607 | dense_matrix[i, outslice] = self.matrix[inslice] 608 | 609 | # iterate the response index for next round 610 | resp_idx += current_num_chans 611 | 612 | return dense_matrix 613 | 614 | 615 | class ARF(object): 616 | """Area response file.""" 617 | 618 | def __init__(self, filename): 619 | """Initialise. 620 | 621 | Parameters 622 | ---------- 623 | filename : str 624 | The file name with the ARF file 625 | """ 626 | self._load_arf(filename) 627 | pass 628 | 629 | def _load_arf(self, filename): 630 | """Load an ARF from a FITS file. 631 | 632 | Parameters 633 | ---------- 634 | filename : str 635 | The file name with the ARF file 636 | """ 637 | # open the FITS file and extract the MATRIX extension 638 | # which contains the redistribution matrix and 639 | # anxillary information 640 | hdulist = fits.open(filename) 641 | h = hdulist["SPECRESP"] 642 | data = h.data 643 | hdr = h.header 644 | hdulist.close() 645 | 646 | # extract + store the attributes described in the docstring 647 | 648 | self.e_low = np.array(data.field("ENERG_LO")) 649 | self.e_high = np.array(data.field("ENERG_HI")) 650 | self.e_unit = data.columns["ENERG_LO"].unit 651 | self.specresp = np.array(data.field("SPECRESP")) 652 | 653 | if "EXPOSURE" in list(hdr.keys()): 654 | self.exposure = float(hdr["EXPOSURE"]) 655 | else: 656 | self.exposure = 1.0 657 | 658 | if "FRACEXPO" in data.columns.names: 659 | self.fracexpo = float(data["FRACEXPO"]) 660 | else: 661 | self.fracexpo = 1.0 662 | 663 | def strip(self, mask): 664 | """Remove unneeded energy ranges. 665 | 666 | Parameters 667 | ---------- 668 | mask: array 669 | Boolean array indicating which energy channel to keep. 670 | """ 671 | self.e_low = self.e_low[mask] 672 | self.e_high = self.e_high[mask] 673 | self.specresp = self.specresp[mask] 674 | 675 | def apply_arf(self, spec, exposure=None): 676 | """Fold the spectrum through the ARF. 677 | 678 | The ARF is a single vector encoding the effective area information 679 | about the detector. A such, applying the ARF is a simple 680 | multiplication with the input spectrum. 681 | 682 | Parameters 683 | ---------- 684 | spec : numpy.ndarray 685 | The (model) spectrum to be folded 686 | 687 | exposure : float, default None 688 | Value for the exposure time. By default, `apply_arf` will use the 689 | exposure keyword from the ARF file. If this exposure time is not 690 | correct (for example when simulated spectra use a different exposure 691 | time and the ARF from a real observation), one can override the 692 | default exposure by setting the `exposure` keyword to the correct 693 | value. 694 | 695 | Returns 696 | ------- 697 | s_arf : numpy.ndarray 698 | The (model) spectrum after folding, in 699 | counts/s/channel 700 | """ 701 | assert spec.shape[0] == self.specresp.shape[0], ( 702 | "Input spectrum and ARF must be of same size.", 703 | spec.shape, 704 | self.specresp.shape, 705 | ) 706 | e = self.exposure if exposure is None else exposure 707 | return spec * self.specresp * e 708 | 709 | 710 | class MockARF(ARF): 711 | """Mock area response file.""" 712 | 713 | def __init__(self, rmf): 714 | """Initialise. 715 | 716 | Parameters 717 | ---------- 718 | rmf: RMF 719 | RMF object to mimic. 720 | """ 721 | self.e_low = rmf.energ_lo 722 | self.e_high = rmf.energ_hi 723 | self.e_unit = rmf.energ_unit 724 | self.exposure = None 725 | self.specresp = np.ones_like(self.e_low) 726 | -------------------------------------------------------------------------------- /multispecopt.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | A new approach to X-ray spectral fitting with Xspec models and 4 | optimized nested sampling. 5 | 6 | This script explains how to set up your model. 7 | 8 | The idea is that you create a function which computes all model components. 9 | 10 | """ 11 | import corner 12 | import os 13 | import numpy as np 14 | import scipy.stats 15 | from matplotlib import pyplot as plt 16 | from optns.sampler import OptNS 17 | from optns.profilelike import GaussianPrior 18 | from astropy.cosmology import LambdaCDM 19 | 20 | cosmo = LambdaCDM(H0=70, Om0=0.3, Ode0=0.730) 21 | 22 | import fastxsf 23 | 24 | # fastxsf.x.chatter(0) 25 | fastxsf.x.abundance('wilm') 26 | fastxsf.x.cross_section('vern') 27 | 28 | # let's take a realistic example of a Chandra + NuSTAR FPMA + FPMB spectrum 29 | # with normalisation cross-calibration uncertainty of +-0.2 dex. 30 | # and a soft apec, a pexmon and a UXCLUMPY model, plus a background of course 31 | 32 | # we want to make pretty plots of the fit and its components, folded and unfolded, 33 | # compute 2-10 keV fluxes of the components 34 | # compute luminosities of the intrinsic power law 35 | 36 | filepath = '/mnt/data/daten/PostDoc2/research/agn/eROSITA/xlf/xrayspectra/NuSTARenhance/COSMOS/spectra/102/' 37 | 38 | data_sets = { 39 | 'Chandra': fastxsf.load_pha(filepath + 'C.pha', 0.5, 8), 40 | 'NuSTAR-FPMA': fastxsf.load_pha(filepath + 'A.pha', 4, 77), 41 | 'NuSTAR-FPMB': fastxsf.load_pha(filepath + 'B.pha', 4, 77), 42 | } 43 | 44 | redshift = data_sets['Chandra']['redshift'] 45 | 46 | # pre-compute the absorption factors -- no need to call this again and again if the parameters do not change! 47 | galabsos = { 48 | k: fastxsf.x.TBabs(energies=data['energies'], pars=[data['galnh']]) 49 | for k, data in data_sets.items() 50 | } 51 | 52 | # load a Table model 53 | tablepath = os.path.join(os.environ.get('MODELDIR', '.'), 'uxclumpy-cutoff.fits') 54 | import time 55 | t0 = time.time() 56 | print("preparing fixed table models...") 57 | absAGN = fastxsf.model.Table(tablepath) 58 | print(f'took {time.time() - t0:.3f}s') 59 | t0 = time.time() 60 | print("preparing folded table models...") 61 | absAGNs = { 62 | k: fastxsf.model.FixedTable( 63 | tablepath, energies=data['energies'], redshift=redshift) 64 | for k, data in data_sets.items() 65 | } 66 | absAGN_folded = { 67 | k: fastxsf.model.FixedFoldedTable( 68 | tablepath, energies=data['energies'], ARF=data['ARF'] * galabsos[k], RMF=data['RMF_src'], redshift=redshift, fix=dict(Ecut=400, Theta_inc=60)) 69 | for k, data in data_sets.items() 70 | } 71 | print(f'took {time.time() - t0:.3f}s') 72 | t0 = time.time() 73 | print("preparing 1d interpolated models...") 74 | scat_folded = { 75 | k: fastxsf.model.prepare_folded_model1d(fastxsf.x.zpowerlw, pars=[np.arange(1.0, 3.1, 0.01), redshift], energies=data['energies'], ARF=data['ARF'] * galabsos[k], RMF=data['RMF_src']) 76 | for k, data in data_sets.items() 77 | } 78 | apec_folded = { 79 | k: fastxsf.model.prepare_folded_model1d(fastxsf.x.apec, pars=[10**np.arange(-2, 1.2, 0.01), 1.0, redshift], energies=data['energies'], ARF=data['ARF'] * galabsos[k], RMF=data['RMF_src']) 80 | for k, data in data_sets.items() 81 | } 82 | print(f'took {time.time() - t0:.3f}s') 83 | print(scat_folded['Chandra'](2.0), scat_folded['Chandra']([2.0]).shape) 84 | assert scat_folded['Chandra'](2.0).shape == data_sets['Chandra']['chan_mask'].shape 85 | print(apec_folded['Chandra'](2.0), apec_folded['Chandra']([2.0]).shape) 86 | assert apec_folded['Chandra'](2.0).shape == data_sets['Chandra']['chan_mask'].shape 87 | 88 | # pre-compute the absorption factors -- no need to call this again and again if the parameters do not change! 89 | Incl = 45.0 90 | Ecut = 400 91 | 92 | # lets now start using optimized nested sampling. 93 | 94 | # set up function which computes the various model components: 95 | # the parameters are: 96 | nonlinear_param_names = ['logNH', 'PhoIndex', 'TORsigma', 'CTKcover', 'kT'] 97 | 98 | # set up a prior transform 99 | PhoIndex_gauss = scipy.stats.truncnorm(loc=1.95, scale=0.15, a=(1.0 - 1.95) / 0.15, b=(3.0 - 1.95) / 0.15) 100 | def nonlinear_param_transform(cube): 101 | params = cube.copy() 102 | params[0] = cube[0] * 6 + 20 # logNH 103 | params[1] = PhoIndex_gauss.ppf(cube[1]) 104 | params[2] = cube[2] * (80 - 7) + 7 105 | params[3] = cube[3] * (0.4) + 0 106 | params[4] = 10**(cube[4] * 2 - 1) # kT 107 | return params 108 | 109 | #component_names = ['pl', 'scat', 'apec'] 110 | linear_param_names = ['Nbkg', 'Npl', 'Nscat', 'Napec'] 111 | #for k in data_sets.keys(): 112 | # for name in component_names + ['bkg']: 113 | # linear_param_names.append(f'norm_{name}_{k}') 114 | 115 | bkg_deviations = 0.2 116 | src_deviations = 0.1 117 | 118 | Nlinear = len(linear_param_names) 119 | Ndatasets = len(data_sets) 120 | 121 | 122 | class LinkedPredictionPacker: 123 | """Map source and background spectral components to counts, 124 | 125 | Identical components for each dataset. 126 | 127 | pred_counts should look like 128 | pred_counts should look like 129 | component1-norm1: [counts_data1, 0, counts_data2, 0, counts_data3, 0] 130 | component2-norm2: [counts_data1, 0, counts_data2, 0, counts_data3, 0] 131 | component3-norm3: [counts_data1, 0, counts_data2, 0, counts_data3, 0] 132 | background-bkg: [counts_srcbkg1, counts_bkgbkg1, counts_srcbkg2, counts_bkgbkg2, counts_srcbkg3, counts_bkgbkg3] 133 | """ 134 | def __init__(self, data_sets, Ncomponents): 135 | """Initialise.""" 136 | self.data_sets = data_sets 137 | self.width = 0 138 | self.Ncomponents = Ncomponents 139 | self.counts_flat = np.hstack(tuple([ 140 | np.hstack((data['src_region_counts'], data['bkg_region_counts'])) 141 | for k, data in data_sets.items()])) 142 | self.width, = self.counts_flat.shape 143 | 144 | def pack(self, pred_fluxes): 145 | # one row for each normalisation, 146 | pred_counts = np.zeros((self.Ncomponents, self.width)) 147 | # now let's apply the response to each component: 148 | left = 0 149 | for k, data in self.data_sets.items(): 150 | Ndata = data['chan_mask'].sum() 151 | for i, component_spec in enumerate(pred_fluxes[k]): 152 | pred_counts[i, left:left + Ndata] = component_spec 153 | # now look at background in the background region 154 | left += Ndata 155 | for i, component_spec in enumerate(pred_fluxes[k + '_bkg']): 156 | # fill in background 157 | pred_counts[i, left:left + Ndata] = component_spec 158 | left += Ndata 159 | return pred_counts 160 | 161 | def unpack(self, pred_counts): 162 | pred_fluxes = {} 163 | # now let's apply the response to each component: 164 | left = 0 165 | for k, data in self.data_sets.items(): 166 | Ndata = data['chan_mask'].sum() 167 | pred_fluxes[k] = pred_counts[:, left:left + Ndata] 168 | # now look at background in the background region 169 | left += Ndata 170 | pred_fluxes[k + '_bkg'] = pred_counts[:, left:left + Ndata] 171 | left += Ndata 172 | return pred_fluxes 173 | 174 | def prior_prediction_producer(self, nsamples=8): 175 | for i in range(nsamples): 176 | u = np.random.uniform(size=len(statmodel.nonlinear_param_names)) 177 | nonlinear_params = statmodel.nonlinear_param_transform(u) 178 | X = statmodel.compute_model_components(nonlinear_params) 179 | statmodel.statmodel.update_components(X) 180 | norms = statmodel.statmodel.norms() 181 | pred_counts = norms @ X.T 182 | yield nonlinear_params, norms, pred_counts, X 183 | 184 | def posterior_prediction_producer(self, samples, ypred, nsamples=8): 185 | for i, (params, pred_counts) in enumerate(zip(samples, ypred)): 186 | nonlinear_params = params[self.Ncomponents:] 187 | X = statmodel.compute_model_components(nonlinear_params) 188 | statmodel.statmodel.update_components(X) 189 | norms = params[:self.Ncomponents] 190 | assert np.allclose(pred_counts, norms @ X.T), (norms, pred_counts, norms @ X.T) 191 | yield nonlinear_params, norms, pred_counts, X 192 | 193 | def prior_predictive_check_plot(self, ax, unit='counts', nsamples=8): 194 | self.predictive_check_plot(ax, self.prior_prediction_producer(nsamples=nsamples), unit=unit) 195 | 196 | def posterior_predictive_check_plot(self, ax, samples, ypred, unit='counts'): 197 | self.predictive_check_plot(ax, self.posterior_prediction_producer(samples, ypred), unit=unit) 198 | 199 | def predictive_check_plot(self, ax, sample_infos, unit='counts', nsamples=8): 200 | src_factor = 1 201 | bkg_factor = 1 202 | colors = {} 203 | ylo = np.inf 204 | yhi = 0 205 | # now we need to unpack again: 206 | key_first_dataset = next(iter(self.data_sets)) 207 | legend_entries_first_dataset = [] # data, bkg; model components, make them all black 208 | legend_entries_first_dataset_labels = [] 209 | legend_entries_across_dataset = [] # take total component from each data set 210 | legend_entries_across_dataset_labels = [] 211 | markers = 'osp> all_energies[-1]): 363 | raise ValueError(f"Energy range for {k} is out of bounds of all_energies.") 364 | indices = all_energies_indices[k] 365 | # models compute the sum between energy_lo and energy_hi 366 | # so a wider bin needs to sum the entries between its energy_lo and energy_hi. 367 | indices_left = indices[:-1] 368 | indices_right = indices[1:] 369 | results[k] = all_spec_cumsum[indices_right] - all_spec_cumsum[indices_left] 370 | return results 371 | 372 | 373 | # define spectral components 374 | def compute_model_components_simple_unfolded(params): 375 | logNH, PhoIndex, TORsigma, CTKcover, kT = params 376 | 377 | # compute model components for each data set: 378 | apec_components = deduplicated_evaluation(fastxsf.x.apec, pars=[kT, 1.0, redshift]) 379 | 380 | pred_spectra = {} 381 | for k, data in data_sets.items(): 382 | energies = data['energies'] 383 | # first component: a absorbed power law 384 | abspl_component = absAGNs[k](energies=energies, pars=[ 385 | 10**(logNH - 22), PhoIndex, Ecut, TORsigma, CTKcover, Incl]) 386 | 387 | # second component, a copy of the unabsorbed power law 388 | scat_component = fastxsf.x.zpowerlw(energies=energies, pars=[PhoIndex, redshift]) 389 | 390 | # third component, a apec model 391 | # apec_component = np.clip(fastxsf.x.apec(energies=energies, pars=[kT, 1.0, redshift]), 0, None) 392 | apec_component = np.clip(apec_components[k], 0, None) 393 | 394 | background_component = data['bkg_model_src_region'] * data['src_expoarea'] 395 | background_component_bkg_region = data['bkg_model_bkg_region'] * data['bkg_expoarea'] 396 | 397 | pred_spectra[k] = [background_component, abspl_component, scat_component, apec_component] 398 | pred_spectra[k + '_bkg'] = [background_component_bkg_region] 399 | 400 | return pred_spectra 401 | 402 | 403 | # fold each spectral component through the appropriate response 404 | def compute_model_components_simple(params): 405 | pred_spectra = compute_model_components_simple_unfolded(params) 406 | pred_counts = {} 407 | for k, data in data_sets.items(): 408 | pred_counts[k] = list(pred_spectra[k]) 409 | pred_counts[k + '_bkg'] = list(pred_spectra[k + '_bkg']) 410 | src_spectral_components = pred_spectra[k][1:] # skip background 411 | for j, src_spectral_component in enumerate(src_spectral_components): 412 | # now let's apply the response to each component: 413 | pred_counts[k][j + 1] = data['RMF_src'].apply_rmf( 414 | data['ARF'] * galabsos[k] * src_spectral_component)[data['chan_mask']] * data['src_expoarea'] 415 | assert np.all(pred_counts[k][j + 1] >= 0), (k, j + 1) 416 | return pred_counts 417 | 418 | 419 | # faster version, based on precomputed tables 420 | def compute_model_components_precomputed(params): 421 | assert np.isfinite(params).all(), params 422 | logNH, PhoIndex, TORsigma, CTKcover, kT = params 423 | 424 | pred_counts = {} 425 | 426 | for k, data in data_sets.items(): 427 | # compute model components for each data set: 428 | pred_counts[k] = [data['bkg_model_src_region'] * data['src_expoarea']] 429 | pred_counts[k + '_bkg'] = [data['bkg_model_bkg_region'] * data['bkg_expoarea']] 430 | 431 | # first component: a absorbed power law 432 | pred_counts[k].append(absAGN_folded[k](pars=[10**(logNH - 22), PhoIndex, TORsigma, CTKcover])[data['chan_mask']] * data['src_expoarea']) 433 | 434 | # second component, a copy of the unabsorbed power law 435 | pred_counts[k].append(scat_folded[k](PhoIndex)[data['chan_mask']] * data['src_expoarea']) 436 | 437 | # third component, a apec model 438 | pred_counts[k].append(apec_folded[k](kT)[data['chan_mask']] * data['src_expoarea']) 439 | return pred_counts 440 | 441 | 442 | # compute_model_components = compute_model_components_precomputed 443 | compute_model_components = compute_model_components_simple 444 | 445 | 446 | def compute_model_components_intrinsic(params, energies): 447 | logNH, PhoIndex, TORsigma, CTKcover, kT = params 448 | pred_spectra = [] 449 | pred_spectra.append(energies[:-1] * 0) 450 | pred_spectra.append(absAGN(energies=energies, pars=[0.01, PhoIndex, Ecut, TORsigma, CTKcover, Incl, redshift])) 451 | scat = fastxsf.x.zpowerlw(energies=energies, pars=[PhoIndex, redshift]) 452 | pred_spectra.append(scat) 453 | apec = np.clip(fastxsf.x.apec(energies=energies, pars=[kT, 1.0, redshift]), 0, None) 454 | pred_spectra.append(apec) 455 | return pred_spectra 456 | 457 | 458 | def fakeit(data_sets, norms, background=True, rng=np.random, verbose=True): 459 | for k, data in data_sets.items(): 460 | counts = rng.poisson(norms @ np.array(X[k])) 461 | if verbose: 462 | print(f' Expected counts for {k}: {np.sum(norms @ np.array(X[k]))}, actual counts: {counts.sum()}') 463 | # write result into the data set 464 | data['src_region_counts'] = counts 465 | if background: 466 | counts_bkg = rng.poisson(data['bkg_model_bkg_region'] * data['bkg_expoarea']) 467 | data['bkg_region_counts'] = counts_bkg 468 | return data_sets 469 | 470 | for k, data in data_sets.items(): 471 | print(k, 'expoarea:', data['src_expoarea'], data['bkg_expoarea']) 472 | if k.startswith('NuSTAR'): 473 | data['src_expoarea'] *= 50 474 | data['bkg_expoarea'] *= 50 475 | 476 | # choose model parameters 477 | X = compute_model_components([24.5, 2.0, 30.0, 0.4, 0.5]) 478 | pp = LinkedPredictionPacker(data_sets, 4) 479 | counts_model = pp.pack(X) 480 | # make it so that spectra have ~10000 counts each 481 | target_counts = np.array([40, 40000, 4, 400]) 482 | norms = target_counts / counts_model.sum(axis=1) 483 | norms[0] = 1.0 484 | 485 | # let's compute some luminosities 486 | print(f'norms: {norms}') 487 | 488 | # simulate spectra and fill in the counts 489 | print('Expected total counts:', norms @ np.sum(counts_model, axis=1)) 490 | fakeit(data_sets, norms, rng=np.random.RandomState(42)) 491 | 492 | # need a new prediction packer, because data changed 493 | pp = LinkedPredictionPacker(data_sets, 4) 494 | 495 | def compute_model_components_unnamed(params): 496 | return pp.pack(compute_model_components(params)).T 497 | 498 | linear_param_prior_Sigma_offset = np.eye(Nlinear * Ndatasets) * 0 499 | linear_param_prior_Sigma = np.eye(Nlinear * Ndatasets) * 0 500 | for j in range(len(data_sets)): 501 | # for all data-sets, set a parameter prior: 502 | linear_param_prior_Sigma[j * Nlinear + 3, j * Nlinear + 3] = bkg_deviations**-2 503 | # across data-sets set a mutual parameter prior for each normalisation 504 | for k in range(j + 1, len(data_sets)): 505 | linear_param_prior_Sigma[j * Nlinear + 0, k * Nlinear + 0] = src_deviations**-2 506 | linear_param_prior_Sigma[j * Nlinear + 1, k * Nlinear + 1] = src_deviations**-2 507 | linear_param_prior_Sigma[j * Nlinear + 2, k * Nlinear + 2] = src_deviations**-2 508 | # set a prior, apply it only to the first data-set 509 | if j == 0: 510 | # -5 +- 2 for ratio of pl and scat normalisations, only on first data set 511 | linear_param_prior_Sigma_offset[j * Nlinear + 3, j * Nlinear + 3] = -5 512 | linear_param_prior_Sigma[j * Nlinear + 3, j * Nlinear + 3] = 2.0**-2 513 | 514 | 515 | # now for the linear (normalisation) parameters: 516 | 517 | # set up a prior log-probability density function for these linear parameters: 518 | def linear_param_logprior_linked(params): 519 | lognorms = np.log(params) 520 | Npl = lognorms[:, linear_param_names.index('Npl')] 521 | Nscat = lognorms[:, linear_param_names.index('Nscat')] 522 | logp = np.where(Nscat > np.log(0.1) + Npl, -np.inf, 0) 523 | return logp 524 | 525 | def linear_param_logprior_independent(params): 526 | assert np.all(params > 0) 527 | lognorms = np.log(params.reshape((-1, Nlinear, len(data_sets)))) 528 | Npl = lognorms[:, linear_param_names.index('Npl'), :] 529 | Nscat = lognorms[:, linear_param_names.index('Nscat'), :] 530 | #Napec = lognorms[:, component_names.index('apec'), :] 531 | #Nbkg = lognorms[:, component_names.index('bkg'), :] 532 | logp = np.where(Nscat > np.log(0.1) + Npl, -np.inf, 0) 533 | return logp 534 | 535 | # we should be able to handle two cases: 536 | # the model normalisations are identical across data sets (LinkedPredictionPacker) 537 | # in that case, we have only few linear parameters 538 | # the model normalisations are free parameters in each data set (IndependentPredictionPacker) 539 | # in that case, we have many linear parameters, and we add a Gaussian prior on the lognorms deviations across data sets 540 | # in both cases, the variation across normalisations can also be a prior 541 | 542 | # put a prior on the ratio of Npl and Nscat 543 | ratio_mutual_priors = [2, 1, -5, 2.0] 544 | 545 | # lognorms_prior = pp.create_linear_param_prior(ratio_mutual_priors) 546 | #linear_param_prior_Sigma_offset, linear_param_prior_Sigma 547 | #lognorms_prior = GaussianPrior(linear_param_prior_Sigma_offset, linear_param_prior_Sigma) 548 | 549 | 550 | 551 | 552 | # create OptNS object, and give it all of these ingredients, 553 | # as well as our data 554 | 555 | # we will need some glue between OptNS and our dictionaries 556 | statmodel = OptNS( 557 | linear_param_names, nonlinear_param_names, compute_model_components_unnamed, 558 | nonlinear_param_transform, linear_param_logprior_linked, 559 | pp.counts_flat, positive=True) 560 | #statmodel.statmodel.minimize_kwargs['options']['ftol'] = 1e-2 561 | 562 | # prior predictive checks: 563 | fig = plt.figure(figsize=(15, 4)) 564 | pp.prior_predictive_check_plot(fig.gca()) 565 | #plt.legend(ncol=4) 566 | plt.savefig('multispecopt-priorpc-counts.pdf') 567 | plt.close() 568 | 569 | fig = plt.figure(figsize=(15, 4)) 570 | pp.prior_predictive_check_plot(fig.gca(), unit='area') 571 | #plt.legend(ncol=4) 572 | plt.savefig('multispecopt-priorpc.pdf') 573 | plt.close() 574 | print("starting benchmark...") 575 | import time, tqdm 576 | t0 = time.time() 577 | for i in tqdm.trange(1000): 578 | u = np.random.uniform(size=len(statmodel.nonlinear_param_names)) 579 | nonlinear_params = statmodel.nonlinear_param_transform(u) 580 | assert np.isfinite(nonlinear_params).all() 581 | X = statmodel.compute_model_components(nonlinear_params) 582 | assert np.isfinite(X).all() 583 | statmodel.statmodel.update_components(X) 584 | norms = statmodel.statmodel.norms() 585 | assert np.isfinite(norms).all() 586 | pred_counts = norms @ X.T 587 | print('Duration:', (time.time() - t0) / 1000) 588 | 589 | # create a UltraNest sampler from this. You can pass additional arguments like here: 590 | optsampler = statmodel.ReactiveNestedSampler( 591 | log_dir='multispecoptjit', resume=True) 592 | # run the UltraNest optimized sampler on the nonlinear parameter space: 593 | optresults = optsampler.run(max_num_improvement_loops=0, frac_remain=0.5) 594 | optsampler.print_results() 595 | optsampler.plot() 596 | 597 | # now for postprocessing the results, we want to get the full posterior: 598 | # this samples up to 1000 normalisations for each nonlinear posterior sample: 599 | fullsamples, weights, y_preds = statmodel.get_weighted_samples(optresults['samples'][:200], 40) 600 | print(f'Obtained {len(fullsamples)} weighted posterior samples') 601 | 602 | print('weights:', weights, np.nanmin(weights), np.nanmax(weights), np.mean(weights)) 603 | # make a corner plot: 604 | mask = weights > 1e-6 * np.nanmax(weights) 605 | fullsamples_selected = fullsamples[mask,:] 606 | fullsamples_selected[:, :len(linear_param_names)] = np.log10(fullsamples_selected[:, :len(linear_param_names)]) 607 | 608 | print(f'Obtained {mask.sum()} with not minuscule weight.') 609 | fig = corner.corner( 610 | fullsamples_selected, weights=weights[mask], 611 | labels=linear_param_names + nonlinear_param_names, 612 | show_titles=True, quiet=True, 613 | plot_datapoints=False, plot_density=False, 614 | levels=[0.9973, 0.9545, 0.6827, 0.3934], quantiles=[0.15866, 0.5, 0.8413], 615 | contour_kwargs=dict(linestyles=['-','-.',':','--'], colors=['navy','navy','navy','purple']), 616 | color='purple' 617 | ) 618 | plt.savefig('multispecopt-corner.pdf') 619 | plt.close() 620 | 621 | # to obtain equally weighted samples, we resample 622 | # this respects the effective sample size. If you get too few samples here, 623 | # crank up the number just above. 624 | samples, y_pred_samples = statmodel.resample(fullsamples, weights, y_preds) 625 | print(f'Obtained {len(samples)} equally weighted posterior samples') 626 | 627 | # posterior predictive checks: 628 | fig = plt.figure(figsize=(15, 4)) 629 | pp.posterior_predictive_check_plot(fig.gca(), samples[:40], y_pred_samples[:40]) 630 | plt.savefig('multispecopt-ppc-counts.pdf') 631 | plt.close() 632 | 633 | fig = plt.figure(figsize=(15, 4)) 634 | pp.posterior_predictive_check_plot(fig.gca(), samples[:40], y_pred_samples[:40], unit='area') 635 | plt.savefig('multispecopt-ppc.pdf') 636 | plt.close() 637 | 638 | from fastxsf.flux import luminosity, energy_flux 639 | import astropy.units as u 640 | 641 | luminosities = [] 642 | energy_fluxes = [] 643 | luminosities2 = [] 644 | energy_fluxes2 = [] 645 | for i, (params, pred_counts) in enumerate(zip(samples[:40], y_pred_samples[:40])): 646 | norms = params[:Nlinear] 647 | nonlinear_params = params[Nlinear:] 648 | X = compute_model_components_intrinsic(nonlinear_params, data_sets['Chandra']['energies']) 649 | X2 = compute_model_components_intrinsic(nonlinear_params, data_sets['NuSTAR-FPMA']['energies']) 650 | 651 | energy_fluxes.append([energy_flux(Ni * Xi, data_sets['Chandra']['energies'], 2, 8) / (u.erg/u.s/u.cm**2) for Ni, Xi in zip(norms, X)]) 652 | luminosities.append([luminosity(Ni * Xi, data_sets['Chandra']['energies'], 2, 10, z=redshift, cosmo=cosmo) / (u.erg/u.s) for Ni, Xi in zip(norms, X)]) 653 | 654 | energy_fluxes2.append([energy_flux(Ni * Xi, data_sets['NuSTAR-FPMA']['energies'], 2, 8) / (u.erg/u.s/u.cm**2) for Ni, Xi in zip(norms, X2)]) 655 | luminosities2.append([luminosity(Ni * Xi, data_sets['NuSTAR-FPMA']['energies'], 2, 10, z=redshift, cosmo=cosmo) / (u.erg/u.s) for Ni, Xi in zip(norms, X2)]) 656 | 657 | luminosities = np.array(luminosities) 658 | luminosities2 = np.array(luminosities2) 659 | energy_fluxes = np.array(energy_fluxes) 660 | energy_fluxes2 = np.array(energy_fluxes2) 661 | print("Luminosities[erg/s]:") 662 | print(np.mean(luminosities.sum(axis=1)), np.std(luminosities.sum(axis=1))) 663 | print(np.mean(luminosities, axis=0)) 664 | print(np.std(luminosities, axis=0)) 665 | print(np.mean(luminosities2, axis=0)) 666 | print(np.std(luminosities2, axis=0)) 667 | print() 668 | 669 | print("Energy fluxes[erg/s/cm^2]:") 670 | print(np.mean(energy_fluxes.sum(axis=1)), np.std(energy_fluxes.sum(axis=1))) 671 | print(np.mean(energy_fluxes, axis=0)) 672 | print(np.std(energy_fluxes, axis=0)) 673 | print(np.mean(energy_fluxes2, axis=0)) 674 | print(np.std(energy_fluxes2, axis=0)) 675 | print() 676 | 677 | 678 | 679 | 680 | 681 | 682 | 683 | --------------------------------------------------------------------------------