├── .flake8 ├── .github └── workflows │ ├── black.yml │ ├── flake8.yml │ └── pytest.yml ├── .gitignore ├── CHANGES.rst ├── CITATION.cff ├── LICENSE ├── README.md ├── docs ├── aperture.md ├── ffi.md ├── index.md ├── machine.md ├── perturbationmatrix.md ├── perturbationmatrix3d.md ├── superstamp.md ├── test.png ├── test1.png ├── tpf.md ├── tutorials │ ├── Tutorial_10_TPFs.ipynb │ ├── Tutorial_11_TPFs.ipynb │ ├── Tutorial_20_FFI.ipynb │ ├── Tutorial_30_k2ss.ipynb │ ├── Tutorial_31_k2ss.ipynb │ └── shape_models_K2_M67_c5.mp4 └── utils.md ├── mkdocs.yml ├── poetry.lock ├── pyproject.toml ├── src └── psfmachine │ ├── __init__.py │ ├── aperture.py │ ├── ffi.py │ ├── machine.py │ ├── perturbation.py │ ├── superstamp.py │ ├── tpf.py │ ├── utils.py │ └── version.py └── tests ├── data ├── kplr-ffi_ch01_test.fits ├── shape_model_ffi.fits ├── tpf_test_00.fits ├── tpf_test_01.fits ├── tpf_test_02.fits ├── tpf_test_03.fits ├── tpf_test_04.fits ├── tpf_test_05.fits ├── tpf_test_06.fits ├── tpf_test_07.fits ├── tpf_test_08.fits └── tpf_test_09.fits ├── test_ffimachine.py ├── test_machine.py ├── test_perturbation.py ├── test_tpfmachine.py └── test_utils.py /.flake8: -------------------------------------------------------------------------------- 1 | [flake8] 2 | max-line-length = 127 3 | max-complexity = 15 4 | count = True 5 | show-source = True 6 | extend-ignore = E203, E741 7 | -------------------------------------------------------------------------------- /.github/workflows/black.yml: -------------------------------------------------------------------------------- 1 | name: black 2 | 3 | on: [push, pull_request] 4 | 5 | jobs: 6 | lint: 7 | runs-on: ubuntu-latest 8 | steps: 9 | - uses: actions/checkout@v2 10 | - uses: actions/setup-python@v2 11 | -------------------------------------------------------------------------------- /.github/workflows/flake8.yml: -------------------------------------------------------------------------------- 1 | name: flake8 2 | 3 | on: [push, pull_request] 4 | 5 | jobs: 6 | build: 7 | runs-on: ubuntu-latest 8 | name: flake8 9 | steps: 10 | - uses: actions/checkout@v1 11 | - name: Set up Python 3.8 12 | uses: actions/setup-python@v1 13 | with: 14 | python-version: 3.8 15 | - name: Install dependencies 16 | run: | 17 | python -m pip install --upgrade pip 18 | python -m pip install poetry 19 | poetry install 20 | - name: Run flake8 21 | run: | 22 | poetry run flake8 src 23 | poetry run flake8 tests 24 | -------------------------------------------------------------------------------- /.github/workflows/pytest.yml: -------------------------------------------------------------------------------- 1 | name: pytest 2 | 3 | on: [push, pull_request] 4 | 5 | jobs: 6 | build: 7 | 8 | runs-on: ubuntu-latest 9 | strategy: 10 | matrix: 11 | python-version: [3.8, 3.9] 12 | 13 | steps: 14 | - uses: actions/checkout@v2 15 | - name: Set up Python ${{ matrix.python-version }} 16 | uses: actions/setup-python@v2 17 | with: 18 | python-version: ${{ matrix.python-version }} 19 | - name: Install dependencies 20 | run: | 21 | python -m pip install --upgrade pip 22 | python -m pip install poetry 23 | poetry install 24 | - name: Test with pytest 25 | run: | 26 | poetry run pytest src tests 27 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | psfmachine/data/ 2 | 3 | # Byte-compiled / optimized / DLL files 4 | __pycache__/ 5 | */__pycache__/ 6 | *.py[cod] 7 | 8 | # Distribution / packaging 9 | .Python 10 | env/ 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | *.egg-info/ 23 | .installed.cfg 24 | *.egg 25 | 26 | tutorials/.ipynb_checkpoints 27 | .python-version 28 | .ipynb_checkpoints/ 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | junit/ 42 | htmlcov/ 43 | .tox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xmld 49 | *,cover 50 | 51 | # Translations 52 | *.mo 53 | *.pot 54 | 55 | # Django stuff: 56 | *.log 57 | 58 | # Sphinx documentation 59 | docs/source/api/* 60 | 61 | # PyBuilder 62 | target/ 63 | 64 | # PyCharm 65 | .idea 66 | 67 | # Custom additions 68 | pyraf 69 | 70 | # Mac OSX 71 | .DS_Store 72 | 73 | # VIM 74 | .*.swp 75 | .*.swo 76 | 77 | sandbox 78 | mastDownload 79 | .pytest_cache 80 | .vscode 81 | site/ 82 | 83 | # shape models from FFI 84 | src/psfmachine/data 85 | 86 | # Jupyter 87 | docs/tutorials/.virtual_documents 88 | -------------------------------------------------------------------------------- /CHANGES.rst: -------------------------------------------------------------------------------- 1 | 1.1.0 (2021-10-26) 2 | ================== 3 | 4 | New Features 5 | ------------ 6 | 7 | psfmachine.aperture 8 | ^^^^^^^^^^^^^^^^^^^^^^^^ 9 | 10 | - Collection of functions that perform aperture photometry. It uses the PSF model to 11 | define the aperture mask and computes photometry. 12 | 13 | - The aperture masks are can be created by optimizing two flux metrics, completeness 14 | and crowdeness. 15 | 16 | - Functions that defines the flux metrics, completeness and crowdeness, and compute 17 | object centroids using momentum. 18 | 19 | - Many of these functions inputs machine and create new attribute to it. 20 | 21 | - Diagnostic functions that plot the flux metrics for a given object. 22 | 23 | psfmachine.Machine 24 | ^^^^^^^^^^^^^^^^^^ 25 | 26 | - A new method (`compute_aperture_photometry`) that computes aperture photometry using 27 | the new functionalities in `psfmachine.aperture_utils`. 28 | 29 | 30 | psfmachine.TPFMachine 31 | ^^^^^^^^^^^^^^^^^^^^^ 32 | 33 | - A new method that computes the centroid of each sources. 34 | -------------------------------------------------------------------------------- /CITATION.cff: -------------------------------------------------------------------------------- 1 | cff-version: 1.2.0 2 | message: "If you use this software, please cite it as below." 3 | authors: 4 | - family-names: "Hedges" 5 | given-names: "Christina" 6 | orcid: "https://orcid.org/0000-0002-3385-8391" 7 | - family-names: "Martinez-Palomera" 8 | given-names: "Jorge" 9 | orcid: "https://orcid.org/0000-0002-7395-4935" 10 | title: "PSF Machine" 11 | version: 1.0.0 12 | doi: 10.5281/zenodo.4784073 13 | date-released: 2021-04-24 14 | url: "https://github.com/SSDataLab/psfmachine" 15 | preferred-citation: 16 | type: article 17 | authors: 18 | - family-names: "Hedges" 19 | given-names: "Christina" 20 | orcid: "https://orcid.org/0000-0002-3385-8391" 21 | - family-names: "Luger" 22 | given-names: "Rodrigo" 23 | orcid: "https://orcid.org/0000-0002-0296-3826" 24 | - family-names: "Martinez-Palomera" 25 | given-names: "Jorge" 26 | orcid: "https://orcid.org/0000-0002-7395-4935" 27 | - family-names: "Dotson" 28 | given-names: "Jessie" 29 | orcid: "https://orcid.org/0000-0003-4206-5649" 30 | - family-names: "Barentsen" 31 | given-names: "Geert" 32 | orcid: "https://orcid.org/0000-0002-3306-3484" 33 | doi: "10.3847/1538-3881/ac0825" 34 | url: "https://doi.org/10.3847/1538-3881/ac0825" 35 | journal: "The Astronomical Journal" 36 | month: 10 37 | title: "Linearized Field Deblending: Point-spread Function Photometry for Impatient Astronomers" 38 | issue: 3 39 | volume: 162 40 | year: 2021 41 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 The Authors 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # PSFMachine 2 | 3 | *PRF photometry with Kepler* 4 | 5 | Test status 6 | pypi status 7 | 8 | 9 | Check out the [documentation](https://ssdatalab.github.io/psfmachine/). 10 | Check out the [paper](#) 11 | 12 | `PSFMachine` is an open source Python tool for creating models of instrument effective Point Spread Functions (ePSFs), a.k.a Pixel Response Functions (PRFs). These models are then used to fit a scene in a stack of astronomical images. `PSFMachine` is able to quickly derive photometry from stacks of *Kepler* images and separate crowded sources. 13 | 14 | # Installation 15 | 16 | ``` 17 | pip install psfmachine 18 | ``` 19 | 20 | # Example use 21 | 22 | Below is an example script that shows how to use `PSFMachine`. Depending on the speed or your computer fitting this sort of model will probably take ~10 minutes to build 200 light curves. You can speed this up by changing some of the input parameters. 23 | 24 | ```python 25 | import psfmachine as psf 26 | import lightkurve as lk 27 | tpfs = lk.search_targetpixelfile('Kepler-16', mission='Kepler', quarter=12, radius=1000, limit=200, cadence='long').download_all(quality_bitmask=None) 28 | machine = psf.TPFMachine.from_TPFs(tpfs, n_r_knots=10, n_phi_knots=12) 29 | machine.fit_lightcurves() 30 | ``` 31 | 32 | Funding for this project is provided by NASA ROSES grant number 80NSSC20K0874. 33 | -------------------------------------------------------------------------------- /docs/aperture.md: -------------------------------------------------------------------------------- 1 | # Documentation for `Aperture` 2 | 3 | ::: psfmachine.aperture 4 | handler: python 5 | selection: 6 | members: 7 | - optimize_aperture 8 | - goodness_metric_obj_fun 9 | - plot_flux_metric_diagnose 10 | - estimate_source_centroids_aperture 11 | - compute_FLFRCSAP 12 | - compute_CROWDSAP 13 | rendering: 14 | show_root_heading: false 15 | show_source: false 16 | -------------------------------------------------------------------------------- /docs/ffi.md: -------------------------------------------------------------------------------- 1 | # Documentation for `FFIMachine` 2 | 3 | ::: psfmachine.ffi.FFIMachine 4 | handler: python 5 | selection: 6 | members: 7 | - __init__ 8 | - from_file 9 | - save_shape_model 10 | - load_shape_model 11 | - save_flux_values 12 | - plot_image 13 | - plot_pixel_masks 14 | - residuals 15 | rendering: 16 | show_root_heading: false 17 | show_source: false 18 | -------------------------------------------------------------------------------- /docs/index.md: -------------------------------------------------------------------------------- 1 | # PSFMachine 2 | 3 | *PRF photometry with Kepler* 4 | 5 | Test status 6 | pypi status 7 | 8 | 9 | `PSFMachine` is an open source Python tool for creating models of instrument effective Point Spread Functions (ePSFs), a.k.a Pixel Response Functions (PRFs). These models are then used to fit a scene in a stack of astronomical images. `PSFMachine` is able to quickly derive photometry from stacks of *Kepler* images and separate crowded sources. 10 | 11 | # Installation 12 | 13 | ``` 14 | pip install psfmachine 15 | ``` 16 | 17 | # What's happening? 18 | 19 | `PSFMachine` is doing a few things for you, it's 20 | 21 | * using the *Gaia* archive to find all the sources in your images 22 | * building a Pixel Response Function (PRF) model based on all the sources 23 | * fitting the PRF to all the sources in an image stack, to find the best fitting flux of all sources, accounting for crowding and contamination 24 | * creating custom apertures masks that follow the PRF profile to performed aperture photometry 25 | * returning neat, [`lightkurve.LightCurve`](https://docs.lightkurve.org/) objects of all the sources in the images, with all the available meta data. 26 | 27 | # What does it look like? 28 | 29 | After you run `PSFmachine` on a stack of images you can retrieve data like this. Here is an example of a very crowded *Kepler* [Object of Interest](https://exoplanetarchive.ipac.caltech.edu/overview/KOI-608) KOI-608. This object looks like a planet transit, but there are actually two targets almost exactly overlapping each other! You can see this on the right hand side, where an image of the target is shown with the two stars identified by *Gaia* highlighted. 30 | 31 | ![Example PSFMachine Output](test1.png) 32 | 33 | The output of `PSFMachine` is on the right hand side, and shows two light curves, one for each source. `PSFMachine` has separated the light curves of these sources, despite the fact that they are separated by less than a pixel, and the background source is over 2 magnitudes fainter. 34 | 35 | We can flatten and fold these light curves at the transit period to find the following 36 | 37 | ![Example PSFMachine Output, Folded](test.png) 38 | 39 | In black we see the original *Kepler* light curve, which looks like an exoplanet transit. When using `PSFMachine` to split these two sources, we see that the transit is actually around the orange target (which is fainter), and has been significantly diluted. `PSFMachine` is able to separate these sources with high confidence, and rule out the blue source as the origin of the transit. 40 | 41 | 42 | # What can I use it on? 43 | 44 | Currently `PSFMachine` is designed to work with *Kepler* data. The tool should work with *K2* or *TESS* data, but some of our key assumptions may break, and so mileage may vary. More work will be done on `PSFMachine` in the future to better integrate these datasets. 45 | 46 | If you'd like to try using `PSFMachine` on a more generic dataset, you can try the `Machine` class instead of the `TPFMachine` class to work with more generic data. 47 | 48 | # Example use 49 | 50 | Below is an example script that shows how to use `PSFMachine`. Depending on the speed or your computer fitting this sort of model will probably take ~10 minutes to build 200 light curves. You can speed this up by changing some of the input parameters. 51 | 52 | ```python 53 | import psfmachine as psf 54 | import lightkurve as lk 55 | tpfs = lk.search_targetpixelfile('Kepler-16', mission='Kepler', quarter=12, radius=1000, limit=200, cadence='long').download_all(quality_bitmask=None) 56 | machine = psf.TPFMachine.from_TPFs(tpfs, n_r_knots=10, n_phi_knots=12) 57 | machine.fit_lightcurves() 58 | ``` 59 | 60 | Funding for this project is provided by NASA ROSES grant number 80NSSC20K0874. 61 | -------------------------------------------------------------------------------- /docs/machine.md: -------------------------------------------------------------------------------- 1 | # Documentation for `Machine` 2 | 3 | ::: psfmachine.machine.Machine 4 | handler: python 5 | selection: 6 | members: 7 | - __init__ 8 | - build_shape_model 9 | - build_time_model 10 | - plot_shape_model 11 | - plot_time_model 12 | - perturbed_model 13 | - fit_model 14 | - get_psf_metrics 15 | - create_aperture_mask 16 | - compute_aperture_photometry 17 | rendering: 18 | show_root_heading: false 19 | show_source: false 20 | -------------------------------------------------------------------------------- /docs/perturbationmatrix.md: -------------------------------------------------------------------------------- 1 | # Documentation for `PerturbationMatrix` 2 | 3 | ::: psfmachine.perturbation.PerturbationMatrix 4 | handler: python 5 | selection: 6 | members: 7 | - __init__ 8 | - plot 9 | - fit 10 | - model 11 | - bin_func 12 | - pca 13 | rendering: 14 | show_root_heading: false 15 | show_source: false 16 | -------------------------------------------------------------------------------- /docs/perturbationmatrix3d.md: -------------------------------------------------------------------------------- 1 | # Documentation for `PerturbationMatrix3D` 2 | 3 | ::: psfmachine.perturbation.PerturbationMatrix3D 4 | handler: python 5 | selection: 6 | members: 7 | - __init__ 8 | - plot_model 9 | - fit 10 | - model 11 | - pca 12 | rendering: 13 | show_root_heading: false 14 | show_source: false 15 | -------------------------------------------------------------------------------- /docs/superstamp.md: -------------------------------------------------------------------------------- 1 | # Documentation for `SSMachine` 2 | 3 | ::: psfmachine.superstamp.SSMachine 4 | handler: python 5 | selection: 6 | members: 7 | - __init__ 8 | - from_file 9 | - build_frame_shape_model 10 | - fit_frame_model 11 | - fit_lightcurves 12 | - plot_image_interactive 13 | rendering: 14 | show_root_heading: false 15 | show_source: false 16 | -------------------------------------------------------------------------------- /docs/test.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSDataLab/psfmachine/6235cc91f2f62c80508668d3ad0648622862450b/docs/test.png -------------------------------------------------------------------------------- /docs/test1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSDataLab/psfmachine/6235cc91f2f62c80508668d3ad0648622862450b/docs/test1.png -------------------------------------------------------------------------------- /docs/tpf.md: -------------------------------------------------------------------------------- 1 | # Documentation for `TPFMachine` 2 | 3 | ::: psfmachine.tpf.TPFMachine 4 | handler: python 5 | selection: 6 | members: 7 | - __init__ 8 | - from_TPFs 9 | - remove_background_model 10 | - plot_background_model 11 | - fit_lightcurves 12 | - plot_tpf 13 | - lcs_in_tpf 14 | - load_shape_model 15 | - save_shape_model 16 | - get_source_centroids 17 | rendering: 18 | show_root_heading: false 19 | show_source: false 20 | -------------------------------------------------------------------------------- /docs/tutorials/shape_models_K2_M67_c5.mp4: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSDataLab/psfmachine/6235cc91f2f62c80508668d3ad0648622862450b/docs/tutorials/shape_models_K2_M67_c5.mp4 -------------------------------------------------------------------------------- /docs/utils.md: -------------------------------------------------------------------------------- 1 | # Documentation for `Utils` 2 | 3 | ::: psfmachine.utils 4 | handler: python 5 | selection: 6 | members: 7 | - get_gaia_sources 8 | - do_tiled_query 9 | - spline1d 10 | - wrapped_spline 11 | - solve_linear_model 12 | - sparse_lessthan 13 | - threshold_bin 14 | - get_breaks 15 | - gaussian_smooth 16 | - bspline_smooth 17 | rendering: 18 | show_root_heading: false 19 | show_source: false 20 | -------------------------------------------------------------------------------- /mkdocs.yml: -------------------------------------------------------------------------------- 1 | site_name: PSFMachine 2 | nav: 3 | - Home: index.md 4 | - API: 5 | - Machine: machine.md 6 | - TPFMachine : tpf.md 7 | - FFIMachine : ffi.md 8 | - Aperture Photometry : aperture.md 9 | - SSMachine : superstamp.md 10 | - PerturbationMatrix : perturbationmatrix.md 11 | - PerturbationMatrix2D : perturbationmatrix3d.md 12 | - Utils : utils.md 13 | - Tutorials: 14 | - Basic PSFMachine: tutorials/Tutorial_10_TPFs.ipynb 15 | - TPFMachine: tutorials/Tutorial_11_TPFs.ipynb 16 | - FFIMachine: tutorials/Tutorial_20_FFI.ipynb 17 | - Basic SSMachine: tutorials/Tutorial_30_k2ss.ipynb 18 | - SSMachine: tutorials/Tutorial_31_k2ss.ipynb 19 | theme: 20 | name: "material" 21 | icon: 22 | logo: material/star-plus 23 | repo_url: https://github.com/SSDataLab/psfmachine 24 | plugins: 25 | - search 26 | - mkdocs-jupyter: 27 | execute: False 28 | include_source: True 29 | ignore_h1_titles: False 30 | - mkdocstrings: 31 | default_handler: python 32 | handlers: 33 | python: 34 | selection: 35 | docstring_style: "numpy" 36 | rendering: 37 | show_source: false 38 | custom_templates: templates 39 | watch: 40 | - src/psfmachine 41 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.poetry] 2 | name = "psfmachine" 3 | version = "1.1.4" 4 | description = "Tool to perform fast PSF photometry of primary and background targets from Kepler/K2 Target Pixel Files" 5 | authors = ["Christina Hedges ", 6 | "Jorge Martinez-Palomera "] 7 | license = "MIT" 8 | readme = "README.md" 9 | homepage = "https://ssdatalab.github.io/psfmachine/" 10 | repository = "https://github.com/ssdatalab/psfmachine" 11 | keywords = ["NASA, Kepler, Astronomy"] 12 | classifiers = [ 13 | "Intended Audience :: Science/Research", 14 | "Topic :: Scientific/Engineering :: Astronomy", 15 | ] 16 | 17 | [tool.coverage.paths] 18 | source = ["src"] 19 | 20 | 21 | [tool.poetry.dependencies] 22 | python = "^3.8" 23 | numpy = "^1.19.4" 24 | scipy = "^1.5.4" 25 | astropy = "^5.2" 26 | tqdm = "^4.54.0" 27 | matplotlib = "^3.3.3" 28 | patsy = "^0.5.1" 29 | pyia = "^1.2" 30 | lightkurve = "^2.0.4" 31 | corner = "^2.1.0" 32 | fitsio = "^1.1.3" 33 | jedi = "0.17.2" 34 | pandas = "^1.1" 35 | photutils = "^1.1.0" 36 | diskcache = "^5.3.0" 37 | imageio = "^2.9.0" 38 | imageio-ffmpeg = "^0.4.5" 39 | kbackground = "^0.1.9" 40 | fbpca = "^1.0" 41 | ipywidgets = "^7.6.3" 42 | 43 | [tool.poetry.dev-dependencies] 44 | pytest = "^6.1.2" 45 | jupyterlab = "^2.2.2" 46 | line-profiler = "^3.0.2" 47 | memory-profiler = "^0.58.0" 48 | flake8 = "^3.8.1" 49 | black = "^20.8b1" 50 | exoplanet = "^0.4.4" 51 | ipywidgets = "^7.6.3" 52 | mkdocs = "1.2.3" 53 | mkdocs-material = "^7.0.6" 54 | mkdocstrings = {version = "^0.15.0", extras = ["numpy-style"]} 55 | pytkdocs = {version = "^0.11.0", extras = ["numpy-style"]} 56 | mkdocs-jupyter = "^0.17.3" 57 | 58 | 59 | [build-system] 60 | requires = ["poetry-core>=1.0.0"] 61 | build-backend = "poetry.core.masonry.api" 62 | -------------------------------------------------------------------------------- /src/psfmachine/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | from __future__ import absolute_import 4 | 5 | import os 6 | 7 | PACKAGEDIR = os.path.abspath(os.path.dirname(__file__)) 8 | 9 | from .version import __version__ # noqa 10 | from .machine import Machine # noqa 11 | from .tpf import TPFMachine # noqa 12 | from .ffi import FFIMachine # noqa 13 | from .superstamp import SSMachine # noqa 14 | from .utils import solve_linear_model # noqa 15 | -------------------------------------------------------------------------------- /src/psfmachine/aperture.py: -------------------------------------------------------------------------------- 1 | """ 2 | Collection of aperture utils lifted from 3 | [Kepler-Apertures](https://github.com/jorgemarpa/kepler-apertures) and adapted to work 4 | with PSFMachine. 5 | 6 | Some this functions inputs and operate on a `Machine` object but we move them out of 7 | `mahine.py` to keep the latter smowhow clean and short. 8 | 9 | """ 10 | 11 | import numpy as np 12 | import matplotlib.pyplot as plt 13 | from scipy import optimize 14 | from tqdm import tqdm 15 | 16 | 17 | def optimize_aperture( 18 | psf_model, 19 | target_complete=0.9, 20 | target_crowd=0.9, 21 | max_iter=100, 22 | percentile_bounds=[0, 100], 23 | quiet=False, 24 | ): 25 | """ 26 | Function to optimize the aperture mask for a given source. 27 | 28 | The optimization is done using scipy Brent's algorithm and it uses a custom 29 | loss function `goodness_metric_obj_fun` that uses a Leaky ReLU term to 30 | achive the target value for both metrics. 31 | 32 | Parameters 33 | ---------- 34 | psf_model : scipy.sparce.csr_matrix 35 | Sparse matrix with the PSF models for all targets in the scene. It has shape 36 | [n_sources, n_pixels]. 37 | target_complete : float 38 | Value of the target completeness metric. 39 | target_crowd : float 40 | Value of the target crowdeness metric. 41 | max_iter : int 42 | Numer of maximum iterations to be performed by the optimizer. 43 | percentile_bounds : tuple 44 | Tuple of minimun and maximun values for allowed percentile values during 45 | the optimization. Default is the widest range of [0, 100]. 46 | 47 | Returns 48 | ------- 49 | optimal_percentile : numpy.ndarray 50 | An array with the percentile value to defines the "optimal" aperture for 51 | each source. 52 | """ 53 | # optimize percentile cut for every source 54 | optimal_percentile = [] 55 | for sdx in tqdm( 56 | range(psf_model.shape[0]), 57 | desc="Optimizing apertures per source", 58 | disable=quiet, 59 | ): 60 | optim_params = { 61 | "percentile_bounds": percentile_bounds, 62 | "target_complete": target_complete, 63 | "target_crowd": target_crowd, 64 | "max_iter": max_iter, 65 | "psf_models": psf_model, 66 | "sdx": sdx, 67 | } 68 | minimize_result = optimize.minimize_scalar( 69 | goodness_metric_obj_fun, 70 | method="Bounded", 71 | bounds=percentile_bounds, 72 | options={"maxiter": max_iter, "disp": False}, 73 | args=(optim_params), 74 | ) 75 | optimal_percentile.append(minimize_result.x) 76 | return np.array(optimal_percentile) 77 | 78 | 79 | def goodness_metric_obj_fun(percentile, optim_params): 80 | """ 81 | The objective function to minimize with scipy.optimize.minimize_scalar called 82 | during optimization of the photometric aperture. 83 | 84 | Parameters 85 | ---------- 86 | percentile : int 87 | Percentile of the normalized flux distribution that defines the isophote. 88 | optim_params : dictionary 89 | Dictionary with the variables needed to evaluate the metric: 90 | * psf_models 91 | * sdx 92 | * target_complete 93 | * target_crowd 94 | 95 | Returns 96 | ------- 97 | penalty : float 98 | Value of the objective function to be used for optiization. 99 | """ 100 | psf_models = optim_params["psf_models"] 101 | sdx = optim_params["sdx"] 102 | # Find the value where to cut 103 | cut = np.nanpercentile(psf_models[sdx].data, percentile) 104 | # create "isophot" mask with current cut 105 | mask = (psf_models[sdx] > cut).toarray()[0] 106 | 107 | # Do not compute and ignore if target score < 0 108 | if optim_params["target_complete"] > 0: 109 | # compute_FLFRCSAP returns an array of size 1 when doing only one source 110 | completMetric = compute_FLFRCSAP(psf_models[sdx], mask)[0] 111 | else: 112 | completMetric = 1.0 113 | 114 | # Do not compute and ignore if target score < 0 115 | if optim_params["target_crowd"] > 0: 116 | crowdMetric = compute_CROWDSAP(psf_models, mask, idx=sdx) 117 | else: 118 | crowdMetric = 1.0 119 | 120 | # Once we hit the target we want to ease-back on increasing the metric 121 | # However, we don't want to ease-back to zero pressure, that will 122 | # unconstrain the penalty term and cause the optmizer to run wild. 123 | # So, use a "Leaky ReLU" 124 | # metric' = threshold + (metric - threshold) * leakFactor 125 | leakFactor = 0.01 126 | if ( 127 | optim_params["target_complete"] > 0 128 | and completMetric >= optim_params["target_complete"] 129 | ): 130 | completMetric = optim_params["target_complete"] + leakFactor * ( 131 | completMetric - optim_params["target_complete"] 132 | ) 133 | 134 | if optim_params["target_crowd"] > 0 and crowdMetric >= optim_params["target_crowd"]: 135 | crowdMetric = optim_params["target_crowd"] + leakFactor * ( 136 | crowdMetric - optim_params["target_crowd"] 137 | ) 138 | 139 | penalty = -(completMetric + crowdMetric) 140 | 141 | return penalty 142 | 143 | 144 | def plot_flux_metric_diagnose(psf_model, idx=0, ax=None, optimal_percentile=None): 145 | """ 146 | Function to evaluate the flux metrics for a single source as a function of 147 | the parameter that controls the aperture size. 148 | The flux metrics are computed by taking into account the PSF models of 149 | neighbor sources. 150 | 151 | This function is meant to be used only to generate diagnostic figures. 152 | 153 | Parameters 154 | ---------- 155 | psf_model : scipy.sparce.csr_matrix 156 | Sparse matrix with the PSF models for all targets in the scene. It has shape 157 | [n_sources, n_pixels]. 158 | idx : int 159 | Index of the source for which the metrcs will be computed. Has to be a 160 | number between 0 and psf_models.shape[0]. 161 | ax : matplotlib.axes 162 | Axis to be used to plot the figure 163 | 164 | Returns 165 | ------- 166 | ax : matplotlib.axes 167 | Figure axes 168 | """ 169 | compl, crowd, cut = [], [], [] 170 | for p in range(0, 101, 1): 171 | cut.append(p) 172 | mask = (psf_model[idx] >= np.nanpercentile(psf_model[idx].data, p)).toarray()[0] 173 | crowd.append(compute_CROWDSAP(psf_model, mask, idx)) 174 | compl.append(compute_FLFRCSAP(psf_model[idx], mask)) 175 | 176 | if ax is None: 177 | fig, ax = plt.subplots(1) 178 | ax.plot(cut, compl, label=r"FLFRCSAP", c="tab:blue") 179 | ax.plot(cut, crowd, label=r"CROWDSAP", c="tab:green") 180 | if optimal_percentile: 181 | ax.axvline(optimal_percentile, c="tab:red", label="optimal") 182 | ax.set_xlabel("Percentile") 183 | ax.set_ylabel("Metric") 184 | ax.legend() 185 | return ax 186 | 187 | 188 | def estimate_source_centroids_aperture(aperture_mask, flux, column, row): 189 | """ 190 | Computes the centroid via 2D moments methods for all sources all times. It needs 191 | `aperture_mask` to be computed first by runing `compute_aperture_photometry`. 192 | 193 | Parameters 194 | ---------- 195 | aperture_mask : numpy.ndarray 196 | Aperture mask, shape is [n_surces, n_pixels] 197 | flux: numpy.ndarray 198 | Flux values at each pixels and times in units of electrons / sec 199 | column : numpy.ndarray 200 | Data array containing the "columns" of the detector that each pixel is on. 201 | row : numpy.ndarray 202 | Data array containing the "rows" of the detector that each pixel is on. 203 | 204 | Returns 205 | ------- 206 | centroid_col : numpy.ndarray 207 | Column pixel number of the moments centroid, shape is [nsources, ntimes]. 208 | centroid_row : numpy.ndarray 209 | Row pixel number of the moments centroid, shape is [nsources, ntimes]. 210 | """ 211 | centroid_col, centroid_row = [], [] 212 | for idx in range(aperture_mask.shape[0]): 213 | total_flux = np.nansum(flux[:, aperture_mask[idx]], axis=1) 214 | centroid_col.append( 215 | np.nansum( 216 | np.tile(column[aperture_mask[idx]], (flux.shape[0], 1)) 217 | * flux[:, aperture_mask[idx]], 218 | axis=1, 219 | ) 220 | / total_flux 221 | ) 222 | centroid_row.append( 223 | np.nansum( 224 | np.tile(row[aperture_mask[idx]], (flux.shape[0], 1)) 225 | * flux[:, aperture_mask[idx]], 226 | axis=1, 227 | ) 228 | / total_flux 229 | ) 230 | return np.array(centroid_col), np.array(centroid_row) 231 | 232 | 233 | def compute_FLFRCSAP(psf_models, aperture_mask): 234 | """ 235 | Compute fraction of target flux enclosed in the optimal aperture to total flux 236 | for a given source (flux completeness). 237 | Follows definition by Kinemuchi at al. 2012. 238 | 239 | Parameters 240 | ---------- 241 | psf_models : scipy.sparce.csr_matrix 242 | Sparse matrix with the PSF models for all targets in the scene. It has shape 243 | [n_sources, n_pixels]. 244 | aperture_mask: numpy.ndarray 245 | Array of boolean indicating the aperture for the target source. It has shape of 246 | [n_sources, n_pixels]. 247 | 248 | Returns 249 | ------- 250 | FLFRCSAP: numpy.ndarray 251 | Completeness metric 252 | """ 253 | return np.array( 254 | psf_models.multiply(aperture_mask.astype(float)).sum(axis=1) 255 | / psf_models.sum(axis=1) 256 | ).ravel() 257 | 258 | 259 | def compute_CROWDSAP(psf_models, aperture_mask, idx=None): 260 | """ 261 | Compute the ratio of target flux relative to flux from all sources within 262 | the photometric aperture (i.e. 1 - Crowdeness). 263 | Follows definition by Kinemuchi at al. 2012. 264 | 265 | Parameters 266 | ---------- 267 | psf_models : scipy.sparce.csr_matrix 268 | Sparse matrix with the PSF models for all targets in the scene. It has shape 269 | [n_sources, n_pixels]. 270 | aperture_mask : numpy.ndarray 271 | Array of boolean indicating the aperture for the target source. It has shape of 272 | [n_sources, n_pixels]. 273 | idx : int 274 | Source index for what the metric is computed. Value has to be betweeen 0 and 275 | psf_model first dimension size. 276 | If None, it returns the metric for all sources (first dimension of psf_model). 277 | 278 | Returns 279 | ------- 280 | CROWDSAP : numpy.ndarray 281 | Crowdeness metric 282 | """ 283 | ratio = psf_models.multiply(1 / psf_models.sum(axis=0)).tocsr() 284 | if idx is None: 285 | return np.array( 286 | ratio.multiply(aperture_mask.astype(float)).sum(axis=1) 287 | ).ravel() / aperture_mask.sum(axis=1) 288 | else: 289 | return ratio[idx].toarray()[0][aperture_mask].sum() / aperture_mask.sum() 290 | 291 | 292 | def aperture_mask_to_2d(tpfs, sources, aperture_mask, column, row): 293 | """ 294 | Convert 1D aperture mask into 2D to match the shape of TPFs. This 2D aperture 295 | masks are useful to plot them with lightkurve TPF plot. 296 | Because a sources can be in more than one TPF, having 2D array masks per object 297 | with the shape of a single TPF is not possible. 298 | 299 | Parameters 300 | ---------- 301 | tpfs: lightkurve TargetPixelFileCollection 302 | Collection of Target Pixel files 303 | tpfs_meta : list 304 | List of source indices for every TPF in `tpfs`. 305 | aperture_mask : numpy.ndarray 306 | Aperture mask, shape is [n_surces, n_pixels] 307 | column : numpy.ndarray 308 | Data array containing the "columns" of the detector that each pixel is on. 309 | row : numpy.ndarray 310 | Data array containing the "rows" of the detector that each pixel is on. 311 | 312 | Returns 313 | ------- 314 | aperture_mask_2d : dictionary 315 | Is a dictionary with key values as 'TPFindex_SOURCEindex', e.g. a source 316 | (idx=10) with multiple TPF (TPF index 1 and 2) data will look '1_10' and '2_10'. 317 | """ 318 | aperture_mask_2d = {} 319 | for k, tpf in enumerate(tpfs): 320 | # find sources in tpf 321 | sources_in = sources[k] 322 | # row_col pix value of TPF 323 | rc = [ 324 | "%i_%i" % (y, x) 325 | for y in np.arange(tpf.row, tpf.row + tpf.shape[1]) 326 | for x in np.arange(tpf.column, tpf.column + tpf.shape[2]) 327 | ] 328 | # iter sources in the TPF 329 | for sdx in sources_in: 330 | # row_col value of pixels inside aperture 331 | rc_in = [ 332 | "%i_%i" 333 | % ( 334 | row[aperture_mask[sdx]][i], 335 | column[aperture_mask[sdx]][i], 336 | ) 337 | for i in range(aperture_mask[sdx].sum()) 338 | ] 339 | # create initial mask 340 | mask = np.zeros(tpf.shape[1:], dtype=bool).ravel() 341 | # populate mask with True when pixel is inside aperture 342 | mask[np.in1d(rc, rc_in)] = True 343 | mask = mask.reshape(tpf.shape[1:]) 344 | aperture_mask_2d["%i_%i" % (k, sdx)] = mask 345 | 346 | return aperture_mask_2d 347 | -------------------------------------------------------------------------------- /src/psfmachine/ffi.py: -------------------------------------------------------------------------------- 1 | """Subclass of `Machine` that Specifically work with FFIs""" 2 | import os 3 | import logging 4 | import numpy as np 5 | 6 | import matplotlib.pyplot as plt 7 | import matplotlib.colors as colors 8 | 9 | from astropy.io import fits 10 | from astropy.time import Time 11 | from astropy.wcs import WCS 12 | from astropy.stats import SigmaClip 13 | import astropy.units as u 14 | from astropy.stats import sigma_clip 15 | from photutils import Background2D, MedianBackground, BkgZoomInterpolator 16 | 17 | # from . import PACKAGEDIR 18 | from .utils import ( 19 | do_tiled_query, 20 | _make_A_cartesian, 21 | solve_linear_model, 22 | _load_ffi_image, 23 | ) 24 | from .tpf import _clean_source_list 25 | 26 | from .machine import Machine 27 | from .version import __version__ 28 | 29 | log = logging.getLogger(__name__) 30 | __all__ = ["FFIMachine"] 31 | 32 | 33 | class FFIMachine(Machine): 34 | """ 35 | Subclass of Machine for working with FFI data. It is a subclass of Machine 36 | """ 37 | 38 | def __init__( 39 | self, 40 | time, 41 | flux, 42 | flux_err, 43 | ra, 44 | dec, 45 | sources, 46 | column, 47 | row, 48 | wcs=None, 49 | limit_radius=32.0, 50 | n_r_knots=10, 51 | n_phi_knots=15, 52 | time_nknots=10, 53 | time_resolution=200, 54 | time_radius=8, 55 | cut_r=6, 56 | rmin=1, 57 | rmax=16, 58 | sparse_dist_lim=40, 59 | quality_mask=None, 60 | meta=None, 61 | ): 62 | """ 63 | Repeated optional parameters are described in `Machine`. 64 | 65 | Parameters 66 | ---------- 67 | time: numpy.ndarray 68 | Time values in JD 69 | flux: numpy.ndarray 70 | Flux values at each pixels and times in units of electrons / sec. Has shape 71 | [n_times, n_rows, n_columns] 72 | flux_err: numpy.ndarray 73 | Flux error values at each pixels and times in units of electrons / sec. 74 | Has shape [n_times, n_rows, n_columns] 75 | ra: numpy.ndarray 76 | Right Ascension coordinate of each pixel 77 | dec: numpy.ndarray 78 | Declination coordinate of each pixel 79 | sources: pandas.DataFrame 80 | DataFrame with source present in the images 81 | column: np.ndarray 82 | Data array containing the "columns" of the detector that each pixel is on. 83 | row: np.ndarray 84 | Data array containing the "columns" of the detector that each pixel is on. 85 | wcs : astropy.wcs 86 | World coordinates system solution for the FFI. Used for plotting. 87 | quality_mask : np.ndarray or booleans 88 | Boolean array of shape time indicating cadences with bad quality. 89 | meta : dictionary 90 | Meta data information related to the FFI 91 | 92 | Attributes 93 | ---------- 94 | meta : dictionary 95 | Meta data information related to the FFI 96 | wcs : astropy.wcs 97 | World coordinates system solution for the FFI. Used for plotting. 98 | flux_2d : numpy.ndarray 99 | 2D image representation of the FFI, used for plotting. Has shape [n_times, 100 | image_height, image_width] 101 | image_shape : tuple 102 | Shape of 2D image 103 | """ 104 | self.column = column 105 | self.row = row 106 | self.ra = ra 107 | self.dec = dec 108 | self.wcs = wcs 109 | self.meta = meta 110 | 111 | # keep 2d image shape 112 | self.image_shape = flux.shape[1:] 113 | # reshape flux and flux_err as [ntimes, npix] 114 | self.flux = flux.reshape(flux.shape[0], -1) 115 | self.flux_err = flux_err.reshape(flux_err.shape[0], -1) 116 | self.sources = sources 117 | 118 | # remove background 119 | self._remove_background() 120 | 121 | # init `machine` object 122 | super().__init__( 123 | time, 124 | self.flux, 125 | self.flux_err, 126 | self.ra, 127 | self.dec, 128 | self.sources, 129 | self.column, 130 | self.row, 131 | n_r_knots=n_r_knots, 132 | n_phi_knots=n_phi_knots, 133 | time_nknots=time_nknots, 134 | time_resolution=time_resolution, 135 | time_radius=time_radius, 136 | cut_r=cut_r, 137 | rmin=rmin, 138 | rmax=rmax, 139 | sparse_dist_lim=sparse_dist_lim, 140 | ) 141 | self._mask_pixels() 142 | if quality_mask is None: 143 | self.quality_mask = np.zeros(len(time), dtype=int) 144 | else: 145 | self.quality_mask = quality_mask 146 | 147 | @property 148 | def flux_2d(self): 149 | return self.flux.reshape((self.flux.shape[0], *self.image_shape)) 150 | 151 | def __repr__(self): 152 | return f"FFIMachine (N sources, N times, N pixels): {self.shape}" 153 | 154 | @staticmethod 155 | def from_file( 156 | fname, 157 | extension=1, 158 | cutout_size=None, 159 | cutout_origin=[0, 0], 160 | correct_offsets=False, 161 | plot_offsets=False, 162 | magnitude_limit=18, 163 | dr=3, 164 | sources=None, 165 | **kwargs, 166 | ): 167 | """ 168 | Reads data from files and initiates a new object of FFIMachine class. 169 | 170 | Parameters 171 | ---------- 172 | fname : str or list of strings 173 | File name or list of file names of the FFI files. 174 | extension : int 175 | Number of HDU extension to be used, for Kepler FFIs this corresponds to the 176 | channel number. For TESS FFIs, it correspond to the HDU extension containing 177 | the image data (1). 178 | cutout_size : int 179 | Size of the cutout in pixels, assumed to be squared 180 | cutout_origin : tuple of ints 181 | Origin pixel coordinates where to start the cut out. Follows matrix indexing 182 | correct_offsets : boolean 183 | Check and correct for coordinate offset due to wrong WCS. It is off by 184 | default. 185 | plot_offsets : boolean 186 | Create diagnostic plot for oordinate offset correction. 187 | magnitude_limit : float 188 | Limiting magnitude to query Gaia catalog. 189 | dr : int 190 | Gaia data release to be use, default is 3, options are DR2 and EDR3 191 | sources : pandas.DataFrame 192 | Catalog with sources to be extracted by PSFMachine 193 | **kwargs : dictionary 194 | Keyword arguments that defines shape model in a `machine` class object. 195 | See `psfmachine.Machine` for details. 196 | 197 | Returns 198 | ------- 199 | FFIMachine : Machine object 200 | A Machine class object built from the FFI. 201 | """ 202 | # load FITS files and parse arrays 203 | ( 204 | wcs, 205 | time, 206 | flux, 207 | flux_err, 208 | ra, 209 | dec, 210 | column, 211 | row, 212 | metadata, 213 | quality_mask, 214 | ) = _load_file( 215 | fname, 216 | extension=extension, 217 | cutout_size=cutout_size, 218 | cutout_origin=cutout_origin, 219 | ) 220 | 221 | # hardcoded: the grid size to do the Gaia tiled query. This is different for 222 | # cutouts and full channel. TESS and Kepler also need different grid sizes. 223 | if metadata["TELESCOP"] == "Kepler": 224 | ngrid = (2, 2) if flux.shape[1] <= 500 else (4, 4) 225 | else: 226 | ngrid = (4, 4) if flux.shape[1] < 500 else (7, 7) 227 | # query Gaia and clean sources. 228 | if sources is None: 229 | sources = _get_sources( 230 | ra, 231 | dec, 232 | wcs, 233 | magnitude_limit=magnitude_limit, 234 | epoch=time.jyear.mean(), 235 | ngrid=ngrid, 236 | dr=dr, 237 | img_limits=[[row.min(), row.max()], [column.min(), column.max()]], 238 | ) 239 | # correct coordinate offset if necessary. 240 | if correct_offsets: 241 | ra, dec, sources = _check_coordinate_offsets( 242 | ra, 243 | dec, 244 | row, 245 | column, 246 | flux[0], 247 | sources, 248 | wcs, 249 | plot=plot_offsets, 250 | cutout_size=100, 251 | ) 252 | 253 | return FFIMachine( 254 | time.jd, 255 | flux, 256 | flux_err, 257 | ra.ravel(), 258 | dec.ravel(), 259 | sources, 260 | column.ravel(), 261 | row.ravel(), 262 | wcs=wcs, 263 | meta=metadata, 264 | quality_mask=quality_mask, 265 | **kwargs, 266 | ) 267 | 268 | def save_shape_model(self, output=None): 269 | """ 270 | Saves the weights of a PRF fit to disk. 271 | 272 | Parameters 273 | ---------- 274 | output : str, None 275 | Output file name. If None, one will be generated. 276 | """ 277 | # asign a file name 278 | if output is None: 279 | output = "./%s_ffi_shape_model_ext%s_q%s.fits" % ( 280 | self.meta["MISSION"], 281 | str(self.meta["EXTENSION"]), 282 | str(self.meta["QUARTER"]), 283 | ) 284 | log.info(f"File name: {output}") 285 | 286 | # create data structure (DataFrame) to save the model params 287 | table = fits.BinTableHDU.from_columns( 288 | [ 289 | fits.Column( 290 | name="psf_w", 291 | array=self.psf_w / np.log10(self.mean_model_integral), 292 | format="D", 293 | ) 294 | ] 295 | ) 296 | # include metadata and descriptions 297 | table.header["object"] = ("PRF shape", "PRF shape parameters") 298 | table.header["datatype"] = ("FFI", "Type of data used to fit shape model") 299 | table.header["origin"] = ("PSFmachine.FFIMachine", "Software of origin") 300 | table.header["version"] = (__version__, "Software version") 301 | table.header["TELESCOP"] = (self.meta["TELESCOP"], "Telescope name") 302 | table.header["mission"] = (self.meta["MISSION"], "Mission name") 303 | table.header["quarter"] = ( 304 | self.meta["QUARTER"], 305 | "Quarter/Campaign/Sector of observations", 306 | ) 307 | table.header["channel"] = (self.meta["EXTENSION"], "Channel/Camera-CCD output") 308 | table.header["MJD-OBS"] = (self.time[0], "MJD of observation") 309 | table.header["n_rknots"] = ( 310 | self.n_r_knots, 311 | "Number of knots for spline basis in radial axis", 312 | ) 313 | table.header["n_pknots"] = ( 314 | self.n_phi_knots, 315 | "Number of knots for spline basis in angle axis", 316 | ) 317 | table.header["rmin"] = (self.rmin, "Minimum value for knot spacing") 318 | table.header["rmax"] = (self.rmax, "Maximum value for knot spacing") 319 | table.header["cut_r"] = ( 320 | self.cut_r, 321 | "Radial distance to remove angle dependency", 322 | ) 323 | # spline degree is hardcoded in `_make_A_polar` implementation. 324 | table.header["spln_deg"] = (3, "Degree of the spline basis") 325 | table.header["norm"] = (str(False), "Normalized model") 326 | 327 | table.writeto(output, checksum=True, overwrite=True) 328 | 329 | def load_shape_model(self, input=None, plot=False): 330 | """ 331 | Loads a PRF model from disk. 332 | 333 | Parameters 334 | ---------- 335 | input : str, None 336 | Input file name. If None, one will be generated. 337 | plot : boolean 338 | Plot the PRF mean model loaded from disk 339 | """ 340 | if input is None: 341 | raise NotImplementedError( 342 | "Loading default model not implemented. Please provide input file." 343 | ) 344 | # check if file exists and is the right format 345 | if not os.path.isfile(input): 346 | raise FileNotFoundError("No shape file: %s" % input) 347 | if not input.endswith(".fits"): 348 | # should use a custom exception for wrong file format 349 | raise ValueError("File format not suported. Please provide a FITS file.") 350 | 351 | # open file 352 | hdu = fits.open(input) 353 | # check if shape parameters are for correct mission, quarter, and channel 354 | if hdu[1].header["MISSION"] != self.meta["MISSION"]: 355 | raise ValueError( 356 | "Wrong shape model: file is for mission '%s'," 357 | % (hdu[1].header["MISSION"]) 358 | + " it should be '%s'." % (self.meta["MISSION"]) 359 | ) 360 | if hdu[1].header["QUARTER"] != self.meta["QUARTER"]: 361 | raise ValueError( 362 | "Wrong shape model: file is for quarter %i," 363 | % (hdu[1].header["QUARTER"]) 364 | + " it should be %i." % (self.meta["QUARTER"]) 365 | ) 366 | if hdu[1].header["CHANNEL"] != self.meta["EXTENSION"]: 367 | raise ValueError( 368 | "Wrong shape model: file is for channel %i," 369 | % (hdu[1].header["CHANNEL"]) 370 | + " it should be %i." % (self.meta["EXTENSION"]) 371 | ) 372 | # load model hyperparameters and weights 373 | self.n_r_knots = hdu[1].header["n_rknots"] 374 | self.n_phi_knots = hdu[1].header["n_pknots"] 375 | self.rmin = hdu[1].header["rmin"] 376 | self.rmax = hdu[1].header["rmax"] 377 | self.cut_r = hdu[1].header["cut_r"] 378 | self.psf_w = hdu[1].data["psf_w"] 379 | # read from header if weights come from a normalized model. 380 | self.normalized_shape_model = ( 381 | True if hdu[1].header.get("norm") in ["True", "T", 1] else False 382 | ) 383 | del hdu 384 | 385 | # create mean model, but PRF shapes from FFI are in pixels! and TPFMachine 386 | # work in arcseconds 387 | self._get_mean_model() 388 | # remove background pixels and recreate mean model 389 | self._update_source_mask_remove_bkg_pixels() 390 | 391 | if plot: 392 | return self.plot_shape_model() 393 | return 394 | 395 | def save_flux_values(self, output=None, format="fits"): 396 | """ 397 | Saves the flux values of all sources to a file. For FITS output files a multi- 398 | extension file is created with each extension containing a single cadence/frame. 399 | 400 | Parameters 401 | ---------- 402 | output : str, None 403 | Output file name. If None, one will be generated. 404 | format : str 405 | Format of the output file. Only FITS is supported for now. 406 | """ 407 | # check if model was fitted 408 | if not hasattr(self, "ws"): 409 | self.fit_model(fit_va=False) 410 | 411 | # asign default output file name 412 | if output is None: 413 | output = "./%s_source_catalog_ext%s_q%s_mjd%s.fits" % ( 414 | self.meta["MISSION"], 415 | str(self.meta["EXTENSION"]), 416 | str(self.meta["QUARTER"]), 417 | str(self.time[0]), 418 | ) 419 | log.info(f"File name: {output}") 420 | 421 | primary_hdu = fits.PrimaryHDU() 422 | primary_hdu.header["object"] = ("Photometric Catalog", "Photometry") 423 | primary_hdu.header["origin"] = ("PSFmachine.FFIMachine", "Software of origin") 424 | primary_hdu.header["version"] = (__version__, "Software version") 425 | primary_hdu.header["TELESCOP"] = (self.meta["TELESCOP"], "Telescope") 426 | primary_hdu.header["mission"] = (self.meta["MISSION"], "Mission name") 427 | primary_hdu.header["DCT_TYPE"] = (self.meta["DCT_TYPE"], "Data type") 428 | primary_hdu.header["quarter"] = ( 429 | self.meta["QUARTER"], 430 | "Quarter/Campaign/Sector of observations", 431 | ) 432 | primary_hdu.header["channel"] = ( 433 | self.meta["EXTENSION"], 434 | "Channel/Camera-CCD output", 435 | ) 436 | primary_hdu.header["aperture"] = ("PSF", "Type of photometry") 437 | primary_hdu.header["N_OBS"] = (self.time.shape[0], "Number of cadences") 438 | primary_hdu.header["DATSETNM"] = (self.meta["DATSETNM"], "data set name") 439 | primary_hdu.header["RADESYS"] = ( 440 | self.meta["RADESYS"], 441 | "reference frame of celestial coordinates", 442 | ) 443 | primary_hdu.header["EQUINOX"] = ( 444 | self.meta["EQUINOX"], 445 | "equinox of celestial coordinate system", 446 | ) 447 | hdul = fits.HDUList([primary_hdu]) 448 | # create bin table with photometry 449 | for k in range(self.time.shape[0]): 450 | id_col = fits.Column( 451 | name="gaia_id", array=self.sources.designation, format="29A" 452 | ) 453 | ra_col = fits.Column( 454 | name="ra", array=self.sources.ra, format="D", unit="deg" 455 | ) 456 | dec_col = fits.Column( 457 | name="dec", array=self.sources.dec, format="D", unit="deg" 458 | ) 459 | flux_col = fits.Column( 460 | name="psf_flux", array=self.ws[k, :], format="D", unit="-e/s" 461 | ) 462 | flux_err_col = fits.Column( 463 | name="psf_flux_err", array=self.werrs[k, :], format="D", unit="-e/s" 464 | ) 465 | table_hdu = fits.BinTableHDU.from_columns( 466 | [id_col, ra_col, dec_col, flux_col, flux_err_col] 467 | ) 468 | table_hdu.header["EXTNAME"] = "CATALOG" 469 | table_hdu.header["MJD-OBS"] = (self.time[k], "MJD of observation") 470 | 471 | hdul.append(table_hdu) 472 | 473 | hdul.writeto(output, checksum=True, overwrite=True) 474 | 475 | return 476 | 477 | def _remove_background(self, mask=None): 478 | """ 479 | Background removal. It models the background using a median estimator, rejects 480 | flux values with sigma clipping. It modiffies the attributes `flux` and 481 | `flux_2d`. The background model are stored in the `background_model` attribute. 482 | 483 | Parameters 484 | ---------- 485 | mask : numpy.ndarray of booleans 486 | Mask to reject pixels containing sources. Default None. 487 | """ 488 | # lets keep this one for now while JMP opens a PR in kbackground with a simpler 489 | # solution 490 | if not self.meta["BACKAPP"]: 491 | # model background for all cadences 492 | self.bkg_model = np.array( 493 | [ 494 | Background2D( 495 | flux_2d, 496 | mask=mask, 497 | box_size=(64, 50), 498 | filter_size=15, 499 | exclude_percentile=20, 500 | sigma_clip=SigmaClip(sigma=3.0, maxiters=5), 501 | bkg_estimator=MedianBackground(), 502 | interpolator=BkgZoomInterpolator(order=3), 503 | ).background 504 | for flux_2d in self.flux_2d 505 | ] 506 | ) 507 | # substract background 508 | flux_2d_corr = self.flux_2d - self.bkg_model 509 | # flatten flix image 510 | 511 | self.flux = flux_2d_corr.reshape(self.flux_2d.shape[0], -1) 512 | self.meta["BACKAPP"] = True 513 | return 514 | 515 | def _saturated_pixels_mask(self, saturation_limit=1.5e5, tolerance=3): 516 | """ 517 | Finds and removes saturated pixels, including bleed columns. 518 | 519 | Parameters 520 | ---------- 521 | saturation_limit : foat 522 | Saturation limit at which pixels are removed. 523 | tolerance : float 524 | Number of pixels masked around the saturated pixel, remove bleeding. 525 | 526 | Returns 527 | ------- 528 | mask : numpy.ndarray 529 | Boolean mask with rejected pixels 530 | """ 531 | # Which pixels are saturated 532 | # this nanpercentile takes forever to compute for a single cadance ffi 533 | # saturated = np.nanpercentile(self.flux, 99, axis=0) 534 | # assume we'll use ffi for 1 single cadence 535 | saturated = np.where(self.flux > saturation_limit)[1] 536 | # Find bad pixels, including allowence for a bleed column. 537 | bad_pixels = np.vstack( 538 | [ 539 | np.hstack( 540 | [ 541 | self.column[saturated] + idx 542 | for idx in np.arange(-tolerance, tolerance) 543 | ] 544 | ), 545 | np.hstack( 546 | [self.row[saturated] for idx in np.arange(-tolerance, tolerance)] 547 | ), 548 | ] 549 | ).T 550 | # Find unique row/column combinations 551 | bad_pixels = bad_pixels[ 552 | np.unique(["".join(s) for s in bad_pixels.astype(str)], return_index=True)[ 553 | 1 554 | ] 555 | ] 556 | # Build a mask of saturated pixels 557 | m = np.zeros(len(self.column), bool) 558 | # this works for FFIs but is slow 559 | for p in bad_pixels: 560 | m |= (self.column == p[0]) & (self.row == p[1]) 561 | 562 | saturated = (self.flux > saturation_limit)[0] 563 | return m 564 | 565 | def _bright_sources_mask(self, magnitude_limit=8, tolerance=30): 566 | """ 567 | Finds and mask pixels with halos produced by bright stars (e.g. <8 mag). 568 | 569 | Parameters 570 | ---------- 571 | magnitude_limit : foat 572 | Magnitude limit at which bright sources are identified. 573 | tolerance : float 574 | Radius limit (in pixels) at which pixels around bright sources are masked. 575 | 576 | Returns 577 | ------- 578 | mask : numpy.ndarray 579 | Boolean mask with rejected pixels 580 | """ 581 | bright_mask = self.sources["phot_g_mean_mag"] <= magnitude_limit 582 | 583 | mask = [ 584 | np.hypot(self.column - s.column, self.row - s.row) < tolerance 585 | for _, s in self.sources[bright_mask].iterrows() 586 | ] 587 | mask = np.array(mask).sum(axis=0) > 0 588 | 589 | return mask 590 | 591 | def _remove_bad_pixels_from_source_mask(self): 592 | """ 593 | Combines source_mask and uncontaminated_pixel_mask with saturated and bright 594 | pixel mask. 595 | """ 596 | self.source_mask = self.source_mask.multiply(self.pixel_mask).tocsr() 597 | self.source_mask.eliminate_zeros() 598 | self.uncontaminated_source_mask = self.uncontaminated_source_mask.multiply( 599 | self.pixel_mask 600 | ).tocsr() 601 | self.uncontaminated_source_mask.eliminate_zeros() 602 | 603 | def _mask_pixels(self, pixel_saturation_limit=1.2e5, magnitude_bright_limit=8): 604 | """ 605 | Mask saturated pixels and halo/difraction pattern from bright sources. 606 | 607 | Parameters 608 | ---------- 609 | pixel_saturation_limit: float 610 | Flux value at which pixels saturate. 611 | magnitude_bright_limit: float 612 | Magnitude limit for sources at which pixels are masked. 613 | """ 614 | 615 | # mask saturated pixels. 616 | self.non_sat_pixel_mask = ~self._saturated_pixels_mask( 617 | saturation_limit=pixel_saturation_limit 618 | ) 619 | # tolerance dependens on pixel scale, TESS pixels are 5 times larger than TESS 620 | tolerance = 5 if self.meta["MISSION"] == "TESS" else 25 621 | self.non_bright_source_mask = ~self._bright_sources_mask( 622 | magnitude_limit=magnitude_bright_limit, tolerance=tolerance 623 | ) 624 | self.pixel_mask = self.non_sat_pixel_mask & self.non_bright_source_mask 625 | 626 | if not hasattr(self, "source_mask"): 627 | self._get_source_mask() 628 | # include saturated pixels in the source mask and uncontaminated mask 629 | self._remove_bad_pixels_from_source_mask() 630 | 631 | return 632 | 633 | def _get_source_mask( 634 | self, 635 | upper_radius_limit=28.0, 636 | lower_radius_limit=4.5, 637 | upper_flux_limit=2e5, 638 | lower_flux_limit=100, 639 | correct_centroid_offset=True, 640 | plot=False, 641 | ): 642 | """ 643 | Adapted version of `machine._get_source_mask()` that masks out saturated and 644 | bright halo pixels in FFIs. See parameter descriptions in `Machine`. 645 | """ 646 | super()._get_source_mask( 647 | upper_radius_limit=upper_radius_limit, 648 | lower_radius_limit=lower_radius_limit, 649 | upper_flux_limit=upper_flux_limit, 650 | lower_flux_limit=lower_flux_limit, 651 | correct_centroid_offset=correct_centroid_offset, 652 | plot=plot, 653 | ) 654 | self._remove_bad_pixels_from_source_mask() 655 | 656 | def build_shape_model( 657 | self, plot=False, flux_cut_off=1, frame_index="mean", bin_data=False, **kwargs 658 | ): 659 | """ 660 | Adapted version of `machine.build_shape_model()` that masks out saturated and 661 | bright halo pixels in FFIs. See parameter descriptions in `Machine`. 662 | """ 663 | # call method from super calss `machine` 664 | super().build_shape_model( 665 | plot=False, 666 | flux_cut_off=flux_cut_off, 667 | frame_index=frame_index, 668 | bin_data=bin_data, 669 | **kwargs, 670 | ) 671 | # include sat/halo pixels again into source_mask 672 | self._remove_bad_pixels_from_source_mask() 673 | if plot: 674 | return self.plot_shape_model(frame_index=frame_index, bin_data=bin_data) 675 | 676 | def residuals(self, plot=False, zoom=False, metric="residuals"): 677 | """ 678 | Get the residuals (model - image) and compute statistics. It creates a model 679 | of the full image using the `mean_model` and the weights computed when fitting 680 | the shape model. 681 | 682 | Parameters 683 | ---------- 684 | plot : bool 685 | Do plotting. 686 | zoom : bool 687 | If plot is True then zoom into a section of the image for better 688 | visualization. 689 | metric : string 690 | Type of metric used to plot. Default is "residuals", "chi2" is also 691 | available. 692 | 693 | Return 694 | ------ 695 | fig : matplotlib figure 696 | Figure. 697 | """ 698 | if not hasattr(self, "ws"): 699 | self.fit_model(fit_va=False) 700 | 701 | # evaluate mean model 702 | ffi_model = self.mean_model.T.dot(self.ws[0]) 703 | ffi_model_err = self.mean_model.T.dot(self.werrs[0]) 704 | # compute residuals 705 | residuals = ffi_model - self.flux[0] 706 | weighted_chi = (ffi_model - self.flux[0]) ** 2 / ffi_model_err 707 | # mask background 708 | source_mask = ffi_model != 0.0 709 | # rms 710 | self.rms = np.sqrt((residuals[source_mask] ** 2).mean()) 711 | self.frac_esidual_median = np.median( 712 | residuals[source_mask] / self.flux[0][source_mask] 713 | ) 714 | self.frac_esidual_std = np.std( 715 | residuals[source_mask] / self.flux[0][source_mask] 716 | ) 717 | 718 | if plot: 719 | fig, ax = plt.subplots(2, 2, figsize=(15, 15)) 720 | 721 | ax[0, 0].scatter( 722 | self.column, 723 | self.row, 724 | c=self.flux[0], 725 | marker="s", 726 | s=7.5 if zoom else 1, 727 | norm=colors.SymLogNorm(linthresh=500, vmin=0, vmax=5000, base=10), 728 | ) 729 | ax[0, 0].set_aspect("equal", adjustable="box") 730 | 731 | ax[0, 1].scatter( 732 | self.column, 733 | self.row, 734 | c=ffi_model, 735 | marker="s", 736 | s=7.5 if zoom else 1, 737 | norm=colors.SymLogNorm(linthresh=500, vmin=0, vmax=5000, base=10), 738 | ) 739 | ax[0, 1].set_aspect("equal", adjustable="box") 740 | 741 | if metric == "residuals": 742 | to_plot = residuals 743 | norm = colors.SymLogNorm(linthresh=500, vmin=-5000, vmax=5000, base=10) 744 | cmap = "RdBu" 745 | elif metric == "chi2": 746 | to_plot = weighted_chi 747 | norm = colors.LogNorm(vmin=1, vmax=5000) 748 | cmap = "viridis" 749 | else: 750 | raise ValueError("wrong type of metric") 751 | 752 | cbar = ax[1, 0].scatter( 753 | self.column[source_mask], 754 | self.row[source_mask], 755 | c=to_plot[source_mask], 756 | marker="s", 757 | s=7.5 if zoom else 1, 758 | cmap=cmap, 759 | norm=norm, 760 | ) 761 | ax[1, 0].set_aspect("equal", adjustable="box") 762 | plt.colorbar( 763 | cbar, ax=ax[1, 0], label=r"Flux ($e^{-}s^{-1}$)", fraction=0.042 764 | ) 765 | 766 | ax[1, 1].hist( 767 | residuals[source_mask] / self.flux[0][source_mask], 768 | bins=50, 769 | log=True, 770 | label=( 771 | "RMS (model - data) = %.3f" % self.rms 772 | + "\nMedian = %.3f" % self.frac_esidual_median 773 | + "\nSTD = %3f" % self.frac_esidual_std 774 | ), 775 | ) 776 | ax[1, 1].legend(loc="best") 777 | 778 | ax[0, 0].set_ylabel("Pixel Row Number") 779 | ax[0, 0].set_xlabel("Pixel Column Number") 780 | ax[0, 1].set_xlabel("Pixel Column Number") 781 | ax[1, 0].set_ylabel("Pixel Row Number") 782 | ax[1, 0].set_xlabel("Pixel Column Number") 783 | ax[1, 1].set_xlabel("(model - data) / data") 784 | ax[1, 0].set_title(metric) 785 | 786 | if zoom: 787 | ax[0, 0].set_xlim(self.column.min(), self.column.min() + 100) 788 | ax[0, 0].set_ylim(self.row.min(), self.row.min() + 100) 789 | ax[0, 1].set_xlim(self.column.min(), self.column.min() + 100) 790 | ax[0, 1].set_ylim(self.row.min(), self.row.min() + 100) 791 | ax[1, 0].set_xlim(self.column.min(), self.column.min() + 100) 792 | ax[1, 0].set_ylim(self.row.min(), self.row.min() + 100) 793 | 794 | return fig 795 | return 796 | 797 | def plot_image(self, ax=None, sources=False, frame_index=0): 798 | """ 799 | Function to plot the Full Frame Image and Gaia sources. 800 | 801 | Parameters 802 | ---------- 803 | ax : matplotlib.axes 804 | Matlotlib axis can be provided, if not one will be created and returned. 805 | sources : boolean 806 | Whether to overplot or not the source catalog. 807 | frame_index : int 808 | Time index used to plot the image data. 809 | 810 | Returns 811 | ------- 812 | ax : matplotlib.axes 813 | Matlotlib axis with the figure. 814 | """ 815 | if ax is None: 816 | fig, ax = plt.subplots(1, figsize=(10, 10)) 817 | 818 | ax = plt.subplot(projection=self.wcs) 819 | row_2d, col_2d = np.mgrid[ 820 | self.row.min() : self.row.max() + 1, 821 | self.column.min() : self.column.max() + 1, 822 | ] 823 | im = ax.pcolormesh( 824 | col_2d, 825 | row_2d, 826 | self.flux_2d[frame_index], 827 | cmap=plt.cm.viridis, 828 | shading="nearest", 829 | # origin="lower", 830 | norm=colors.SymLogNorm(linthresh=200, vmin=0, vmax=2000, base=10), 831 | rasterized=True, 832 | ) 833 | plt.colorbar(im, ax=ax, label=r"Flux ($e^{-}s^{-1}$)", fraction=0.042) 834 | 835 | ax.set_title( 836 | "%s %s Ch/CCD %s MJD %f" 837 | % ( 838 | self.meta["MISSION"], 839 | self.meta["OBJECT"] 840 | if "OBJECT" in self.meta.keys() 841 | else self.meta["DCT_TYPE"], 842 | self.meta["EXTENSION"], 843 | self.time[frame_index], 844 | ) 845 | ) 846 | ax.set_xlabel("R.A. [hh:mm]") 847 | ax.set_ylabel("Decl. [deg]") 848 | ax.grid(True, which="major", axis="both", ls="-", color="w", alpha=0.7) 849 | ax.set_xlim(self.column.min() - 2, self.column.max() + 2) 850 | ax.set_ylim(self.row.min() - 2, self.row.max() + 2) 851 | 852 | ax.set_aspect("equal", adjustable="box") 853 | 854 | if sources: 855 | ax.scatter( 856 | self.sources.column, 857 | self.sources.row, 858 | facecolors="none", 859 | edgecolors="r", 860 | linewidths=0.5 if self.sources.shape[0] > 1000 else 1, 861 | alpha=0.9, 862 | ) 863 | return ax 864 | 865 | def plot_pixel_masks(self, ax=None): 866 | """ 867 | Function to plot the mask used to reject saturated and bright pixels. 868 | 869 | Parameters 870 | ---------- 871 | ax : matplotlib.axes 872 | Matlotlib axis can be provided, if not one will be created and returned. 873 | 874 | Returns 875 | ------- 876 | ax : matplotlib.axes 877 | Matlotlib axis with the figure. 878 | """ 879 | row_2d, col_2d = np.mgrid[: self.flux_2d.shape[1], : self.flux_2d.shape[2]] 880 | 881 | if ax is None: 882 | fig, ax = plt.subplots(1, figsize=(10, 10)) 883 | if hasattr(self, "non_bright_source_mask"): 884 | ax.scatter( 885 | col_2d.ravel()[~self.non_bright_source_mask], 886 | row_2d.ravel()[~self.non_bright_source_mask], 887 | c="y", 888 | marker="s", 889 | s=1, 890 | label="bright mask", 891 | ) 892 | if hasattr(self, "non_sat_pixel_mask"): 893 | ax.scatter( 894 | col_2d.ravel()[~self.non_sat_pixel_mask], 895 | row_2d.ravel()[~self.non_sat_pixel_mask], 896 | c="r", 897 | marker="s", 898 | s=1, 899 | label="saturated pixels", 900 | ) 901 | ax.legend(loc="best") 902 | 903 | ax.set_xlabel("Column Pixel Number") 904 | ax.set_ylabel("Row Pixel Number") 905 | ax.set_title("Pixel Mask") 906 | 907 | return ax 908 | 909 | 910 | def _load_file(fname, extension=1, cutout_size=256, cutout_origin=[0, 0]): 911 | """ 912 | Helper function to load FFI files and parse data. It parses the FITS files to 913 | extract the image data and metadata. It checks that all files provided in fname 914 | correspond to FFIs from the same mission. 915 | 916 | Parameters 917 | ---------- 918 | fname : string or list of strings 919 | Name of the FFI files 920 | extension : int 921 | Number of HDU extension to use, for Kepler FFIs this corresponds to the channel 922 | cutout_size: int 923 | Size of (square) portion of FFIs to cut out 924 | cutout_origin: tuple 925 | Coordinates of the origin of the cut out 926 | 927 | Returns 928 | ------- 929 | wcs : astropy.wcs 930 | World coordinates system solution for the FFI. Used to convert RA, Dec to pixels 931 | time : numpy.array 932 | Array with time values in MJD 933 | flux_2d : numpy.ndarray 934 | Array with 2D (image) representation of flux values 935 | flux_err_2d : numpy.ndarray 936 | Array with 2D (image) representation of flux errors 937 | ra_2d : numpy.ndarray 938 | Array with 2D (image) representation of flux RA 939 | dec_2d : numpy.ndarray 940 | Array with 2D (image) representation of flux Dec 941 | col_2d : numpy.ndarray 942 | Array with 2D (image) representation of pixel column 943 | row_2d : numpy.ndarray 944 | Array with 2D (image) representation of pixel row 945 | meta : dict 946 | Dictionary with metadata 947 | """ 948 | if not isinstance(fname, (list, np.ndarray)): 949 | fname = np.sort([fname]) 950 | flux = [] 951 | flux_err = [] 952 | times = [] 953 | telescopes = [] 954 | dct_types = [] 955 | quarters = [] 956 | extensions = [] 957 | quality_mask = [] 958 | cadenceno = [] 959 | for i, f in enumerate(fname): 960 | if not os.path.isfile(f): 961 | raise FileNotFoundError("FFI calibrated fits file does not exist: ", f) 962 | 963 | hdul = fits.open(f, lazy_load_hdus=None) 964 | primary_header = hdul[0].header 965 | hdr = hdul[extension].header 966 | telescopes.append(primary_header["TELESCOP"]) 967 | # kepler 968 | if f.split("/")[-1].startswith("kplr"): 969 | dct_types.append(primary_header["DCT_TYPE"]) 970 | quarters.append(primary_header["QUARTER"]) 971 | extensions.append(hdr["CHANNEL"]) 972 | times.append((hdr["MJDEND"] + hdr["MJDSTART"]) / 2) 973 | # K2 974 | elif f.split("/")[-1].startswith("ktwo"): 975 | dct_types.append(primary_header["DCT_TYPE"]) 976 | quarters.append(primary_header["CAMPAIGN"]) 977 | extensions.append(hdr["CHANNEL"]) 978 | times.append((hdr["MJDEND"] + hdr["MJDSTART"]) / 2) 979 | # TESS 980 | elif f.split("/")[-1].startswith("tess"): 981 | dct_types.append(primary_header["CREATOR"].split(" ")[-1].upper()) 982 | quarters.append(f.split("/")[-1].split("-")[1]) 983 | times.append((hdr["TSTART"] + hdr["TSTOP"]) / 2) 984 | extensions.append("%i.%i" % (hdr["CAMERA"], hdr["CCD"])) 985 | quality_mask.append(hdr["DQUALITY"]) 986 | cadenceno.append(primary_header["FFIINDEX"]) 987 | else: 988 | raise ValueError("FFI is not from Kepler or TESS.") 989 | 990 | if i == 0: 991 | wcs = WCS(hdr) 992 | col_2d, row_2d, f2d = _load_ffi_image( 993 | telescopes[-1], 994 | f, 995 | extension, 996 | cutout_size, 997 | cutout_origin, 998 | return_coords=True, 999 | ) 1000 | flux.append(f2d) 1001 | else: 1002 | flux.append( 1003 | _load_ffi_image( 1004 | telescopes[-1], f, extension, cutout_size, cutout_origin 1005 | ) 1006 | ) 1007 | if telescopes[-1].lower() in ["tess"]: 1008 | flux_err.append( 1009 | _load_ffi_image(telescopes[-1], f, 2, cutout_size, cutout_origin) 1010 | ) 1011 | else: 1012 | flux_err.append(flux[-1] ** 0.5) 1013 | 1014 | # check for integrity of files, same telescope, all FFIs and same quarter/campaign 1015 | if len(set(telescopes)) != 1: 1016 | raise ValueError("All FFIs must be from same telescope") 1017 | if len(set(dct_types)) != 1 or "FFI" not in set(dct_types).pop(): 1018 | raise ValueError("All images must be FFIs") 1019 | if len(set(quarters)) != 1: 1020 | raise ValueError("All FFIs must be of same quarter/campaign/sector.") 1021 | 1022 | # collect meta data, get everthing from one header. 1023 | attrs = [ 1024 | "TELESCOP", 1025 | "INSTRUME", 1026 | "MISSION", 1027 | "DATSETNM", 1028 | ] 1029 | meta = {k: primary_header[k] for k in attrs if k in primary_header.keys()} 1030 | attrs = [ 1031 | "RADESYS", 1032 | "EQUINOX", 1033 | "BACKAPP", 1034 | ] 1035 | meta.update({k: hdr[k] for k in attrs if k in hdr.keys()}) 1036 | # we use "EXTENSION" to combine channel/camera keywords and "QUARTERS" to refer to 1037 | # Kepler quarters and TESS campaigns 1038 | meta.update( 1039 | { 1040 | "EXTENSION": extensions[0], 1041 | "CHANNEL": extensions[0], 1042 | "CAMERA": extensions[0], 1043 | "QUARTER": quarters[0], 1044 | "CAMPAIGN": quarters[0], 1045 | "DCT_TYPE": "FFI", 1046 | } 1047 | ) 1048 | if "MISSION" not in meta.keys(): 1049 | meta["MISSION"] = meta["TELESCOP"] 1050 | 1051 | # sort by times in case fnames aren't 1052 | times = Time(times, format="mjd" if meta["TELESCOP"] == "Kepler" else "btjd") 1053 | tdx = np.argsort(times) 1054 | times = times[tdx] 1055 | if len(quality_mask) == 0: 1056 | quality_mask = np.zeros(len(times), dtype=int) 1057 | cadenceno = np.arange(len(times), dtype=int) 1058 | quality_mask = np.array(quality_mask)[tdx] 1059 | cadenceno = np.array(cadenceno)[tdx] 1060 | flux = np.asarray(flux)[tdx] 1061 | flux_err = np.asarray(flux_err)[tdx] 1062 | 1063 | # convert to RA and Dec 1064 | ra, dec = wcs.all_pix2world(np.vstack([col_2d.ravel(), row_2d.ravel()]).T, 0.0).T 1065 | # some Kepler Channels/Modules have image data but no WCS (e.g. ch 5-8). If the WCS 1066 | # doesn't exist or is wrong, it could produce RA Dec values out of bound. 1067 | if ra.min() < 0.0 or ra.max() > 360 or dec.min() < -90 or dec.max() > 90: 1068 | raise ValueError("WCS lead to out of bound RA and Dec coordinates.") 1069 | ra_2d = ra.reshape(col_2d.shape) 1070 | dec_2d = dec.reshape(col_2d.shape) 1071 | 1072 | del hdul, primary_header, hdr, ra, dec 1073 | 1074 | return ( 1075 | wcs, 1076 | times, 1077 | flux, 1078 | flux_err, 1079 | ra_2d, 1080 | dec_2d, 1081 | col_2d, 1082 | row_2d, 1083 | meta, 1084 | quality_mask, 1085 | ) 1086 | 1087 | 1088 | def _get_sources(ra, dec, wcs, img_limits=[[0, 0], [0, 0]], square=True, **kwargs): 1089 | """ 1090 | Query Gaia catalog in a tiled manner and clean sources off sensor. 1091 | 1092 | Parameters 1093 | ---------- 1094 | ra : numpy.ndarray 1095 | Data array with pixel RA values used to create the grid for tiled query and 1096 | compute centers and radius of cone search 1097 | dec : numpy.ndarray 1098 | Data array with pixel Dec values used to create the grid for tiled query and 1099 | compute centers and radius of cone search 1100 | wcs : astropy.wcs 1101 | World coordinates system solution for the FFI. Used to convert RA, Dec to pixels 1102 | img_limits : 1103 | Image limits in pixel numbers to remove sources outside the CCD. 1104 | square : boolean 1105 | True if the original data is square (e.g. FFIs), False if contains empty pixels 1106 | (e.g. some K2 SuperStamps). This helps to speedup off-sensor source cleaning. 1107 | **kwargs 1108 | Keyword arguments to be passed to `psfmachine.utils.do_tiled_query()`. 1109 | 1110 | Returns 1111 | ------- 1112 | sources : pandas.DataFrame 1113 | Data Frame with query result 1114 | """ 1115 | sources = do_tiled_query(ra, dec, **kwargs) 1116 | sources["column"], sources["row"] = wcs.all_world2pix( 1117 | sources.loc[:, ["ra", "dec"]].values, 0.0 1118 | ).T 1119 | 1120 | if square: 1121 | # remove sources outiside the ccd with a tolerance 1122 | tolerance = 0 1123 | inside = ( 1124 | (sources.row > img_limits[0][0] - tolerance) 1125 | & (sources.row < img_limits[0][1] + tolerance) 1126 | & (sources.column > img_limits[1][0] - tolerance) 1127 | & (sources.column < img_limits[1][1] + tolerance) 1128 | ) 1129 | sources = sources[inside].reset_index(drop=True) 1130 | else: 1131 | sources, _ = _clean_source_list(sources, ra, dec, pixel_tolerance=2) 1132 | return sources 1133 | 1134 | 1135 | def _compute_coordinate_offset(ra, dec, flux, sources, plot=True): 1136 | """ 1137 | Compute coordinate offsets if the RA Dec of objects in source catalog don't align 1138 | with the RA Dec values of the image. 1139 | How it works: first compute dra, ddec and radius of each pixel respect to the 1140 | objects listed in sources. Then masks out all pixels further than ~25 arcsecs around 1141 | each source. It uses spline basis to model the flux as a function of the spatial 1142 | coord and find the scene centroid offsets. 1143 | 1144 | Parameters 1145 | ---------- 1146 | ra : numpy.ndarray 1147 | Data array with pixel RA coordinates. 1148 | dec : numpy.ndarray 1149 | Data array with pixel Dec coordinates. 1150 | flux : numpy.ndarray 1151 | Data array with flux values. 1152 | sources : pandas DataFrame 1153 | Catalog with sources detected in the image. 1154 | plot : boolean 1155 | Create diagnostic plots. 1156 | 1157 | Returns 1158 | ------- 1159 | ra_offset : float 1160 | RA coordinate offset 1161 | dec_offset : float 1162 | Dec coordinate offset 1163 | """ 1164 | # diagnostic plot 1165 | if plot: 1166 | fig, ax = plt.subplots(1, 3, figsize=(15, 4)) 1167 | ax[0].pcolormesh( 1168 | ra, 1169 | dec, 1170 | flux, 1171 | cmap=plt.cm.viridis, 1172 | shading="nearest", 1173 | norm=colors.SymLogNorm(linthresh=200, vmin=0, vmax=2000, base=10), 1174 | rasterized=True, 1175 | ) 1176 | ax[0].scatter( 1177 | sources.ra, 1178 | sources.dec, 1179 | facecolors="none", 1180 | edgecolors="r", 1181 | linewidths=1, 1182 | alpha=0.9, 1183 | ) 1184 | 1185 | # create a temporal mask of ~25 (6 pix) arcsec around each source 1186 | ra, dec, flux = ra.ravel(), dec.ravel(), flux.ravel() 1187 | dra, ddec = np.asarray( 1188 | [ 1189 | [ 1190 | ra - sources["ra"][idx], 1191 | dec - sources["dec"][idx], 1192 | ] 1193 | for idx in range(len(sources)) 1194 | ] 1195 | ).transpose(1, 0, 2) 1196 | dra = dra * (u.deg) 1197 | ddec = ddec * (u.deg) 1198 | r = np.hypot(dra, ddec).to("arcsec") 1199 | source_rad = 0.5 * np.log10(sources.phot_g_mean_flux) ** 1.5 + 25 1200 | tmp_mask = r.value < source_rad.values[:, None] 1201 | flx = np.tile(flux, (sources.shape[0], 1))[tmp_mask] 1202 | 1203 | # design matrix in cartesian coord to model flux(dra, ddec) 1204 | A = _make_A_cartesian( 1205 | dra.value[tmp_mask], 1206 | ddec.value[tmp_mask], 1207 | radius=np.percentile(r[tmp_mask], 90) / 3600, 1208 | n_knots=8, 1209 | ) 1210 | prior_sigma = np.ones(A.shape[1]) * 10 1211 | prior_mu = np.zeros(A.shape[1]) + 10 1212 | w = solve_linear_model( 1213 | A, 1214 | flx, 1215 | y_err=np.sqrt(np.abs(flx)), 1216 | prior_mu=prior_mu, 1217 | prior_sigma=prior_sigma, 1218 | ) 1219 | # iterate to reject outliers from nearby sources using (data - model) 1220 | for k in range(3): 1221 | bad = sigma_clip(flx - A.dot(w), sigma=3).mask 1222 | w = solve_linear_model( 1223 | A, 1224 | flx, 1225 | y_err=np.sqrt(np.abs(flx)), 1226 | k=~bad, 1227 | prior_mu=prior_mu, 1228 | prior_sigma=prior_sigma, 1229 | ) 1230 | # flux model 1231 | flx_mdl = A.dot(w) 1232 | # mask flux values from model to be used as weights 1233 | k = flx_mdl > np.percentile(flx_mdl, 90) 1234 | 1235 | # compute centroid offsets in arcseconds 1236 | ra_offset = np.average(dra[tmp_mask][k], weights=np.sqrt(flx_mdl[k])).to("arcsec") 1237 | dec_offset = np.average(ddec[tmp_mask][k], weights=np.sqrt(flx_mdl[k])).to("arcsec") 1238 | 1239 | # diagnostic plots 1240 | if plot: 1241 | ax[1].scatter( 1242 | dra[tmp_mask] * 3600, 1243 | ddec[tmp_mask] * 3600, 1244 | c=np.log10(flx), 1245 | s=2, 1246 | vmin=2.5, 1247 | vmax=3, 1248 | ) 1249 | 1250 | ax[2].scatter( 1251 | dra[tmp_mask][k] * 3600, 1252 | ddec[tmp_mask][k] * 3600, 1253 | c=np.log10(flx_mdl[k]), 1254 | s=2, 1255 | ) 1256 | ax[1].set_xlim(-30, 30) 1257 | ax[1].set_ylim(-30, 30) 1258 | ax[2].set_xlim(-30, 30) 1259 | ax[2].set_ylim(-30, 30) 1260 | 1261 | ax[1].set_xlabel("R.A.") 1262 | ax[1].set_ylabel("Dec") 1263 | ax[1].set_xlabel(r"$\delta x$") 1264 | ax[1].set_ylabel(r"$\delta y$") 1265 | ax[2].set_xlabel(r"$\delta x$") 1266 | ax[2].set_ylabel(r"$\delta y$") 1267 | 1268 | ax[1].axvline(ra_offset.value, c="r", ls="-") 1269 | ax[1].axhline(dec_offset.value, c="r", ls="-") 1270 | 1271 | plt.show() 1272 | 1273 | return ra_offset, dec_offset 1274 | 1275 | 1276 | def _check_coordinate_offsets( 1277 | ra, dec, row, column, flux, sources, wcs, cutout_size=50, plot=False 1278 | ): 1279 | """ 1280 | Checks if there is any offset between the pixel coordinates and the Gaia sources 1281 | due to wrong WCS. It checks all 4 corners and image center, compute coordinates 1282 | offsets and sees if offsets are consistent in all regions. 1283 | 1284 | Parameters 1285 | ---------- 1286 | ra : numpy.ndarray 1287 | Data array with pixel RA coordinates. 1288 | dec : numpy.ndarray 1289 | Data array with pixel Dec coordinates. 1290 | flux : numpy.ndarray 1291 | Data array with flux values. 1292 | sources : pandas DataFrame 1293 | Catalog with sources detected in the image. 1294 | wcs : astropy.wcs 1295 | World coordinates system solution for the FFI. 1296 | cutout_size : int 1297 | Size of the cutouts in each corner and center to be used to compute offsets. 1298 | Use larger cutouts for regions with low number of sources detected. 1299 | plot : boolean 1300 | Create diagnostic plots. 1301 | 1302 | Returns 1303 | ------- 1304 | ra : numpy.ndarray 1305 | Data arrays with corrected coordinates. 1306 | dec : numpy.ndarray 1307 | Data arrays with corrected coordinates. 1308 | sources : pandas DataFrame 1309 | Catalog with corrected pixel row and column coordinates. 1310 | """ 1311 | # define cutout origins for corners and image center 1312 | cutout_org = [ 1313 | [0, 0], 1314 | [flux.shape[0] - cutout_size, 0], 1315 | [0, flux.shape[1] - cutout_size], 1316 | [flux.shape[0] - cutout_size, flux.shape[1] - cutout_size], 1317 | [(flux.shape[0] - cutout_size) // 2, (flux.shape[1] - cutout_size) // 2], 1318 | ] 1319 | ra_offsets, dec_offsets = [], [] 1320 | # iterate over cutouts to get offsets 1321 | for cdx, c_org in enumerate(cutout_org): 1322 | # create cutouts and sources inside 1323 | cutout_f = flux[ 1324 | c_org[0] : c_org[0] + cutout_size, c_org[1] : c_org[1] + cutout_size 1325 | ] 1326 | cutout_ra = ra[ 1327 | c_org[0] : c_org[0] + cutout_size, c_org[1] : c_org[1] + cutout_size 1328 | ] 1329 | cutout_dec = dec[ 1330 | c_org[0] : c_org[0] + cutout_size, c_org[1] : c_org[1] + cutout_size 1331 | ] 1332 | cutout_row = row[ 1333 | c_org[0] : c_org[0] + cutout_size, c_org[1] : c_org[1] + cutout_size 1334 | ] 1335 | cutout_col = column[ 1336 | c_org[0] : c_org[0] + cutout_size, c_org[1] : c_org[1] + cutout_size 1337 | ] 1338 | inside = ( 1339 | (sources.row > cutout_row.min()) 1340 | & (sources.row < cutout_row.max()) 1341 | & (sources.column > cutout_col.min()) 1342 | & (sources.column < cutout_col.max()) 1343 | ) 1344 | sources_in = sources[inside].reset_index(drop=True) 1345 | 1346 | ra_offset, dec_offset = _compute_coordinate_offset( 1347 | cutout_ra, cutout_dec, cutout_f, sources_in, plot=plot 1348 | ) 1349 | ra_offsets.append(ra_offset.value) 1350 | dec_offsets.append(dec_offset.value) 1351 | 1352 | ra_offsets = np.asarray(ra_offsets) * u.arcsec 1353 | dec_offsets = np.asarray(dec_offsets) * u.arcsec 1354 | 1355 | # diagnostic plot 1356 | if plot: 1357 | plt.plot(ra_offsets, label="RA offset") 1358 | plt.plot(dec_offsets, label="Dec offset") 1359 | plt.legend() 1360 | plt.xlabel("Cutout number") 1361 | plt.ylabel(r"$\delta$ [arcsec]") 1362 | plt.show() 1363 | 1364 | # if offsets are > 1 arcsec and all within 1" from each other, then apply offsets 1365 | # to source coordinates 1366 | if ( 1367 | (np.abs(ra_offsets.mean()) > 1 * u.arcsec) 1368 | and (np.abs(dec_offsets.mean()) > 1 * u.arcsec) 1369 | and (np.abs(ra_offsets - ra_offsets.mean()) < 1 * u.arcsec).all() 1370 | and (np.abs(dec_offsets - dec_offsets.mean()) < 1 * u.arcsec).all() 1371 | ): 1372 | log.info("All offsets are > 1'' and in the same direction") 1373 | # correct the pix coord of sources 1374 | sources["column"], sources["row"] = wcs.all_world2pix( 1375 | np.array( 1376 | [ 1377 | sources.ra + ra_offsets.mean().to("deg"), 1378 | sources.dec + dec_offsets.mean().to("deg"), 1379 | ] 1380 | ).T, 1381 | 0.0, 1382 | ).T 1383 | # correct the ra, dec grid with the offsets 1384 | ra -= ra_offsets.mean().to("deg").value 1385 | dec -= dec_offsets.mean().to("deg").value 1386 | 1387 | return ra, dec, sources 1388 | 1389 | 1390 | def buildKeplerPRFDatabase(fnames): 1391 | """Procedure to build the database of Kepler PRF shape models. 1392 | Parameters 1393 | --------- 1394 | fnames: list of str 1395 | List of filenames for Kepler FFIs. 1396 | """ 1397 | 1398 | # This proceedure should be stored as part of the module, because it will 1399 | # be vital for reproducability. 1400 | 1401 | # 1. Do some basic checks on FFI files that they are Kepler FFIs, and that 1402 | # all 53 are present, all same channel etc. 1403 | 1404 | # 2. Iterate through files 1405 | # for fname in fnames: 1406 | # f = FFIMachine.from_file(fname, HARD_CODED_PARAMETERS) 1407 | # f.build_shape_model() 1408 | # f.fit_model() 1409 | # 1410 | # output = ( 1411 | # PACKAGEDIR 1412 | # + f"src/psfmachine/data/q{quarter}_ch{channel}_{params}.csv" 1413 | # ) 1414 | # f.save_shape_model(output=output) 1415 | raise NotImplementedError 1416 | -------------------------------------------------------------------------------- /src/psfmachine/perturbation.py: -------------------------------------------------------------------------------- 1 | """Classes to deal with perturbation matrices""" 2 | 3 | import numpy as np 4 | import numpy.typing as npt 5 | from typing import Optional 6 | from scipy import sparse 7 | from psfmachine.utils import _make_A_cartesian 8 | import matplotlib.pyplot as plt 9 | from fbpca import pca 10 | from .utils import spline1d 11 | 12 | 13 | class PerturbationMatrix(object): 14 | """ 15 | Class to handle perturbation matrices in PSFMachine 16 | 17 | Parameters 18 | ---------- 19 | time : np.ndarray 20 | Array of time values 21 | other_vectors: list or np.ndarray 22 | Other detrending vectors (e.g. centroids) 23 | poly_order: int 24 | Polynomial order to use for detrending, default 3 25 | focus : bool 26 | Whether to correct focus using a simple exponent model 27 | segments: bool 28 | Whether to fit portions of data where there is a significant time break as separate segments 29 | resolution: int 30 | How many cadences to bin down via `bin_method` 31 | bin_method: str 32 | How to bin the data under the hood. Default is by mean binning. Options are 'downsample' and 'bin' 33 | focus_exptime: float 34 | Time for the exponent for focus change, if used 35 | """ 36 | 37 | def __init__( 38 | self, 39 | time: npt.ArrayLike, 40 | other_vectors: Optional[list] = None, 41 | poly_order: int = 3, 42 | focus: bool = False, 43 | segments: bool = True, 44 | resolution: int = 10, 45 | bin_method: str = "bin", 46 | focus_exptime=2, 47 | ): 48 | 49 | self.time = time 50 | self.other_vectors = np.nan_to_num(other_vectors) 51 | self.poly_order = poly_order 52 | self.focus = focus 53 | self.segments = segments 54 | self.resolution = resolution 55 | self.bin_method = bin_method 56 | self.focus_exptime = focus_exptime 57 | self._vectors = np.vstack( 58 | [ 59 | (self.time - self.time.mean()) ** idx 60 | for idx in range(self.poly_order + 1) 61 | ] 62 | ).T 63 | if self.focus: 64 | self._get_focus_change() 65 | if self.other_vectors is not None: 66 | if isinstance(self.other_vectors, (list, np.ndarray)): 67 | self.other_vectors = np.atleast_2d(self.other_vectors) 68 | if self.other_vectors.shape[0] != len(self.time): 69 | if self.other_vectors.shape[1] == len(self.time): 70 | self.other_vectors = self.other_vectors.T 71 | else: 72 | raise ValueError("Must pass other vectors in the right shape") 73 | else: 74 | raise ValueError("Must pass a list as other vectors") 75 | self.vectors = np.hstack([self._vectors.copy(), self.other_vectors]) 76 | else: 77 | self.vectors = self._vectors.copy() 78 | if self.segments: 79 | self.vectors = self._cut_segments(self.vectors) 80 | # self._clean_vectors() 81 | self.matrix = sparse.csr_matrix(self.bin_func(self.vectors)) 82 | 83 | def __repr__(self): 84 | return "PerturbationMatrix" 85 | 86 | @property 87 | def prior_mu(self): 88 | return np.ones(self.shape[1]) 89 | 90 | @property 91 | def prior_sigma(self): 92 | return np.ones(self.shape[1]) * 0.5 93 | 94 | @property 95 | def breaks(self): 96 | return np.where(np.diff(self.time) / np.median(np.diff(self.time)) > 5)[0] + 1 97 | 98 | @property 99 | def segment_masks(self): 100 | x = np.array_split(np.arange(len(self.time)), self.breaks) 101 | return np.asarray( 102 | [np.in1d(np.arange(len(self.time)), x1).astype(float) for x1 in x] 103 | ).T 104 | 105 | def _cut_segments(self, vectors): 106 | """ 107 | Cuts the data into "segments" wherever there is a break. Breaks are defined 108 | as anywhere where there is a gap in data of more than 5 times the median 109 | time between observations. 110 | 111 | Parameters 112 | ---------- 113 | vectors : np.ndarray 114 | Vector arrays to be break into segments. 115 | """ 116 | return np.hstack( 117 | [ 118 | vectors[:, idx][:, None] * self.segment_masks 119 | for idx in range(vectors.shape[1]) 120 | ] 121 | ) 122 | 123 | def _get_focus_change(self): 124 | """Finds a simple model for the focus change""" 125 | focus = np.asarray( 126 | [ 127 | np.exp(-self.focus_exptime * (self.time - self.time[b])) 128 | for b in np.hstack([0, self.breaks]) 129 | ] 130 | ) 131 | focus *= np.asarray( 132 | [ 133 | ((self.time - self.time[b]) >= 0).astype(float) 134 | for b in np.hstack([0, self.breaks]) 135 | ] 136 | ) 137 | focus[focus < 1e-10] = 0 138 | self._vectors = np.hstack([self._vectors, np.nansum(focus, axis=0)[:, None]]) 139 | return 140 | 141 | # def _clean_vectors(self): 142 | # """Remove time polynomial from other vectors""" 143 | # nvec = self.poly_order + 1 144 | # if self.focus: 145 | # nvec += 1 146 | # if self.segments: 147 | # s = nvec * (len(self.breaks) + 1) 148 | # else: 149 | # s = nvec 150 | # 151 | # if s != self.vectors.shape[1]: 152 | # X = self.vectors[:, :s] 153 | # w = np.linalg.solve(X.T.dot(X), X.T.dot(self.vectors[:, s:])) 154 | # self.vectors[:, s:] -= X.dot(w) 155 | # # Each segment has mean zero 156 | # self.vectors[:, s:] -= np.asarray( 157 | # [v[v != 0].mean() * (v != 0) for v in self.vectors[:, s:].T] 158 | # ).T 159 | # return 160 | 161 | def plot(self): 162 | """Plot basis vectors""" 163 | fig, ax = plt.subplots() 164 | ax.plot(self.time, self.vectors + np.arange(self.vectors.shape[1]) * 0.1) 165 | ax.set(xlabel="Time", ylabel="Vector", yticks=[], title="Vectors") 166 | return fig 167 | 168 | def _fit_linalg(self, y, ye, k=None): 169 | """Hidden method to fit data with linalg""" 170 | if k is None: 171 | k = np.ones(y.shape[0], bool) 172 | X = self.matrix[k] 173 | sigma_w_inv = X.T.dot(X.multiply(1 / ye[k, None] ** 2)) + np.diag( 174 | 1 / self.prior_sigma ** 2 175 | ) 176 | B = X.T.dot(y[k] / ye[k] ** 2) + self.prior_mu / self.prior_sigma ** 2 177 | return np.linalg.solve(sigma_w_inv, B) 178 | 179 | def fit(self, flux: npt.ArrayLike, flux_err: Optional[npt.ArrayLike] = None): 180 | """ 181 | Fits flux to find the best fit model weights. Optionally will include flux errors. 182 | Sets the `self.weights` attribute with best fit weights. 183 | 184 | Parameters 185 | ---------- 186 | flux: npt.ArrayLike 187 | Array of flux values. Should have shape ntimes. 188 | flux: npt.ArrayLike 189 | Optional flux errors. Should have shape ntimes. 190 | 191 | Returns 192 | ------- 193 | weights: npt.ArrayLike 194 | Array with computed weights 195 | """ 196 | if flux_err is None: 197 | flux_err = np.ones_like(flux) 198 | 199 | y, ye = self.bin_func(flux).ravel(), self.bin_func(flux_err, quad=True).ravel() 200 | self.weights = self._fit_linalg(y, ye) 201 | return self.weights 202 | 203 | def model(self, time_indices: Optional[list] = None): 204 | """Returns the best fit model at given `time_indices`. 205 | 206 | Parameters 207 | ---------- 208 | time_indices: list 209 | Optionally pass a list of integers. Model will be evaluated at those indices. 210 | 211 | Returns 212 | ------- 213 | model: npt.ArrayLike 214 | Array of values with the same shape as the `flux` used in `self.fit` 215 | """ 216 | if not hasattr(self, "weights"): 217 | raise ValueError("Run `fit` first.") 218 | if time_indices is None: 219 | time_indices = np.ones(len(self.time), bool) 220 | return self.vectors[time_indices].dot(self.weights) 221 | 222 | @property 223 | def shape(self): 224 | return self.vectors.shape 225 | 226 | @property 227 | def nvec(self): 228 | return self.vectors.shape[1] 229 | 230 | @property 231 | def ntime(self): 232 | return self.time.shape[0] 233 | 234 | def bin_func(self, var, **kwargs): 235 | """ 236 | Bins down an input variable to the same time resolution as `self` 237 | 238 | Parameters 239 | ---------- 240 | var: npt.ArrayLike 241 | Array of values with at least 1 dimension. The first dimension must be 242 | the same shape as `self.time` 243 | 244 | Returns 245 | ------- 246 | func: object 247 | An object function according to `self.bin_method` 248 | """ 249 | if self.bin_method.lower() == "downsample": 250 | func = self._get_downsample_func() 251 | elif self.bin_method.lower() == "bin": 252 | func = self._get_bindown_func() 253 | else: 254 | raise NotImplementedError 255 | return func(var, **kwargs) 256 | 257 | def _get_downsample_func(self): 258 | """Builds a function to lower the resolution of the data through downsampling""" 259 | points = [] 260 | b = np.hstack([0, self.breaks, len(self.time) - 1]) 261 | for b1, b2 in zip(b[:-1], b[1:]): 262 | p = np.arange(b1, b2, self.resolution) 263 | if p[-1] != b2: 264 | p = np.hstack([p, b2]) 265 | points.append(p) 266 | points = np.unique(np.hstack(points)) 267 | self.nbins = len(points) 268 | 269 | def func(x, quad=False): 270 | """ 271 | Bins down an input variable to the same time resolution as `self` 272 | """ 273 | if x.shape[0] == len(self.time): 274 | return x[points] 275 | else: 276 | raise ValueError("Wrong size to bin") 277 | 278 | return func 279 | 280 | def _get_bindown_func(self): 281 | """Builds a function to lower the resolution of the data through binning""" 282 | b = np.hstack([0, self.breaks, len(self.time) - 1]) 283 | points = np.hstack( 284 | [np.arange(b1, b2, self.resolution) for b1, b2 in zip(b[:-1], b[1:])] 285 | ) 286 | points = points[~np.in1d(points, np.hstack([0, len(self.time) - 1]))] 287 | points = np.unique(np.hstack([points, self.breaks])) 288 | self.nbins = len(points) + 1 289 | 290 | def func(x, quad=False): 291 | """ 292 | Bins down an input variable to the same time resolution as `self` 293 | """ 294 | if x.shape[0] == len(self.time): 295 | if not quad: 296 | return np.asarray( 297 | [i.mean(axis=0) for i in np.array_split(x, points)] 298 | ) 299 | else: 300 | return ( 301 | np.asarray( 302 | [ 303 | np.sum(i ** 2, axis=0) / (len(i) ** 2) 304 | for i in np.array_split(x, points) 305 | ] 306 | ) 307 | ** 0.5 308 | ) 309 | else: 310 | raise ValueError("Wrong size to bin") 311 | 312 | return func 313 | 314 | def pca(self, y, ncomponents=5, smooth_time_scale=0): 315 | """Adds the first `ncomponents` principal components of `y` to the design 316 | matrix. `y` is smoothen with a spline function and scale `smooth_time_scale`. 317 | 318 | Parameters 319 | ---------- 320 | y: np.ndarray 321 | Input flux array to take PCA of. 322 | ncomponents: int 323 | Number of principal components to use 324 | smooth_time_scale: float 325 | Amount to smooth the components, using a spline in time. 326 | If 0, the components will not be smoothed. 327 | """ 328 | return self._pca( 329 | y, ncomponents=ncomponents, smooth_time_scale=smooth_time_scale 330 | ) 331 | 332 | def _pca(self, y, ncomponents=3, smooth_time_scale=0): 333 | """This hidden method allows us to update the pca method for other classes""" 334 | if not y.ndim == 2: 335 | raise ValueError("Must pass a 2D `y`") 336 | if not y.shape[0] == len(self.time): 337 | raise ValueError(f"Must pass a `y` with shape ({len(self.time)}, X)") 338 | 339 | # Clean out any time series have significant contribution from one component 340 | k = np.nansum(y, axis=0) != 0 341 | 342 | if smooth_time_scale != 0: 343 | X = sparse.hstack( 344 | [ 345 | spline1d( 346 | self.time, 347 | np.linspace( 348 | self.time[m].min(), 349 | self.time[m].max(), 350 | int( 351 | np.ceil( 352 | (self.time[m].max() - self.time[m].min()) 353 | / smooth_time_scale 354 | ) 355 | ), 356 | ), 357 | degree=3, 358 | )[:, 1:] 359 | for m in self.segment_masks.astype(bool).T 360 | ] 361 | ) 362 | X = sparse.hstack([X, sparse.csr_matrix(np.ones(X.shape[0])).T]).tocsr() 363 | X = X[:, np.asarray(X.sum(axis=0) != 0)[0]] 364 | smoothed_y = X.dot( 365 | np.linalg.solve( 366 | X.T.dot(X).toarray() + np.diag(1 / (np.ones(X.shape[1]) * 1e10)), 367 | X.T.dot(y), 368 | ) 369 | ) 370 | else: 371 | smoothed_y = np.copy(y) 372 | 373 | for count in range(3): 374 | self._pca_components, s, V = pca( 375 | np.nan_to_num(smoothed_y)[:, k], ncomponents, n_iter=30 376 | ) 377 | k[k] &= (np.abs(V) < 0.5).all(axis=0) 378 | 379 | if self.other_vectors is not None: 380 | self.vectors = np.hstack( 381 | [self._vectors, self.other_vectors, self._pca_components] 382 | ) 383 | else: 384 | self.vectors = np.hstack([self._vectors, self._pca_components]) 385 | if self.segments: 386 | self.vectors = self._cut_segments(self.vectors) 387 | # self._clean_vectors() 388 | self.matrix = sparse.csr_matrix(self.bin_func(self.vectors)) 389 | 390 | 391 | class PerturbationMatrix3D(PerturbationMatrix): 392 | """Class to handle 3D perturbation matrices in PSFMachine 393 | 394 | Parameters 395 | ---------- 396 | time : np.ndarray 397 | Array of time values 398 | dx: np.ndarray 399 | Pixel positions in x separation from source center 400 | dy : np.ndaray 401 | Pixel positions in y separation from source center 402 | other_vectors: list or np.ndarray 403 | Other detrending vectors (e.g. centroids) 404 | poly_order: int 405 | Polynomial order to use for detrending, default 3 406 | nknots: int 407 | Number of knots for the cartesian spline 408 | radius: float 409 | Radius out to which to calculate the cartesian spline 410 | focus : bool 411 | Whether to correct focus using a simple exponent model 412 | segments: bool 413 | Whether to fit portions of data where there is a significant time break as 414 | separate segments 415 | resolution: int 416 | How many cadences to bin down via `bin_method` 417 | bin_method: str 418 | How to bin the data under the hood. Default is by mean binning. 419 | focus_exptime: float 420 | Time for the exponent for focus change, if used 421 | degree: int 422 | Polynomial degree used to build the row/column cartesian design matrix 423 | knot_spacing_type: str 424 | Type of spacing bewtwen knots used for cartesian design matrix, options are 425 | {"linear", "sqrt"} 426 | """ 427 | 428 | def __init__( 429 | self, 430 | time: npt.ArrayLike, 431 | dx: npt.ArrayLike, 432 | dy: npt.ArrayLike, 433 | other_vectors: Optional[list] = None, 434 | poly_order: int = 3, 435 | nknots: int = 7, 436 | radius: float = 8, 437 | focus: bool = False, 438 | segments: bool = True, 439 | resolution: int = 30, 440 | bin_method: str = "downsample", 441 | focus_exptime: float = 2, 442 | degree: int = 2, 443 | knot_spacing_type: str = "linear", 444 | ): 445 | self.dx = dx 446 | self.dy = dy 447 | self.nknots = nknots 448 | self.radius = radius 449 | self.degree = degree 450 | self.knot_spacing_type = knot_spacing_type 451 | self.cartesian_matrix = _make_A_cartesian( 452 | self.dx, 453 | self.dy, 454 | n_knots=self.nknots, 455 | radius=self.radius, 456 | knot_spacing_type=self.knot_spacing_type, 457 | degree=self.degree, 458 | ) 459 | super().__init__( 460 | time=time, 461 | other_vectors=other_vectors, 462 | poly_order=poly_order, 463 | focus=focus, 464 | segments=segments, 465 | resolution=resolution, 466 | bin_method=bin_method, 467 | focus_exptime=focus_exptime, 468 | ) 469 | self._get_cartesian_stacked() 470 | 471 | def _get_cartesian_stacked(self): 472 | """ 473 | Stacks cartesian design matrix in preparation to be combined with 474 | time basis vectors. 475 | """ 476 | self._cartesian_stacked = sparse.hstack( 477 | [self.cartesian_matrix for idx in range(self.vectors.shape[1])], 478 | format="csr", 479 | ) 480 | repeat1d = np.repeat( 481 | self.bin_func(self.vectors), self.cartesian_matrix.shape[1], axis=1 482 | ) 483 | repeat2d = np.repeat(repeat1d, self.cartesian_matrix.shape[0], axis=0) 484 | self.matrix = ( 485 | sparse.vstack([self._cartesian_stacked] * self.nbins) 486 | .multiply(repeat2d) 487 | .tocsr() 488 | ) 489 | self.matrix.eliminate_zeros() 490 | 491 | def __repr__(self): 492 | return "PerturbationMatrix3D" 493 | 494 | @property 495 | def shape(self): 496 | return ( 497 | self.cartesian_matrix.shape[0] * self.time.shape[0], 498 | self.cartesian_matrix.shape[1] * self.vectors.shape[1], 499 | ) 500 | 501 | def fit( 502 | self, 503 | flux: npt.ArrayLike, 504 | flux_err: Optional[npt.ArrayLike] = None, 505 | pixel_mask: Optional[npt.ArrayLike] = None, 506 | ): 507 | """ 508 | Fits flux to find the best fit model weights. Optionally will include flux 509 | errors. Sets the `self.weights` attribute with best fit weights. 510 | 511 | Parameters 512 | ---------- 513 | flux: npt.ArrayLike 514 | Array of flux values. Should have shape ntimes x npixels. 515 | flux_err: npt.ArrayLike 516 | Optional flux errors. Should have shape ntimes x npixels. 517 | pixel_mask: npt.ArrayLike 518 | Pixel mask to apply. Values that are `True` will be used in the fit. 519 | Values that are `False` will be masked. Should have shape npixels. 520 | """ 521 | if pixel_mask is not None: 522 | if not isinstance(pixel_mask, np.ndarray): 523 | raise ValueError("`pixel_mask` must be an `np.ndarray`") 524 | if not pixel_mask.shape[0] == flux.shape[-1]: 525 | raise ValueError( 526 | f"`pixel_mask` must be shape {flux.shape[-1]} (npixels)" 527 | ) 528 | else: 529 | pixel_mask = np.ones(flux.shape[-1], bool) 530 | if flux_err is None: 531 | flux_err = np.ones_like(flux) 532 | 533 | y, ye = self.bin_func(flux).ravel(), self.bin_func(flux_err, quad=True).ravel() 534 | k = (np.ones(self.nbins, bool)[:, None] * pixel_mask).ravel() 535 | self.weights = self._fit_linalg(y, ye, k=k) 536 | return 537 | 538 | def model(self, time_indices: Optional[list] = None): 539 | """Returns the best fit model at given `time_indices`. 540 | 541 | Parameters 542 | ---------- 543 | time_indices: list 544 | Optionally pass a list of integers. Model will be evaluated at those indices. 545 | 546 | Returns 547 | ------- 548 | model: npt.ArrayLike 549 | Array of values with the same shape as the `flux` used in `self.fit` 550 | """ 551 | if not hasattr(self, "weights"): 552 | raise ValueError("Run `fit` first") 553 | if time_indices is None: 554 | time_indices = np.arange(len(self.time)) 555 | time_indices = np.atleast_1d(time_indices) 556 | if isinstance(time_indices[0], bool): 557 | time_indices = np.where(time_indices[0])[0] 558 | 559 | return np.asarray( 560 | [ 561 | self._cartesian_stacked.multiply( 562 | np.repeat(self.vectors[time_index], self.cartesian_matrix.shape[1]) 563 | ).dot(self.weights) 564 | for time_index in time_indices 565 | ] 566 | ) 567 | 568 | def pca(self, y, ncomponents=3, smooth_time_scale=0): 569 | """Adds the first `ncomponents` principal components of `y` to the design 570 | matrix. `y` is smoothen with a spline function and scale `smooth_time_scale`. 571 | 572 | Parameters 573 | ---------- 574 | y: np.ndarray 575 | Input flux array to take PCA of. 576 | n_components: int 577 | Number of components to take 578 | smooth_time_scale: float 579 | Amount to smooth the components, using a spline in time. 580 | If 0, the components will not be smoothed. 581 | """ 582 | self._pca( 583 | y, 584 | ncomponents=ncomponents, 585 | smooth_time_scale=smooth_time_scale, 586 | ) 587 | self._get_cartesian_stacked() 588 | 589 | def plot_model(self, time_index=0): 590 | """ 591 | Plot perturbation model 592 | 593 | Parameters 594 | ---------- 595 | time_index : int 596 | Time index to plot the perturbed model 597 | 598 | """ 599 | if not hasattr(self, "weights"): 600 | raise ValueError("Run `fit` first.") 601 | fig, ax = plt.subplots() 602 | ax.scatter(self.dx, self.dy, c=self.model(time_index)[0]) 603 | ax.set( 604 | xlabel=r"$\delta$x", 605 | ylabel=r"$\delta$y", 606 | title=f"Perturbation Model [Cadence {time_index}]", 607 | ) 608 | return fig 609 | -------------------------------------------------------------------------------- /src/psfmachine/superstamp.py: -------------------------------------------------------------------------------- 1 | """Subclass of `Machine` that Specifically work with FFIs""" 2 | import os 3 | import numpy as np 4 | import pandas as pd 5 | import lightkurve as lk 6 | from tqdm import tqdm 7 | 8 | import matplotlib.pyplot as plt 9 | import matplotlib.colors as colors 10 | import imageio 11 | from ipywidgets import interact 12 | 13 | from astropy.io import fits 14 | from astropy.time import Time 15 | from astropy.wcs import WCS 16 | import astropy.units as u 17 | 18 | from .ffi import FFIMachine, _get_sources 19 | from .utils import _do_image_cutout 20 | 21 | __all__ = ["SSMachine"] 22 | 23 | 24 | class SSMachine(FFIMachine): 25 | """ 26 | Subclass of Machine for working with FFI data. It is a subclass of Machine 27 | """ 28 | 29 | def __init__( 30 | self, 31 | time, 32 | flux, 33 | flux_err, 34 | ra, 35 | dec, 36 | sources, 37 | column, 38 | row, 39 | wcs=None, 40 | limit_radius=32.0, 41 | n_r_knots=10, 42 | n_phi_knots=15, 43 | time_nknots=10, 44 | time_resolution=200, 45 | time_radius=8, 46 | rmin=1, 47 | rmax=16, 48 | cut_r=6, 49 | pos_corr1=None, 50 | pos_corr2=None, 51 | meta=None, 52 | ): 53 | """ 54 | Class to work with K2 Supersampts produced by 55 | [Cody et al. 2018](https://archive.stsci.edu/prepds/k2superstamp/) 56 | 57 | Parameters and sttributes are the same as `FFIMachine`. 58 | """ 59 | # init `FFImachine` object 60 | super().__init__( 61 | time, 62 | flux, 63 | flux_err, 64 | ra, 65 | dec, 66 | sources, 67 | column, 68 | row, 69 | wcs=wcs, 70 | limit_radius=limit_radius, 71 | n_r_knots=n_r_knots, 72 | n_phi_knots=n_phi_knots, 73 | time_nknots=time_nknots, 74 | time_resolution=time_resolution, 75 | time_radius=time_radius, 76 | rmin=rmin, 77 | rmax=rmax, 78 | cut_r=cut_r, 79 | meta=meta, 80 | ) 81 | 82 | if pos_corr1 is not None and pos_corr1 is not None: 83 | self.pos_corr1 = np.nan_to_num(np.array(pos_corr1)[None, :]) 84 | self.pos_corr2 = np.nan_to_num(np.array(pos_corr2)[None, :]) 85 | self.time_corrector = "pos_corr" 86 | else: 87 | self.time_corrector = "centroid" 88 | self.poscorr_filter_size = 0 89 | self.meta["DCT_TYPE"] = "SuperStamp" 90 | 91 | def build_frame_shape_model(self, plot=False, **kwargs): 92 | """ 93 | Compute shape model for every cadence (frame) using `Machine.build_shape_model()` 94 | 95 | Parameters 96 | ---------- 97 | plot : boolean 98 | If `True` will create a video file in the working directory with the PSF 99 | model at each frame. It uses `imageio` and `imageio-ffmpeg`. 100 | **kwargs 101 | Keyword arguments to be passed to `build_shape_model()` 102 | """ 103 | self.mean_model_frame = [] 104 | images = [] 105 | self._get_source_mask() 106 | self._get_uncontaminated_pixel_mask() 107 | org_sm = self.source_mask 108 | org_usm = self.uncontaminated_source_mask 109 | for tdx in tqdm(range(self.nt), desc="Building shape model per frame"): 110 | fig = self.build_shape_model(frame_index=tdx, plot=plot, **kwargs) 111 | self.mean_model_frame.append(self.mean_model) 112 | if plot: 113 | fig.canvas.draw() # draw the canvas, cache the render 114 | image = np.frombuffer(fig.canvas.tostring_rgb(), dtype="uint8") 115 | image = image.reshape(fig.canvas.get_width_height()[::-1] + (3,)) 116 | images.append(image) 117 | plt.close() 118 | # we reset the source mask because sources may move slightly, e.g. K2 119 | self.source_mask = org_sm 120 | self.uncontaminated_source_mask = org_usm 121 | if plot: 122 | if hasattr(self, "meta"): 123 | gif_name = "./shape_models_%s_%s_c%i.mp4" % ( 124 | self.meta["MISSION"], 125 | self.meta["OBJECT"].replace(" ", ""), 126 | self.meta["QUARTER"], 127 | ) 128 | else: 129 | gif_name = "./shape_models_%s_q%i.mp4" % ( 130 | self.tpf_meta["mission"][0], 131 | self.tpf_meta["quarter"][0], 132 | ) 133 | imageio.mimsave(gif_name, images, format="mp4", fps=24) 134 | 135 | def fit_frame_model(self): 136 | """ 137 | Fits shape model per frame (cadence). It creates 3 attributes: 138 | * `self.model_flux_frame` has the scene model at every cadence. 139 | * `self.ws_frame` and `self.werrs_frame` have the flux values of all sources 140 | at every cadence. 141 | """ 142 | prior_mu = self.source_flux_estimates # np.zeros(A.shape[1]) 143 | prior_sigma = ( 144 | np.ones(self.mean_model.shape[0]) 145 | * 5 146 | * np.abs(self.source_flux_estimates) ** 0.5 147 | ) 148 | self.model_flux_frame = np.zeros(self.flux.shape) * np.nan 149 | self.ws_frame = np.zeros((self.nt, self.nsources)) 150 | self.werrs_frame = np.zeros((self.nt, self.nsources)) 151 | f = self.flux 152 | fe = self.flux_err 153 | for tdx in tqdm( 154 | range(self.nt), 155 | desc=f"Fitting {self.nsources} Sources (per frame model)", 156 | disable=self.quiet, 157 | ): 158 | X = self.mean_model_frame[tdx].copy().T 159 | sigma_w_inv = X.T.dot(X.multiply(1 / fe[tdx][:, None] ** 2)).toarray() 160 | sigma_w_inv += np.diag(1 / (prior_sigma ** 2)) 161 | B = X.T.dot((f[tdx] / fe[tdx] ** 2)) 162 | B += prior_mu / (prior_sigma ** 2) 163 | self.ws_frame[tdx] = np.linalg.solve(sigma_w_inv, np.nan_to_num(B)) 164 | self.werrs_frame[tdx] = np.linalg.inv(sigma_w_inv).diagonal() ** 0.5 165 | self.model_flux_frame[tdx] = X.dot(self.ws_frame[tdx]) 166 | nodata = np.asarray(self.source_mask.sum(axis=1))[:, 0] == 0 167 | # These sources are poorly estimated 168 | nodata |= (self.mean_model_frame[tdx].max(axis=1) > 1).toarray()[:, 0] 169 | self.ws_frame[tdx, nodata] *= np.nan 170 | self.werrs_frame[tdx, nodata] *= np.nan 171 | 172 | @staticmethod 173 | def from_file( 174 | fname, 175 | magnitude_limit=18, 176 | dr=2, 177 | sources=None, 178 | cutout_size=None, 179 | cutout_origin=[0, 0], 180 | **kwargs, 181 | ): 182 | """ 183 | Reads data from files and initiates a new SSMachine class. SuperStamp file 184 | paths are passed as a string (single frame) or a list of paths (multiple 185 | frames). A samaller cutout of the full SuperSatamp can also be loaded by 186 | passing argumnts `cutout_size` and `cutout_origin`. 187 | 188 | Parameters 189 | ---------- 190 | fname : string or list of strings 191 | Path to the FITS files to be parsed. For only one frame, pass a string, 192 | for multiple frames pass a list of paths. 193 | magnitude_limit : float 194 | Limiting magnitude to query Gaia catalog. 195 | dr : int 196 | Gaia data release to be use, default is 2, options are DR2 and EDR3. 197 | sources : pandas.DataFrame 198 | DataFrame with sources present in the images, optional. If None, then guery 199 | Gaia. 200 | cutout_size : int 201 | Size in pixels of the cutout, assumed to be squared. Default is 100. 202 | cutout_origin : tuple of ints 203 | Origin of the cutout following matrix indexing. Default is [0 ,0]. 204 | **kwargs : dictionary 205 | Keyword arguments that defines shape model in a `Machine` object. 206 | Returns 207 | ------- 208 | SSMachine : Machine object 209 | A Machine class object built from the SuperStamps files. 210 | """ 211 | ( 212 | wcs, 213 | time, 214 | flux, 215 | flux_err, 216 | ra, 217 | dec, 218 | column, 219 | row, 220 | poscorr1, 221 | poscorr2, 222 | metadata, 223 | ) = _load_file(fname) 224 | # do cutout if asked 225 | if cutout_size is not None: 226 | flux, flux_err, ra, dec, column, row = _do_image_cutout( 227 | flux, 228 | flux_err, 229 | ra, 230 | dec, 231 | column, 232 | row, 233 | cutout_size=cutout_size, 234 | cutout_origin=cutout_origin, 235 | ) 236 | 237 | # we pass only non-empy pixels to the Gaia query and cleaning routines 238 | valid_pix = np.isfinite(flux).sum(axis=0).astype(bool) 239 | if sources is None or not isinstance(sources, pd.DataFrame): 240 | sources = _get_sources( 241 | ra[valid_pix], 242 | dec[valid_pix], 243 | wcs, 244 | magnitude_limit=magnitude_limit, 245 | epoch=time.jyear.mean(), 246 | ngrid=(2, 2) if flux.shape[1] <= 500 else (5, 5), 247 | dr=dr, 248 | img_limits=[[row.min(), row.max()], [column.min(), column.max()]], 249 | square=False, 250 | ) 251 | 252 | return SSMachine( 253 | time.jd, 254 | flux, 255 | flux_err, 256 | ra.ravel(), 257 | dec.ravel(), 258 | sources, 259 | column.ravel(), 260 | row.ravel(), 261 | wcs=wcs, 262 | meta=metadata, 263 | pos_corr1=poscorr1, 264 | pos_corr2=poscorr2, 265 | **kwargs, 266 | ) 267 | 268 | def fit_lightcurves( 269 | self, 270 | plot=False, 271 | iter_negative=False, 272 | fit_mean_shape_model=False, 273 | fit_va=False, 274 | sap=False, 275 | ): 276 | """ 277 | Fit the sources in the data to get its light curves. 278 | By default it only uses the per cadence PSF model to do the photometry. 279 | Alternatively it can fit the mean-PSF and the mean-PSF with time model to 280 | the data, this is the original method implemented in `PSFmachine` and described 281 | in the paper. Aperture Photometry is also available by creating aperture masks 282 | that follow the mean-PSF shape. 283 | 284 | This function creates the `lcs` attribuite that contains a collection of light 285 | curves in the form of `lightkurve.LightCurveCollection`. Each entry in the 286 | collection is a `lightkurve.KeplerLightCurve` object with the different type 287 | of photometry (PSF per cadence, SAP, mean-PSF, and mean-PSF velocity-aberration 288 | corrected). Also each `lightkurve.KeplerLightCurve` object includes its 289 | asociated metadata. 290 | 291 | The photometry can also be accessed independently from the following attribuites 292 | that `fit_lightcurves` create: 293 | * `ws` and `werrs` have the uncorrected PSF flux and flux errors. 294 | * `ws_va` and `werrs_va` have the PSF flux and flux errors corrected by 295 | velocity aberration. 296 | * `sap_flux` and `sap_flux_err` have the flux and flux errors computed 297 | using aperture mask. 298 | * `ws_frame` and `werrs_frame` have the flux from PSF at each cadence. 299 | 300 | Parameters 301 | ---------- 302 | plot : bool 303 | Whether or not to show some diagnostic plots. These can be helpful 304 | for a user to see if the PRF and time dependent models are being calculated 305 | correctly. 306 | iter_negative : bool 307 | When fitting light curves, it isn't possible to force the flux to be 308 | positive. As such, when we find there are light curves that deviate into 309 | negative flux values, we can clip these targets out of the analysis and 310 | rerun the model. 311 | If iter_negative is True, PSFmachine will run up to 3 times, clipping out 312 | any negative targets each round. This is used when 313 | `fit_mean_shape_model` is `True`. 314 | fit_mean_shape_model : bool 315 | Will do PSF photmetry using the mean-PSF. 316 | fit_va : bool 317 | Whether or not to fit Velocity Aberration (which implicitly will try to fit 318 | other kinds of time variability). `fit_mean_shape_model` must set to `True` 319 | ortherwise will be ignored. This will try to fit the "long term" 320 | trends in the dataset. If True, this will take slightly longer to fit. 321 | If you are interested in short term phenomena, like transits, you may 322 | find you do not need this to be set to True. If you have the time, it 323 | is recommended to run it. 324 | sap : boolean 325 | Compute or not Simple Aperture Photometry. See 326 | `Machine.compute_aperture_photometry()` for details. 327 | """ 328 | # create mean shape model to be used by SAP and mean-PSF 329 | self.build_shape_model(plot=plot, frame_index="mean") 330 | # do SAP first 331 | if sap: 332 | self.compute_aperture_photometry( 333 | aperture_size="optimal", target_complete=1, target_crowd=1 334 | ) 335 | 336 | # do mean-PSF photometry and time model if asked 337 | if fit_mean_shape_model: 338 | self.build_time_model(plot=plot, downsample=True) 339 | # fit the OG time model 340 | self.fit_model(fit_va=fit_va) 341 | if iter_negative: 342 | # More than 2% negative cadences 343 | negative_sources = (self.ws_va < 0).sum(axis=0) > (0.02 * self.nt) 344 | idx = 1 345 | while len(negative_sources) > 0: 346 | self.mean_model[negative_sources] *= 0 347 | self.fit_model(fit_va=fit_va) 348 | negative_sources = np.where((self.ws_va < 0).all(axis=0))[0] 349 | idx += 1 350 | if idx >= 3: 351 | break 352 | 353 | # fit shape model at each cadence 354 | self.build_frame_shape_model() 355 | self.fit_frame_model() 356 | 357 | self.lcs = [] 358 | for idx, s in self.sources.iterrows(): 359 | meta = { 360 | "ORIGIN": "PSFMACHINE", 361 | "APERTURE": "PSF + SAP" if sap else "PSF", 362 | "LABEL": s.designation, 363 | "MISSION": self.meta["MISSION"], 364 | "RA": s.ra, 365 | "DEC": s.dec, 366 | "PMRA": s.pmra / 1000, 367 | "PMDEC": s.pmdec / 1000, 368 | "PARALLAX": s.parallax, 369 | "GMAG": s.phot_g_mean_mag, 370 | "RPMAG": s.phot_rp_mean_mag, 371 | "BPMAG": s.phot_bp_mean_mag, 372 | } 373 | 374 | attrs = [ 375 | "channel", 376 | "module", 377 | "ccd", 378 | "camera", 379 | "quarter", 380 | "campaign", 381 | "quarter", 382 | "row", 383 | "column", 384 | "mission", 385 | ] 386 | for attr in attrs: 387 | if attr in self.meta.keys(): 388 | meta[attr.upper()] = self.meta[attr] 389 | 390 | lc = lk.KeplerLightCurve( 391 | time=(self.time) * u.d, 392 | flux=self.ws_frame[:, idx] * (u.electron / u.second), 393 | flux_err=self.werrs_frame[:, idx] * (u.electron / u.second), 394 | meta=meta, 395 | time_format="jd", 396 | ) 397 | if fit_mean_shape_model: 398 | lc["flux_NVA"] = (self.ws[:, idx]) * u.electron / u.second 399 | lc["flux_err_NVA"] = (self.werrs[:, idx]) * u.electron / u.second 400 | if fit_va: 401 | lc["flux_VA"] = (self.ws_va[:, idx]) * u.electron / u.second 402 | lc["flux_err_VA"] = (self.werrs_va[:, idx]) * u.electron / u.second 403 | if sap: 404 | lc["sap_flux"] = (self.sap_flux[:, idx]) * u.electron / u.second 405 | lc["sap_flux_err"] = (self.sap_flux_err[:, idx]) * u.electron / u.second 406 | 407 | self.lcs.append(lc) 408 | self.lcs = lk.LightCurveCollection(self.lcs) 409 | return 410 | 411 | def plot_image_interactive(self, ax=None, sources=False): 412 | """ 413 | Function to plot the super stamp and Gaia Sources and interact by changing the 414 | cadence. 415 | 416 | Parameters 417 | ---------- 418 | ax : matplotlib.axes 419 | Matlotlib axis can be provided, if not one will be created and returned 420 | sources : boolean 421 | Whether to overplot or not the source catalog 422 | Returns 423 | ------- 424 | ax : matplotlib.axes 425 | Matlotlib axis with the figure 426 | """ 427 | if ax is None: 428 | fig, ax = plt.subplots(1, figsize=(10, 10)) 429 | 430 | ax = plt.subplot(projection=self.wcs) 431 | row_2d, col_2d = np.mgrid[ 432 | self.row.min() : self.row.max() + 1, 433 | self.column.min() : self.column.max() + 1, 434 | ] 435 | im = ax.pcolormesh( 436 | col_2d, 437 | row_2d, 438 | self.flux_2d[0], 439 | cmap=plt.cm.viridis, 440 | shading="nearest", 441 | norm=colors.SymLogNorm(linthresh=200, vmin=0, vmax=2000, base=10), 442 | rasterized=True, 443 | ) 444 | plt.colorbar(im, ax=ax, label=r"Flux ($e^{-}s^{-1}$)", fraction=0.042) 445 | 446 | ax.set_xlabel("R.A. [hh:mm]") 447 | ax.set_ylabel("Decl. [deg]") 448 | ax.set_xlim(self.column.min() - 2, self.column.max() + 2) 449 | ax.set_ylim(self.row.min() - 2, self.row.max() + 2) 450 | 451 | ax.set_aspect("equal", adjustable="box") 452 | 453 | def update(t): 454 | ax.pcolormesh( 455 | col_2d, 456 | row_2d, 457 | self.flux_2d[t], 458 | cmap=plt.cm.viridis, 459 | shading="nearest", 460 | norm=colors.SymLogNorm(linthresh=200, vmin=0, vmax=2000, base=10), 461 | rasterized=True, 462 | ) 463 | ax.set_title( 464 | "%s %s Ch/CCD %s MJD %f" 465 | % ( 466 | self.meta["MISSION"], 467 | self.meta["OBJECT"], 468 | self.meta["EXTENSION"], 469 | self.time[t], 470 | ) 471 | ) 472 | if sources: 473 | ax.scatter( 474 | self.sources.column, 475 | self.sources.row, 476 | facecolors="none", 477 | edgecolors="r", 478 | linewidths=0.5 if self.sources.shape[0] > 1000 else 1, 479 | alpha=0.9, 480 | ) 481 | ax.grid(True, which="major", axis="both", ls="-", color="w", alpha=0.7) 482 | fig.canvas.draw_idle() 483 | 484 | interact(update, t=(0, self.flux_2d.shape[0] - 1, 1)) 485 | 486 | return 487 | 488 | 489 | def _load_file(fname): 490 | """ 491 | Helper function to load K2 SuperStamp files files and parse data. This function 492 | works with K2 SS files created by Cody et al. 2018, which are single cadence FITS 493 | file, then a full campaign has many FITS files (e.g. M67 has 3620 files). 494 | 495 | Parameters 496 | ---------- 497 | fname : string or list of strings 498 | Path to the FITS files to be parsed. For only one frame, pass a string, 499 | for multiple frames pass a list of paths. 500 | Returns 501 | ------- 502 | wcs : astropy.wcs 503 | World coordinates system solution for the FFI. Used to convert RA, Dec to pixels 504 | time : numpy.array 505 | Array with time values in MJD 506 | flux_2d : numpy.ndarray 507 | Array with 2D (image) representation of flux values 508 | flux_err_2d : numpy.ndarray 509 | Array with 2D (image) representation of flux errors 510 | ra_2d : numpy.ndarray 511 | Array with 2D (image) representation of flux RA 512 | dec_2d : numpy.ndarray 513 | Array with 2D (image) representation of flux Dec 514 | col_2d : numpy.ndarray 515 | Array with 2D (image) representation of pixel column 516 | row_2d : numpy.ndarray 517 | Array with 2D (image) representation of pixel row 518 | meta : dict 519 | Dictionary with metadata 520 | """ 521 | if not isinstance(fname, list): 522 | fname = np.sort([fname]) 523 | flux, flux_err = [], [] 524 | times = [] 525 | telescopes = [] 526 | campaigns = [] 527 | channels = [] 528 | quality = [] 529 | wcs_b = True 530 | poscorr1, poscorr2 = [], [] 531 | for i, f in enumerate(fname): 532 | if not os.path.isfile(f): 533 | raise FileNotFoundError("FFI calibrated fits file does not exist: ", f) 534 | 535 | with fits.open(f) as hdul: 536 | # hdul = fits.open(f) 537 | header = hdul[0].header 538 | telescopes.append(header["TELESCOP"]) 539 | campaigns.append(header["CAMPAIGN"]) 540 | quality.append(header["QUALITY"]) 541 | 542 | # clusters only have one ext, bulge have multi-extensions 543 | if len(hdul) > 1: 544 | img_ext = 1 545 | else: 546 | img_ext = 0 547 | channels.append(hdul[img_ext].header["CHANNEL"]) 548 | hdr = hdul[img_ext].header 549 | # times.append(Time([hdr["DATE-OBS"], hdr["DATE-END"]], format="isot").mjd.mean()) 550 | poscorr1.append(float(hdr["POSCORR1"])) 551 | poscorr2.append(float(hdr["POSCORR2"])) 552 | times.append(Time([hdr["TSTART"], hdr["TSTOP"]], format="jd").mjd.mean()) 553 | flux.append(hdul[img_ext].data) 554 | if img_ext == 1: 555 | flux_err.append(hdul[2].data) 556 | else: 557 | flux_err.append(np.sqrt(np.abs(hdul[img_ext].data))) 558 | 559 | if header["QUALITY"] == 0 and wcs_b: 560 | wcs_b = False 561 | wcs = WCS(hdr) 562 | 563 | # check for integrity of files, same telescope, all FFIs and same quarter/campaign 564 | if len(set(telescopes)) != 1: 565 | raise ValueError("All FFIs must be from same telescope") 566 | 567 | # collect meta data, I get everthing from one header. 568 | attrs = [ 569 | "TELESCOP", 570 | "INSTRUME", 571 | "MISSION", 572 | "DATSETNM", 573 | "OBSMODE", 574 | "OBJECT", 575 | ] 576 | meta = {k: header[k] for k in attrs if k in header.keys()} 577 | attrs = [ 578 | "OBJECT", 579 | "RADESYS", 580 | "EQUINOX", 581 | "BACKAPP", 582 | ] 583 | meta.update({k: hdr[k] for k in attrs if k in hdr.keys()}) 584 | meta.update({"EXTENSION": channels[0], "QUARTER": campaigns[0], "DCT_TYPE": "FFI"}) 585 | if "MISSION" not in meta.keys(): 586 | meta["MISSION"] = meta["TELESCOP"] 587 | 588 | # mask by quality and sort by times 589 | qual_mask = lk.utils.KeplerQualityFlags.create_quality_mask( 590 | np.array(quality), 1 | 2 | 4 | 8 | 32 | 16384 | 32768 | 65536 | 1048576 591 | ) 592 | times = Time(times, format="mjd", scale="tdb") 593 | tdx = np.argsort(times)[qual_mask] 594 | times = times[tdx] 595 | row_2d, col_2d = np.mgrid[: flux[0].shape[0], : flux[0].shape[1]] 596 | flux_2d = np.array(flux)[tdx] 597 | flux_err_2d = np.array(flux_err)[tdx] 598 | poscorr1 = np.array(poscorr1)[tdx] 599 | poscorr2 = np.array(poscorr2)[tdx] 600 | 601 | ra, dec = wcs.all_pix2world(np.vstack([col_2d.ravel(), row_2d.ravel()]).T, 0.0).T 602 | ra_2d = ra.reshape(flux_2d.shape[1:]) 603 | dec_2d = dec.reshape(flux_2d.shape[1:]) 604 | 605 | del hdul, header, hdr, ra, dec 606 | 607 | return ( 608 | wcs, 609 | times, 610 | flux_2d, 611 | flux_err_2d, 612 | ra_2d, 613 | dec_2d, 614 | col_2d, 615 | row_2d, 616 | poscorr1, 617 | poscorr2, 618 | meta, 619 | ) 620 | -------------------------------------------------------------------------------- /src/psfmachine/utils.py: -------------------------------------------------------------------------------- 1 | """ Collection of utility functions""" 2 | 3 | import numpy as np 4 | import pandas as pd 5 | import diskcache 6 | 7 | from scipy import sparse 8 | from patsy import dmatrix 9 | from scipy.ndimage import gaussian_filter1d 10 | import pyia 11 | import fitsio 12 | 13 | # size_limit is 1GB 14 | cache = diskcache.Cache(directory="~/.psfmachine-cache") 15 | 16 | 17 | # cache during 30 days 18 | @cache.memoize(expire=2.592e06) 19 | def get_gaia_sources(ras, decs, rads, magnitude_limit=18, epoch=2020, dr=3): 20 | """ 21 | Will find gaia sources using a TAP query, accounting for proper motions. 22 | 23 | Inputs have be hashable, e.g. tuples 24 | 25 | Parameters 26 | ---------- 27 | ras : tuple 28 | Tuple with right ascension coordinates to be queried 29 | shape nsources 30 | decs : tuple 31 | Tuple with declination coordinates to be queried 32 | shape nsources 33 | rads : tuple 34 | Tuple with radius query 35 | shape nsources 36 | magnitude_limit : int 37 | Limiting magnitued for query 38 | epoch : float 39 | Epoch to be used for propper motion correction during Gaia crossmatch. 40 | dr : int or string 41 | Gaia Data Release version, if Early DR 3 (aka EDR3) is wanted use `"edr3"`. 42 | 43 | Returns 44 | ------- 45 | Pandas DatFrame with number of result sources (rows) and Gaia columns 46 | 47 | """ 48 | if not hasattr(ras, "__iter__"): 49 | ras = [ras] 50 | if not hasattr(decs, "__iter__"): 51 | decs = [decs] 52 | if not hasattr(rads, "__iter__"): 53 | rads = [rads] 54 | if dr not in [1, 2, 3, "edr3"]: 55 | raise ValueError("Please pass a valid data release") 56 | if isinstance(dr, int): 57 | dr = f"dr{dr}" 58 | wheres = [ 59 | f"""1=CONTAINS( 60 | POINT('ICRS',{ra},{dec}), 61 | CIRCLE('ICRS',ra,dec,{rad}))""" 62 | for ra, dec, rad in zip(ras, decs, rads) 63 | ] 64 | 65 | where = """\n\tOR """.join(wheres) 66 | 67 | gd = pyia.GaiaData.from_query( 68 | f"""SELECT designation, 69 | coord1(prop) AS ra, coord2(prop) AS dec, parallax, 70 | parallax_error, pmra, pmdec, 71 | phot_g_mean_flux, 72 | phot_g_mean_mag, 73 | phot_bp_mean_mag, 74 | phot_rp_mean_mag, 75 | phot_bp_mean_flux, 76 | phot_rp_mean_flux FROM ( 77 | SELECT *, 78 | EPOCH_PROP_POS(ra, dec, parallax, pmra, pmdec, 0, ref_epoch, {epoch}) AS prop 79 | FROM gaia{dr}.gaia_source 80 | WHERE {where} 81 | ) AS subquery 82 | WHERE phot_g_mean_mag<={magnitude_limit} 83 | """ 84 | ) 85 | 86 | return gd.data.to_pandas() 87 | 88 | 89 | def do_tiled_query(ra, dec, ngrid=(5, 5), magnitude_limit=18, epoch=2020, dr=3): 90 | """ 91 | Find the centers and radius of tiled queries when the sky area is large. 92 | This function divides the data into `ngrid` tiles and compute the ra, dec 93 | coordinates for each tile as well as its radius. 94 | This is meant to be used with dense data, e.g. FFI or cluster fields, and it is not 95 | optimized for sparse data, e.g. TPF stacks. For the latter use 96 | `psfmachine.tpf._get_coord_and_query_gaia()`. 97 | 98 | Parameters 99 | ---------- 100 | ra : numpy.ndarray 101 | Data array with values of Right Ascension. Array can be 2D image or flatten. 102 | dec : numpy.ndarray 103 | Data array with values of Declination. Array can be 2D image or flatten. 104 | ngrid : tuple 105 | Tuple with number of bins in each axis. Default is (5, 5). 106 | magnitude_limit : int 107 | Limiting magnitude for query 108 | epoch : float 109 | Year of the observation (Julian year) used for proper motion correction. 110 | dr : int 111 | Gaia Data Release to be used, DR2 or EDR3. Default is EDR3. 112 | 113 | Returns 114 | ------- 115 | sources : pandas.DatFrame 116 | Pandas DatFrame with number of result sources (rows) and Gaia columns 117 | """ 118 | # find edges of the bins 119 | ra_edges = np.histogram_bin_edges(ra, ngrid[0]) 120 | dec_edges = np.histogram_bin_edges(dec, ngrid[1]) 121 | sources = [] 122 | # iterate over 2d bins 123 | for idx in range(1, len(ra_edges)): 124 | for jdx in range(1, len(dec_edges)): 125 | # check if image data fall in the bin 126 | _in = ( 127 | (ra_edges[idx - 1] <= ra) 128 | & (ra <= ra_edges[idx]) 129 | & (dec_edges[jdx - 1] <= dec) 130 | & (dec <= dec_edges[jdx]) 131 | ) 132 | if not _in.any(): 133 | continue 134 | # get the center coord of the query and radius to 7th decimal precision 135 | # (3 miliarcsec) to avoid not catching get_gaia_sources() due to 136 | # floating point error. 137 | ra_in = ra[_in] 138 | dec_in = dec[_in] 139 | # we use 50th percentile to get the centers and avoid 360-0 boundary 140 | ra_q = np.round(np.percentile(ra_in, 50), decimals=7) 141 | dec_q = np.round(np.percentile(dec_in, 50), decimals=7) 142 | # HARDCODED: +10/3600 to add a 10 arcsec to search radius, this is to get 143 | # sources off sensor up to 10" distance from sensor border. 144 | rad_q = np.round( 145 | np.hypot(ra_in - ra_q, dec_in - dec_q).max() + 10 / 3600, decimals=7 146 | ) 147 | # query gaia with ra, dec, rad, epoch 148 | result = get_gaia_sources( 149 | tuple([ra_q]), 150 | tuple([dec_q]), 151 | tuple([rad_q]), 152 | magnitude_limit=magnitude_limit, 153 | epoch=epoch, 154 | dr=dr, 155 | ) 156 | sources.append(result) 157 | # concat results and remove duplicated sources 158 | sources = pd.concat(sources, axis=0).drop_duplicates(subset=["designation"]) 159 | return sources.reset_index(drop=True) 160 | 161 | 162 | def _make_A_polar(phi, r, cut_r=6, rmin=1, rmax=18, n_r_knots=12, n_phi_knots=15): 163 | """ 164 | Creates a design matrix (DM) in polar coordinates (r, phi). It will enforce r-only 165 | dependency within `cut_r` radius. This is useful when less data points are available 166 | near the center. 167 | 168 | Parameters 169 | ---------- 170 | phi : np.ndarray 171 | Array of angle (phi) values in polar coordinates. Must have values in the 172 | [-pi, pi] range. 173 | r : np.ndarray 174 | Array of radii values in polar coordinates. 175 | cut_r : float 176 | Radius (units consistent with `r`) whitin the DM only has radius dependency 177 | and not angle. 178 | rmin : float 179 | Radius where the DM starts. 180 | rmax : float 181 | Radius where the DM ends. 182 | n_r_knots : int 183 | Number of knots used for the spline in radius. 184 | n_phi_knots : int 185 | Number of knots used for the spline in angle. 186 | Returns 187 | ------- 188 | X : sparse CSR matrix 189 | A DM with bspline basis in polar coordinates. 190 | """ 191 | # create the spline bases for radius and angle 192 | phi_spline = sparse.csr_matrix(wrapped_spline(phi, order=3, nknots=n_phi_knots).T) 193 | r_knots = np.linspace(rmin ** 0.5, rmax ** 0.5, n_r_knots) ** 2 194 | cut_r_int = np.where(r_knots <= cut_r)[0].max() 195 | r_spline = sparse.csr_matrix( 196 | np.asarray( 197 | dmatrix( 198 | "bs(x, knots=knots, degree=3, include_intercept=True)", 199 | {"x": list(np.hstack([r, rmin, rmax])), "knots": r_knots}, 200 | ) 201 | ) 202 | )[:-2] 203 | 204 | # build full desing matrix 205 | X = sparse.hstack( 206 | [phi_spline.multiply(r_spline[:, idx]) for idx in range(r_spline.shape[1])], 207 | format="csr", 208 | ) 209 | # find and remove the angle dependency for all basis for radius < 6 210 | cut = np.arange(0, phi_spline.shape[1] * cut_r_int) 211 | a = list(set(np.arange(X.shape[1])) - set(cut)) 212 | X1 = sparse.hstack( 213 | [X[:, a], r_spline[:, 1:cut_r_int], sparse.csr_matrix(np.ones(X.shape[0])).T], 214 | format="csr", 215 | ) 216 | return X1 217 | 218 | 219 | def spline1d(x, knots, degree=3, include_knots=False): 220 | """ 221 | Make a bspline design matrix (DM) for 1D variable `x`. 222 | 223 | Parameters 224 | ---------- 225 | x : np.ndarray 226 | Array of values to create the DM. 227 | knots : np.ndarray 228 | Array of knots to be used in the DM. 229 | degree : int 230 | Degree of the spline, default is 3. 231 | include_knots : boolean 232 | Include or not the knots in the `x` vector, this forces knots in case 233 | out of bound values. 234 | 235 | Returns 236 | ------- 237 | X : sparse CSR matrix 238 | A DM with bspline basis for vector `x`. 239 | """ 240 | if include_knots: 241 | x = np.hstack([knots.min(), x, knots.max()]) 242 | X = sparse.csr_matrix( 243 | np.asarray( 244 | dmatrix( 245 | "bs(x, knots=knots, degree=degree, include_intercept=True)", 246 | {"x": list(x), "knots": knots, "degree": degree}, 247 | ) 248 | ) 249 | ) 250 | if include_knots: 251 | X = X[1:-1] 252 | x = x[1:-1] 253 | if not X.shape[0] == x.shape[0]: 254 | raise ValueError("`patsy` has made the wrong matrix.") 255 | X = X[:, np.asarray(X.sum(axis=0) != 0)[0]] 256 | return X 257 | 258 | 259 | def _make_A_cartesian(x, y, n_knots=10, radius=3.0, knot_spacing_type="sqrt", degree=3): 260 | """ 261 | Creates a design matrix (DM) in Cartersian coordinates (r, phi). 262 | 263 | Parameters 264 | ---------- 265 | x : np.ndarray 266 | Array of x values in Cartersian coordinates. 267 | y : np.ndarray 268 | Array of y values in Cartersian coordinates. 269 | n_knots : int 270 | Number of knots used for the spline. 271 | radius : float 272 | Distance from 0 to the furthes knot. 273 | knot_spacing_type : string 274 | Type of spacing betwen knots, options are "linear" or "sqrt". 275 | degree : int 276 | Degree of the spline, default is 3. 277 | 278 | Returns 279 | ------- 280 | X : sparse CSR matrix 281 | A DM with bspline basis in Cartersian coordinates. 282 | """ 283 | # Must be odd 284 | n_odd_knots = n_knots if n_knots % 2 == 1 else n_knots + 1 285 | if knot_spacing_type == "sqrt": 286 | x_knots = np.linspace(-np.sqrt(radius), np.sqrt(radius), n_odd_knots) 287 | x_knots = np.sign(x_knots) * x_knots ** 2 288 | y_knots = np.linspace(-np.sqrt(radius), np.sqrt(radius), n_odd_knots) 289 | y_knots = np.sign(y_knots) * y_knots ** 2 290 | else: 291 | x_knots = np.linspace(-radius, radius, n_odd_knots) 292 | y_knots = np.linspace(-radius, radius, n_odd_knots) 293 | 294 | x_spline = spline1d(x, knots=x_knots, degree=degree, include_knots=True) 295 | y_spline = spline1d(y, knots=y_knots, degree=degree, include_knots=True) 296 | 297 | x_spline = x_spline[:, np.asarray(x_spline.sum(axis=0))[0] != 0] 298 | y_spline = y_spline[:, np.asarray(y_spline.sum(axis=0))[0] != 0] 299 | X = sparse.hstack( 300 | [x_spline.multiply(y_spline[:, idx]) for idx in range(y_spline.shape[1])], 301 | format="csr", 302 | ) 303 | return X 304 | 305 | 306 | def wrapped_spline(input_vector, order=2, nknots=10): 307 | """ 308 | Creates a vector of folded-spline basis according to the input data. This is meant 309 | to be used to build the basis vectors for periodic data, like the angle in polar 310 | coordinates. 311 | 312 | Parameters 313 | ---------- 314 | input_vector : numpy.ndarray 315 | Input data to create basis, angle values MUST BE BETWEEN -PI and PI. 316 | order : int 317 | Order of the spline basis 318 | nknots : int 319 | Number of knots for the splines 320 | 321 | Returns 322 | ------- 323 | folded_basis : numpy.ndarray 324 | Array of folded-spline basis 325 | """ 326 | 327 | if not ((input_vector >= -np.pi) & (input_vector <= np.pi)).all(): 328 | raise ValueError("Must be between -pi and pi") 329 | x = np.copy(input_vector) 330 | x1 = np.hstack([x, x + np.pi * 2]) 331 | nt = (nknots * 2) + 1 332 | 333 | t = np.linspace(-np.pi, 3 * np.pi, nt) 334 | dt = np.median(np.diff(t)) 335 | # Zeroth order basis 336 | basis = np.asarray( 337 | [ 338 | ((x1 >= t[idx]) & (x1 < t[idx + 1])).astype(float) 339 | for idx in range(len(t) - 1) 340 | ] 341 | ) 342 | # Higher order basis 343 | for order in np.arange(1, 4): 344 | basis_1 = [] 345 | for idx in range(len(t) - 1): 346 | a = ((x1 - t[idx]) / (dt * order)) * basis[idx] 347 | 348 | if ((idx + order + 1)) < (nt - 1): 349 | b = (-(x1 - t[(idx + order + 1)]) / (dt * order)) * basis[ 350 | (idx + 1) % (nt - 1) 351 | ] 352 | else: 353 | b = np.zeros(len(x1)) 354 | basis_1.append(a + b) 355 | basis = np.vstack(basis_1) 356 | 357 | folded_basis = np.copy(basis)[: nt // 2, : len(x)] 358 | for idx in np.arange(-order, 0): 359 | folded_basis[idx, :] += np.copy(basis)[nt // 2 + idx, len(x) :] 360 | return folded_basis 361 | 362 | 363 | def solve_linear_model( 364 | A, y, y_err=None, prior_mu=None, prior_sigma=None, k=None, errors=False 365 | ): 366 | """ 367 | Solves a linear model with design matrix A and observations y: 368 | Aw = y 369 | return the solutions w for the system assuming Gaussian priors. 370 | Alternatively the observation errors, priors, and a boolean mask for the 371 | observations (row axis) can be provided. 372 | 373 | Adapted from Luger, Foreman-Mackey & Hogg, 2017 374 | (https://ui.adsabs.harvard.edu/abs/2017RNAAS...1....7L/abstract) 375 | 376 | Parameters 377 | ---------- 378 | A: numpy ndarray or scipy sparce csr matrix 379 | Desging matrix with solution basis 380 | shape n_observations x n_basis 381 | y: numpy ndarray 382 | Observations 383 | shape n_observations 384 | y_err: numpy ndarray, optional 385 | Observation errors 386 | shape n_observations 387 | prior_mu: float, optional 388 | Mean of Gaussian prior values for the weights (w) 389 | prior_sigma: float, optional 390 | Standard deviation of Gaussian prior values for the weights (w) 391 | k: boolean, numpy ndarray, optional 392 | Mask that sets the observations to be used to solve the system 393 | shape n_observations 394 | errors: boolean 395 | Whether to return error estimates of the best fitting weights 396 | 397 | Returns 398 | ------- 399 | w: numpy ndarray 400 | Array with the estimations for the weights 401 | shape n_basis 402 | werrs: numpy ndarray 403 | Array with the error estimations for the weights, returned if `error` is True 404 | shape n_basis 405 | """ 406 | if k is None: 407 | k = np.ones(len(y), dtype=bool) 408 | 409 | if y_err is not None: 410 | sigma_w_inv = A[k].T.dot(A[k].multiply(1 / y_err[k, None] ** 2)) 411 | B = A[k].T.dot((y[k] / y_err[k] ** 2)) 412 | else: 413 | sigma_w_inv = A[k].T.dot(A[k]) 414 | B = A[k].T.dot(y[k]) 415 | 416 | if prior_mu is not None and prior_sigma is not None: 417 | sigma_w_inv += np.diag(1 / prior_sigma ** 2) 418 | B += prior_mu / prior_sigma ** 2 419 | 420 | if isinstance(sigma_w_inv, (sparse.csr_matrix, sparse.csc_matrix)): 421 | sigma_w_inv = sigma_w_inv.toarray() 422 | if isinstance(sigma_w_inv, np.matrix): 423 | sigma_w_inv = np.asarray(sigma_w_inv) 424 | 425 | w = np.linalg.solve(sigma_w_inv, B) 426 | if errors is True: 427 | w_err = np.linalg.inv(sigma_w_inv).diagonal() ** 0.5 428 | return w, w_err 429 | return w 430 | 431 | 432 | def sparse_lessthan(arr, limit): 433 | """ 434 | Compute less than operation on sparse array by evaluating only non-zero values 435 | and reconstructing the sparse array. This function return a sparse array, which is 436 | crutial to keep operating large matrices. 437 | 438 | Notes: when doing `x < a` for a sparse array `x` and `a > 0` it effectively compares 439 | all zero and non-zero values. Then we get a dense boolean array with `True` where 440 | the condition is met but also `True` where the sparse array was zero. 441 | To avoid this we evaluate the condition only for non-zero values in the sparse 442 | array and later reconstruct the sparse array with the right shape and content. 443 | When `x` is a [N * M] matrix and `a` is [N] array, and we want to evaluate the 444 | condition per row, we need to iterate over rows to perform the evaluation and then 445 | reconstruct the masked sparse array. 446 | 447 | Parameters 448 | ---------- 449 | arr : scipy.sparse 450 | Sparse array to be masked, is a 2D matrix. 451 | limit : float, numpy.array 452 | Upper limit to evaluate less than. If float will do `arr < limit`. If array, 453 | shape has to match first dimension of `arr` to do `arr < limi[:, None]`` and 454 | evaluate the condition per row. 455 | 456 | Returns 457 | ------- 458 | masked_arr : scipy.sparse.csr_matrix 459 | Sparse array after less than evaluation. 460 | """ 461 | nonz_idx = arr.nonzero() 462 | # apply condition for each row 463 | if isinstance(limit, np.ndarray) and limit.shape[0] == arr.shape[0]: 464 | mask = [arr[s].data < limit[s] for s in set(nonz_idx[0])] 465 | # flatten mask 466 | mask = [x for sub in mask for x in sub] 467 | else: 468 | mask = arr.data < limit 469 | # reconstruct sparse array 470 | masked_arr = sparse.csr_matrix( 471 | (arr.data[mask], (nonz_idx[0][mask], nonz_idx[1][mask])), 472 | shape=arr.shape, 473 | ).astype(bool) 474 | return masked_arr 475 | 476 | 477 | def _combine_A(A, poscorr=None, time=None): 478 | """ 479 | Combines a design matrix A (cartesian) with a time corrector type. 480 | If poscorr is provided, A will be combined with both axis of the pos corr as a 481 | 1st degree polynomial. 482 | If time is provided, A will be combined with the time values as a 3rd degree 483 | polynomialin time. 484 | 485 | Parameters 486 | ---------- 487 | A : sparse.csr_matrix 488 | A sparse design matix in of cartesian coordinates created with _make_A_cartesian 489 | poscorr : list 490 | A list of pos_corr arrays for axis 1 and 2 491 | time : numpy.array 492 | An array with time values 493 | """ 494 | if poscorr: 495 | # Cartesian spline with poscor dependence 496 | A2 = sparse.hstack( 497 | [ 498 | A, 499 | A.multiply(poscorr[0].ravel()[:, None]), 500 | A.multiply(poscorr[1].ravel()[:, None]), 501 | A.multiply((poscorr[0] * poscorr[1]).ravel()[:, None]), 502 | ], 503 | format="csr", 504 | ) 505 | return A2 506 | elif time is not None: 507 | # Cartesian spline with time dependence 508 | A2 = sparse.hstack( 509 | [ 510 | A, 511 | A.multiply(time.ravel()[:, None]), 512 | A.multiply(time.ravel()[:, None] ** 2), 513 | A.multiply(time.ravel()[:, None] ** 3), 514 | ], 515 | format="csr", 516 | ) 517 | return A2 518 | 519 | 520 | def threshold_bin(x, y, z, z_err=None, abs_thresh=10, bins=15, statistic=np.nanmedian): 521 | """ 522 | Function to bin 2D data and compute array statistic based on density. 523 | This function inputs 2D coordinates, e.g. `X` and `Y` locations, and a number value 524 | `Z` for each point in the 2D space. It bins the 2D spatial data to then compute 525 | a `statistic`, e.g. median, on the Z value based on bin members. The `statistic` 526 | is computed only for bins with more than `abs_thresh` members. It preserves data 527 | when the number of bin memebers is lower than `abs_thresh`. 528 | 529 | Parameters 530 | ---------- 531 | x : numpy.ndarray 532 | Data array with spatial coordinate 1. 533 | y : numpy.ndarray 534 | Data array with spatial coordinate 2. 535 | z : numpy.ndarray 536 | Data array with the number values for each (X, Y) point. 537 | z_err : numpy.ndarray 538 | Array with errors values for z. 539 | abs_thresh : int 540 | Absolute threshold is the number of bib members to compute the statistic, 541 | otherwise data will be preserved. 542 | bins : int or list of ints 543 | Number of bins. If int, both axis will have same number of bins. If list, number 544 | of bins for first (x) and second (y) dimension. 545 | statistic : callable() 546 | The statistic as a callable function that will be use in each bin. 547 | Default is `numpy.nanmedian`. 548 | 549 | Returns 550 | ------- 551 | bin_map : numpy.ndarray 552 | 2D histogram values 553 | new_x : numpy.ndarray 554 | Binned X data. 555 | new_y : numpy.ndarray 556 | Binned Y data. 557 | new_z : numpy.ndarray 558 | Binned Z data. 559 | new_z_err : numpy.ndarray 560 | BInned Z_err data if errors were provided. If no, inverse of the number of 561 | bin members are returned as weights. 562 | """ 563 | if bins < 2 or bins > len(x): 564 | raise ValueError( 565 | "Number of bins is negative or higher than number of points in (x, y, z)" 566 | ) 567 | if abs_thresh < 1: 568 | raise ValueError( 569 | "Absolute threshold is 0 or negative, please input a value > 0" 570 | ) 571 | if isinstance(bins, int): 572 | bins = [bins, bins] 573 | 574 | xedges = np.linspace(np.nanmin(x), np.nanmax(x), num=bins[0] + 1) 575 | yedges = np.linspace(np.nanmin(y), np.nanmax(y), num=bins[1] + 1) 576 | bin_mask = np.zeros_like(z, dtype=bool) 577 | new_x, new_y, new_z, new_z_err, bin_map = [], [], [], [], [] 578 | 579 | for j in range(1, len(xedges)): 580 | for k in range(1, len(yedges)): 581 | idx = np.where( 582 | (x >= xedges[j - 1]) 583 | & (x < xedges[j]) 584 | & (y >= yedges[k - 1]) 585 | & (y < yedges[k]) 586 | )[0] 587 | if len(idx) >= abs_thresh: 588 | bin_mask[idx] = True 589 | # we agregate bin memebers 590 | new_x.append((xedges[j - 1] + xedges[j]) / 2) 591 | new_y.append((yedges[k - 1] + yedges[k]) / 2) 592 | new_z.append(statistic(z[idx])) 593 | bin_map.append(len(idx)) 594 | if isinstance(z_err, np.ndarray): 595 | # agregate errors if provided and sccale by bin member number 596 | new_z_err.append(np.sqrt(np.nansum(z_err[idx] ** 2)) / len(idx)) 597 | 598 | # adding non-binned datapoints 599 | new_x.append(x[~bin_mask]) 600 | new_y.append(y[~bin_mask]) 601 | new_z.append(z[~bin_mask]) 602 | bin_map.append(np.ones_like(z)[~bin_mask]) 603 | 604 | if isinstance(z_err, np.ndarray): 605 | # keep original z errors if provided 606 | new_z_err.append(z_err[~bin_mask]) 607 | else: 608 | new_z_err = 1 / np.hstack(bin_map) 609 | 610 | return ( 611 | np.hstack(bin_map), 612 | np.hstack(new_x), 613 | np.hstack(new_y), 614 | np.hstack(new_z), 615 | np.hstack(new_z_err), 616 | ) 617 | 618 | 619 | def get_breaks(time, include_ext=False): 620 | """ 621 | Finds discontinuity in the time array and return the break indexes. 622 | 623 | Parameters 624 | ---------- 625 | time : numpy.ndarray 626 | Array with time values 627 | 628 | Returns 629 | ------- 630 | splits : numpy.ndarray 631 | An array of indexes with the break positions 632 | """ 633 | dts = np.diff(time) 634 | if include_ext: 635 | return np.hstack([0, np.where(dts > 5 * np.median(dts))[0] + 1, len(time)]) 636 | else: 637 | return np.where(dts > 5 * np.median(dts))[0] + 1 638 | 639 | 640 | def gaussian_smooth( 641 | y, x=None, do_segments=False, filter_size=13, mode="mirror", breaks=None 642 | ): 643 | """ 644 | Applies a Gaussian smoothing to a curve. 645 | 646 | Parameters 647 | ---------- 648 | y : numpy.ndarray or list of numpy.ndarray 649 | Arrays to be smoothen in the last axis 650 | x : numpy.ndarray 651 | Time array of same shape of `y` last axis used to find data discontinuity. 652 | filter_size : int 653 | Filter window size 654 | mode : str 655 | The `mode` parameter determines how the input array is extended 656 | beyond its boundaries. Options are {'reflect', 'constant', 'nearest', 'mirror', 657 | 'wrap'}. Default is 'mirror' 658 | 659 | Returns 660 | ------- 661 | y_smooth : numpy.ndarray 662 | Smooth array. 663 | """ 664 | if isinstance(y, list): 665 | y = np.asarray(y) 666 | else: 667 | y = np.atleast_2d(y) 668 | 669 | if do_segments: 670 | if breaks is None and x is None: 671 | raise ValueError("Please provide `x` or `breaks` to have splits.") 672 | elif breaks is None and x is not None: 673 | splits = get_breaks(x, include_ext=True) 674 | else: 675 | splits = np.array(breaks) 676 | # find discontinuity in y according to x if provided 677 | if x is not None: 678 | grads = np.gradient(y, x, axis=1) 679 | # the 7-sigma here is hardcoded and found to work ok 680 | splits = np.unique( 681 | np.concatenate( 682 | [splits, np.hstack([np.where(g > 7 * g.std())[0] for g in grads])] 683 | ) 684 | ) 685 | else: 686 | splits = [0, y.shape[-1]] 687 | 688 | y_smooth = [] 689 | for i in range(1, len(splits)): 690 | y_smooth.append( 691 | gaussian_filter1d( 692 | y[:, splits[i - 1] : splits[i]], 693 | filter_size, 694 | mode=mode, 695 | axis=1, 696 | ) 697 | ) 698 | return np.hstack(y_smooth) 699 | 700 | 701 | def bspline_smooth(y, x=None, degree=3, do_segments=False, breaks=None, n_knots=100): 702 | """ 703 | Applies a spline smoothing to a curve. 704 | 705 | Parameters 706 | ---------- 707 | y : numpy.ndarray or list of numpy.ndarray 708 | Arrays to be smoothen in the last axis 709 | x : numpy.ndarray 710 | Optional. x array, as `y = f(x)`` used to find discontinuities in `f(x)`. If x 711 | is given then splits will be computed, if not `breaks` argument as to be provided. 712 | degree : int 713 | Degree of the psline fit, default is 3. 714 | do_segments : boolean 715 | Do the splines per segments with splits computed from data `x` or given in `breaks`. 716 | breaks : list of ints 717 | List of break indexes in `y`. 718 | nknots : int 719 | Number of knots for the B-Spline. If `do_segments` is True, knots will be 720 | distributed in each segment. 721 | 722 | Returns 723 | ------- 724 | y_smooth : numpy.ndarray 725 | Smooth array. 726 | """ 727 | if isinstance(y, list): 728 | y = np.asarray(y) 729 | else: 730 | y = np.atleast_2d(y) 731 | 732 | if do_segments: 733 | if breaks is None and x is None: 734 | raise ValueError("Please provide `x` or `breaks` to have splits.") 735 | elif breaks is None and x is not None: 736 | splits = get_breaks(x) 737 | else: 738 | splits = np.array(breaks) 739 | # find discontinuity in y according to x if provided 740 | if x is not None: 741 | grads = np.gradient(y, x, axis=1) 742 | # the 7-sigma here is hardcoded and found to work ok 743 | splits = np.unique( 744 | np.concatenate( 745 | [splits, np.hstack([np.where(g > 7 * g.std())[0] for g in grads])] 746 | ) 747 | ) 748 | else: 749 | splits = [0, y.shape[-1]] 750 | 751 | y_smooth = [] 752 | v = np.arange(y.shape[-1]) 753 | DM = spline1d(v, np.linspace(v.min(), v.max(), n_knots)).toarray() 754 | # do segments 755 | arr_splits = np.array_split(np.arange(len(v)), splits) 756 | masks = np.asarray( 757 | [np.in1d(np.arange(len(v)), x1).astype(float) for x1 in arr_splits] 758 | ).T 759 | DM = np.hstack([DM[:, idx][:, None] * masks for idx in range(DM.shape[1])]) 760 | 761 | prior_mu = np.zeros(DM.shape[1]) 762 | prior_sigma = np.ones(DM.shape[1]) * 1e5 763 | # iterate over vectors in y 764 | for v in range(y.shape[0]): 765 | weights = solve_linear_model( 766 | DM, 767 | y[v], 768 | prior_mu=prior_mu, 769 | prior_sigma=prior_sigma, 770 | ) 771 | y_smooth.append(DM.dot(weights)) 772 | return np.array(y_smooth) 773 | 774 | 775 | def _load_ffi_image( 776 | telescope, 777 | fname, 778 | extension, 779 | cutout_size=None, 780 | cutout_origin=[0, 0], 781 | return_coords=False, 782 | ): 783 | """ 784 | Use fitsio to load an image and return positions and flux values. 785 | It can do a smal cutout of size `cutout_size` with a defined origin. 786 | 787 | Parameters 788 | ---------- 789 | telescope: str 790 | String for the telescope 791 | fname: str 792 | Path to the filename 793 | extension: int 794 | Extension to cut out of the image 795 | cutout_size : int 796 | Size of the cutout in pixels (e.g. 200) 797 | cutout_origin : tuple of ints 798 | Origin coordinates in pixels from where the cutout stars. Pattern is 799 | [row, column]. 800 | return_coords : bool 801 | Return or not pixel coordinates. 802 | 803 | Return 804 | ------ 805 | f: np.ndarray 806 | Array of flux values read from the file. Shape is [row, column]. 807 | col_2d: np.ndarray 808 | Array of column values read from the file. Shape is [row, column]. Optional. 809 | row_2d: np.ndarray 810 | Array of row values read from the file. Shape is [row, column]. Optional. 811 | """ 812 | f = fitsio.FITS(fname)[extension] 813 | if telescope.lower() == "kepler": 814 | # CCD overscan for Kepler 815 | r_min = 20 816 | r_max = 1044 817 | c_min = 12 818 | c_max = 1112 819 | elif telescope.lower() == "tess": 820 | # CCD overscan for TESS 821 | r_min = 0 822 | r_max = 2048 823 | c_min = 45 824 | c_max = 2093 825 | else: 826 | raise TypeError("File is not from Kepler or TESS mission") 827 | 828 | # If the image dimension is not the FFI shape, we change the r_max and c_max 829 | dims = f.get_dims() 830 | if dims != [r_max, c_max]: 831 | r_max, c_max = np.asarray(dims) 832 | r_min += cutout_origin[0] 833 | c_min += cutout_origin[1] 834 | if (r_min > r_max) | (c_min > c_max): 835 | raise ValueError("`cutout_origin` must be within the image.") 836 | if cutout_size is not None: 837 | r_max = np.min([r_min + cutout_size, r_max]) 838 | c_max = np.min([c_min + cutout_size, c_max]) 839 | if return_coords: 840 | row_2d, col_2d = np.mgrid[r_min:r_max, c_min:c_max] 841 | return col_2d, row_2d, f[r_min:r_max, c_min:c_max] 842 | return f[r_min:r_max, c_min:c_max] 843 | 844 | 845 | def _do_image_cutout( 846 | flux, flux_err, ra, dec, column, row, cutout_size=100, cutout_origin=[0, 0] 847 | ): 848 | """ 849 | Creates a cutout of the full image. Return data arrays corresponding to the cutout. 850 | 851 | Parameters 852 | ---------- 853 | flux : numpy.ndarray 854 | Data array with Flux values, correspond to full size image. 855 | flux_err : numpy.ndarray 856 | Data array with Flux errors values, correspond to full size image. 857 | ra : numpy.ndarray 858 | Data array with RA values, correspond to full size image. 859 | dec : numpy.ndarray 860 | Data array with Dec values, correspond to full size image. 861 | column : numpy.ndarray 862 | Data array with pixel column values, correspond to full size image. 863 | row : numpy.ndarray 864 | Data array with pixel raw values, correspond to full size image. 865 | cutout_size : int 866 | Size in pixels of the cutout, assumedto be squared. Default is 100. 867 | cutout_origin : tuple of ints 868 | Origin of the cutout following matrix indexing. Default is [0 ,0]. 869 | 870 | Returns 871 | ------- 872 | flux : numpy.ndarray 873 | Data array with Flux values of the cutout. 874 | flux_err : numpy.ndarray 875 | Data array with Flux errors values of the cutout. 876 | ra : numpy.ndarray 877 | Data array with RA values of the cutout. 878 | dec : numpy.ndarray 879 | Data array with Dec values of the cutout. 880 | column : numpy.ndarray 881 | Data array with pixel column values of the cutout. 882 | row : numpy.ndarray 883 | Data array with pixel raw values of the cutout. 884 | """ 885 | if (cutout_size + cutout_origin[0] <= flux.shape[1]) and ( 886 | cutout_size + cutout_origin[1] <= flux.shape[2] 887 | ): 888 | column = column[ 889 | cutout_origin[0] : cutout_origin[0] + cutout_size, 890 | cutout_origin[1] : cutout_origin[1] + cutout_size, 891 | ] 892 | row = row[ 893 | cutout_origin[0] : cutout_origin[0] + cutout_size, 894 | cutout_origin[1] : cutout_origin[1] + cutout_size, 895 | ] 896 | flux = flux[ 897 | :, 898 | cutout_origin[0] : cutout_origin[0] + cutout_size, 899 | cutout_origin[1] : cutout_origin[1] + cutout_size, 900 | ] 901 | flux_err = flux_err[ 902 | :, 903 | cutout_origin[0] : cutout_origin[0] + cutout_size, 904 | cutout_origin[1] : cutout_origin[1] + cutout_size, 905 | ] 906 | ra = ra[ 907 | cutout_origin[0] : cutout_origin[0] + cutout_size, 908 | cutout_origin[1] : cutout_origin[1] + cutout_size, 909 | ] 910 | dec = dec[ 911 | cutout_origin[0] : cutout_origin[0] + cutout_size, 912 | cutout_origin[1] : cutout_origin[1] + cutout_size, 913 | ] 914 | else: 915 | raise ValueError("Cutout size is larger than image shape ", flux.shape) 916 | 917 | return flux, flux_err, ra, dec, column, row 918 | -------------------------------------------------------------------------------- /src/psfmachine/version.py: -------------------------------------------------------------------------------- 1 | # It is important to store the version number in a separate file 2 | # so that we can read it from setup.py without importing the package 3 | __version__ = "1.1.4" 4 | -------------------------------------------------------------------------------- /tests/data/kplr-ffi_ch01_test.fits: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSDataLab/psfmachine/6235cc91f2f62c80508668d3ad0648622862450b/tests/data/kplr-ffi_ch01_test.fits -------------------------------------------------------------------------------- /tests/data/shape_model_ffi.fits: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSDataLab/psfmachine/6235cc91f2f62c80508668d3ad0648622862450b/tests/data/shape_model_ffi.fits -------------------------------------------------------------------------------- /tests/data/tpf_test_00.fits: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSDataLab/psfmachine/6235cc91f2f62c80508668d3ad0648622862450b/tests/data/tpf_test_00.fits -------------------------------------------------------------------------------- /tests/data/tpf_test_01.fits: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSDataLab/psfmachine/6235cc91f2f62c80508668d3ad0648622862450b/tests/data/tpf_test_01.fits -------------------------------------------------------------------------------- /tests/data/tpf_test_02.fits: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSDataLab/psfmachine/6235cc91f2f62c80508668d3ad0648622862450b/tests/data/tpf_test_02.fits -------------------------------------------------------------------------------- /tests/data/tpf_test_03.fits: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSDataLab/psfmachine/6235cc91f2f62c80508668d3ad0648622862450b/tests/data/tpf_test_03.fits -------------------------------------------------------------------------------- /tests/data/tpf_test_04.fits: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSDataLab/psfmachine/6235cc91f2f62c80508668d3ad0648622862450b/tests/data/tpf_test_04.fits -------------------------------------------------------------------------------- /tests/data/tpf_test_05.fits: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSDataLab/psfmachine/6235cc91f2f62c80508668d3ad0648622862450b/tests/data/tpf_test_05.fits -------------------------------------------------------------------------------- /tests/data/tpf_test_06.fits: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSDataLab/psfmachine/6235cc91f2f62c80508668d3ad0648622862450b/tests/data/tpf_test_06.fits -------------------------------------------------------------------------------- /tests/data/tpf_test_07.fits: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSDataLab/psfmachine/6235cc91f2f62c80508668d3ad0648622862450b/tests/data/tpf_test_07.fits -------------------------------------------------------------------------------- /tests/data/tpf_test_08.fits: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSDataLab/psfmachine/6235cc91f2f62c80508668d3ad0648622862450b/tests/data/tpf_test_08.fits -------------------------------------------------------------------------------- /tests/data/tpf_test_09.fits: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SSDataLab/psfmachine/6235cc91f2f62c80508668d3ad0648622862450b/tests/data/tpf_test_09.fits -------------------------------------------------------------------------------- /tests/test_ffimachine.py: -------------------------------------------------------------------------------- 1 | """ 2 | This contains a collection of functions to test the Machine API 3 | """ 4 | import os 5 | 6 | # import numpy as np 7 | import pytest 8 | from astropy.io import fits 9 | from astropy.table import Table 10 | from astropy.utils.data import get_pkg_data_filename 11 | 12 | from psfmachine import Machine, FFIMachine 13 | 14 | 15 | @pytest.mark.remote_data 16 | def test_ffi_from_file(): 17 | 18 | ffi_path = get_pkg_data_filename("./data/kplr-ffi_ch01_test.fits") 19 | ffi = FFIMachine.from_file(ffi_path, extension=1) 20 | # test `FFIMachine.from_file` is of Machine class 21 | assert isinstance(ffi, Machine) 22 | # test attributes have the right shapes 23 | assert ffi.time.shape == (1,) 24 | assert ffi.flux.shape == (1, 33840) 25 | assert ffi.flux_2d.shape == (1, 180, 188) 26 | assert ffi.flux_err.shape == (1, 33840) 27 | assert ffi.column.shape == (33840,) 28 | assert ffi.row.shape == (33840,) 29 | assert ffi.ra.shape == (33840,) 30 | assert ffi.dec.shape == (33840,) 31 | assert ffi.sources.shape == (259, 15) 32 | 33 | 34 | @pytest.mark.remote_data 35 | def test_save_shape_model(): 36 | ffi_path = get_pkg_data_filename("./data/kplr-ffi_ch01_test.fits") 37 | ffi = FFIMachine.from_file(ffi_path, extension=1) 38 | # create a shape model 39 | ffi.build_shape_model() 40 | file_name = "%s/data/test_ffi_shape_model.fits" % os.path.abspath( 41 | os.path.dirname(__file__) 42 | ) 43 | # save shape model to tmp file 44 | ffi.save_shape_model(file_name) 45 | 46 | # test that the saved shape model has the right metadata and shape 47 | shape_model = fits.open(file_name) 48 | assert shape_model[1].header["OBJECT"] == "PRF shape" 49 | assert shape_model[1].header["TELESCOP"] == ffi.meta["TELESCOP"] 50 | assert shape_model[1].header["MJD-OBS"] == ffi.time[0] 51 | 52 | assert shape_model[1].header["N_RKNOTS"] == ffi.n_r_knots 53 | assert shape_model[1].header["N_PKNOTS"] == ffi.n_phi_knots 54 | assert shape_model[1].header["RMIN"] == ffi.rmin 55 | assert shape_model[1].header["RMAX"] == ffi.rmax 56 | assert shape_model[1].header["CUT_R"] == ffi.cut_r 57 | assert shape_model[1].data.shape == ffi.psf_w.shape 58 | 59 | os.remove(file_name) 60 | 61 | 62 | @pytest.mark.remote_data 63 | def test_save_flux_values(): 64 | ffi_path = get_pkg_data_filename("./data/kplr-ffi_ch01_test.fits") 65 | ffi = FFIMachine.from_file(ffi_path, extension=1) 66 | ffi.build_shape_model() 67 | file_name = "%s/data/ffi_test_phot.fits" % os.path.abspath( 68 | os.path.dirname(__file__) 69 | ) 70 | # fit the shape model to sources, compute psf photometry, save catalog to disk 71 | ffi.save_flux_values(file_name) 72 | 73 | # check FITS file has the right metadata 74 | hdu = fits.open(file_name) 75 | assert hdu[0].header["TELESCOP"] == ffi.meta["TELESCOP"] 76 | assert hdu[0].header["DCT_TYPE"] == ffi.meta["DCT_TYPE"] 77 | assert hdu[1].header["MJD-OBS"] == ffi.time[0] 78 | 79 | # check FITS table has the right shapes, columns and units 80 | table = Table.read(file_name) 81 | assert len(table) == 259 82 | assert "psf_flux" in table.keys() 83 | assert "psf_flux_err" in table.keys() 84 | assert "gaia_id" in table.keys() 85 | assert "ra" in table.keys() 86 | assert "dec" in table.keys() 87 | assert table["psf_flux"].unit == "-e/s" 88 | assert table["psf_flux_err"].unit == "-e/s" 89 | assert table["ra"].unit == "deg" 90 | assert table["dec"].unit == "deg" 91 | 92 | os.remove(file_name) 93 | -------------------------------------------------------------------------------- /tests/test_machine.py: -------------------------------------------------------------------------------- 1 | """ 2 | This contains a collection of functions to test the Machine API 3 | """ 4 | import numpy as np 5 | from scipy import sparse 6 | import pytest 7 | import lightkurve as lk 8 | from astropy.utils.data import get_pkg_data_filename 9 | import astropy.units as u 10 | 11 | from psfmachine import TPFMachine 12 | 13 | tpfs = [] 14 | for idx in range(10): 15 | tpfs.append(lk.read(get_pkg_data_filename(f"data/tpf_test_{idx:02}.fits"))) 16 | tpfs = lk.collections.TargetPixelFileCollection(tpfs) 17 | 18 | 19 | @pytest.mark.remote_data 20 | def test_create_delta_sparse_arrays(): 21 | machine = TPFMachine.from_TPFs(tpfs) 22 | # create numpy arrays 23 | machine._create_delta_arrays() 24 | non_sparse_arr = machine.__dict__.copy() 25 | 26 | # check for main attrs shape 27 | assert non_sparse_arr["time"].shape == (10,) 28 | assert non_sparse_arr["flux"].shape == (10, 287) 29 | assert non_sparse_arr["flux_err"].shape == (10, 287) 30 | assert non_sparse_arr["column"].shape == (287,) 31 | assert non_sparse_arr["row"].shape == (287,) 32 | assert non_sparse_arr["ra"].shape == (287,) 33 | assert non_sparse_arr["dec"].shape == (287,) 34 | assert non_sparse_arr["sources"].shape == (19, 15) 35 | 36 | assert non_sparse_arr["dra"].shape == (19, 287) 37 | assert non_sparse_arr["ddec"].shape == (19, 287) 38 | assert non_sparse_arr["r"].shape == (19, 287) 39 | assert non_sparse_arr["phi"].shape == (19, 287) 40 | # check dra ddec r and phi are numpy arrays 41 | assert isinstance(non_sparse_arr["dra"], np.ndarray) 42 | assert isinstance(non_sparse_arr["ddec"], np.ndarray) 43 | assert isinstance(non_sparse_arr["r"], np.ndarray) 44 | assert isinstance(non_sparse_arr["phi"], np.ndarray) 45 | 46 | # manually mask numpy arrays to compare them vs sparse array 47 | dist_lim = 40 48 | mask = (np.abs(non_sparse_arr["dra"].value) <= dist_lim / 3600) & ( 49 | np.abs(non_sparse_arr["ddec"].value) <= dist_lim / 3600 50 | ) 51 | 52 | # create sparse arrays 53 | machine.sparse_dist_lim = dist_lim * u.arcsecond 54 | machine._create_delta_sparse_arrays() 55 | sparse_arr = machine.__dict__.copy() 56 | 57 | assert sparse_arr["dra"].shape == non_sparse_arr["dra"].shape 58 | assert sparse_arr["ddec"].shape == non_sparse_arr["ddec"].shape 59 | assert sparse_arr["r"].shape == non_sparse_arr["r"].shape 60 | assert sparse_arr["phi"].shape == non_sparse_arr["phi"].shape 61 | # check dra ddec r and phi are sparse arrays 62 | assert isinstance(sparse_arr["dra"], sparse.csr_matrix) 63 | assert isinstance(sparse_arr["ddec"], sparse.csr_matrix) 64 | assert isinstance(sparse_arr["r"], sparse.csr_matrix) 65 | assert isinstance(sparse_arr["phi"], sparse.csr_matrix) 66 | # check for non-zero values shape 67 | assert sparse_arr["dra"].data.shape == (861,) 68 | assert sparse_arr["ddec"].data.shape == (861,) 69 | assert sparse_arr["r"].data.shape == (861,) 70 | assert sparse_arr["phi"].data.shape == (861,) 71 | 72 | # compare sparse array vs numpy array values 73 | assert (non_sparse_arr["dra"][mask].value == sparse_arr["dra"].data).all() 74 | assert (non_sparse_arr["ddec"][mask].value == sparse_arr["ddec"].data).all() 75 | assert (non_sparse_arr["r"][mask].value == sparse_arr["r"].data).all() 76 | assert (non_sparse_arr["phi"][mask].value == sparse_arr["phi"].data).all() 77 | 78 | 79 | @pytest.mark.remote_data 80 | def test_compute_aperture_photometry(): 81 | # it tests aperture mask creation and flux metric computation 82 | machine = TPFMachine.from_TPFs(tpfs, apply_focus_mask=False) 83 | # load FFI shape model from file 84 | machine.load_shape_model( 85 | input=get_pkg_data_filename("data/shape_model_ffi.fits"), 86 | plot=False, 87 | ) 88 | # compute max aperture 89 | machine.create_aperture_mask(percentile=0) 90 | assert machine.aperture_mask.shape == (19, 287) 91 | # some sources way outside the TPF can have 0-size aperture 92 | assert (machine.aperture_mask.sum(axis=1) >= 0).all() 93 | assert machine.FLFRCSAP.shape == (19,) 94 | assert machine.CROWDSAP.shape == (19,) 95 | # full apereture shoudl lead to FLFRCSAP = 1 96 | assert np.allclose(machine.FLFRCSAP[np.isfinite(machine.FLFRCSAP)], 1) 97 | assert (machine.CROWDSAP[np.isfinite(machine.CROWDSAP)] >= 0).all() 98 | assert (machine.CROWDSAP[np.isfinite(machine.CROWDSAP)] <= 1).all() 99 | 100 | # compute min aperture, here CROWDSAP not always will be 1, e.g. 2 sources in the 101 | # same pixel. 102 | machine.create_aperture_mask(percentile=99) 103 | assert (machine.CROWDSAP[np.isfinite(machine.CROWDSAP)] >= 0).all() 104 | assert (machine.CROWDSAP[np.isfinite(machine.CROWDSAP)] <= 1).all() 105 | assert (machine.FLFRCSAP[np.isfinite(machine.FLFRCSAP)] >= 0).all() 106 | assert (machine.FLFRCSAP[np.isfinite(machine.FLFRCSAP)] <= 1).all() 107 | 108 | machine.compute_aperture_photometry(aperture_size="optimal") 109 | assert machine.aperture_mask.shape == (19, 287) 110 | # some sources way outside the TPF can have 0-size aperture 111 | assert (machine.aperture_mask.sum(axis=1) >= 0).all() 112 | # check aperture size is within range 113 | assert (machine.optimal_percentile >= 0).all() and ( 114 | machine.optimal_percentile <= 100 115 | ).all() 116 | assert machine.sap_flux.shape == (10, 19) 117 | assert machine.sap_flux_err.shape == (10, 19) 118 | assert (machine.sap_flux >= 0).all() 119 | assert (machine.sap_flux_err >= 0).all() 120 | 121 | 122 | @pytest.mark.remote_data 123 | def test_psf_metrics(): 124 | machine = TPFMachine.from_TPFs(tpfs, apply_focus_mask=False) 125 | # load FFI shape model from file 126 | machine.build_shape_model(plot=False) 127 | 128 | machine.build_time_model(segments=False, bin_method="bin", focus=False) 129 | 130 | # finite & possitive normalization 131 | assert machine.normalized_shape_model 132 | assert np.isfinite(machine.mean_model_integral) 133 | assert machine.mean_model_integral > 0 134 | 135 | machine.get_psf_metrics() 136 | assert np.isfinite(machine.source_psf_fraction).all() 137 | assert (machine.source_psf_fraction >= 0).all() 138 | 139 | # this ratio is nan for sources with no pixels in the source_mask then zero-division 140 | assert np.isclose( 141 | machine.perturbed_ratio_mean[np.isfinite(machine.perturbed_ratio_mean)], 142 | 1, 143 | atol=1e-2, 144 | ).all() 145 | 146 | # all should be finite because std(0s) = 0 147 | assert np.isfinite(machine.perturbed_std).all() 148 | assert (machine.perturbed_std >= 0).all() 149 | 150 | 151 | @pytest.mark.remote_data 152 | def test_poscorr_smooth(): 153 | machine = TPFMachine.from_TPFs(tpfs, apply_focus_mask=False) 154 | machine.build_shape_model(plot=False) 155 | # no segments 156 | machine.poscorr_filter_size = 1 157 | machine.build_time_model( 158 | segments=False, bin_method="bin", focus=False, positions="poscorr" 159 | ) 160 | 161 | median_pc1 = np.nanmedian(machine.pos_corr1, axis=0) 162 | median_pc2 = np.nanmedian(machine.pos_corr2, axis=0) 163 | median_pc1 = (median_pc1 - median_pc1.mean()) / ( 164 | median_pc1.max() - median_pc1.mean() 165 | ) 166 | median_pc2 = (median_pc2 - median_pc2.mean()) / ( 167 | median_pc2.max() - median_pc2.mean() 168 | ) 169 | 170 | assert np.isclose(machine.P.other_vectors[:, 0], median_pc1, atol=0.5).all() 171 | assert np.isclose(machine.P.other_vectors[:, 1], median_pc2, atol=0.5).all() 172 | 173 | 174 | @pytest.mark.remote_data 175 | def test_segment_time_model(): 176 | # testing segment with the current test dataset we have that only has 10 cadences 177 | # isn't the best, but we can still do some sanity checks. 178 | machine = TPFMachine.from_TPFs(tpfs, apply_focus_mask=False, time_resolution=3) 179 | machine.build_shape_model(plot=False) 180 | # no segments 181 | machine.build_time_model(segments=False, bin_method="bin", focus=False) 182 | assert machine.P.vectors.shape == (10, 4) 183 | 184 | # fake 2 time breaks 185 | machine.time[4:] += 0.5 186 | machine.time[7:] += 0.41 187 | # user defined segments 188 | machine.build_time_model(segments=True, bin_method="bin", focus=False) 189 | assert machine.P.vectors.shape == (10, 4 * 3) 190 | -------------------------------------------------------------------------------- /tests/test_perturbation.py: -------------------------------------------------------------------------------- 1 | """test perturbation matrix""" 2 | import numpy as np 3 | from scipy import sparse 4 | import pytest 5 | from psfmachine.perturbation import PerturbationMatrix, PerturbationMatrix3D 6 | 7 | 8 | def test_perturbation_matrix(): 9 | time = np.arange(0, 10, 0.1) 10 | p = PerturbationMatrix(time=time, focus=False) 11 | assert p.vectors.shape == (100, 4) 12 | p = PerturbationMatrix(time=time, focus=True) 13 | assert p.vectors.shape == (100, 5) 14 | 15 | with pytest.raises(ValueError): 16 | p = PerturbationMatrix( 17 | time=time, other_vectors=np.random.normal(size=(2, 10)), focus=False 18 | ) 19 | p = PerturbationMatrix(time=time, other_vectors=1, focus=False) 20 | p = PerturbationMatrix( 21 | time=time, other_vectors=np.random.normal(size=(2, 100)), focus=False 22 | ) 23 | p = PerturbationMatrix( 24 | time=time, other_vectors=np.random.normal(size=(100, 2)), focus=False 25 | ) 26 | assert p.vectors.shape == (100, 6) 27 | time = np.hstack([np.arange(0, 10, 0.1), np.arange(15, 25, 0.1)]) 28 | p = PerturbationMatrix(time=time, focus=False) 29 | assert p.vectors.shape == (200, 8) 30 | p = PerturbationMatrix(time=time, focus=True) 31 | assert p.vectors.shape == (200, 10) 32 | 33 | assert p.matrix.shape == (200 / 10, 10) 34 | assert sparse.issparse(p.matrix) 35 | 36 | res = 10 37 | p = PerturbationMatrix(time=time, focus=False, resolution=res) 38 | y, ye = np.random.normal(1, 0.01, size=200), np.ones(200) * 0.01 39 | p.fit(y, ye) 40 | w = p.weights 41 | assert w.shape[0] == p.shape[1] 42 | assert np.isfinite(w).all() 43 | model = p.model() 44 | assert model.shape == y.shape 45 | chi = np.sum((y - model) ** 2 / (ye ** 2)) / (p.shape[0] - p.shape[1] - 1) 46 | assert chi < 3 47 | 48 | y, ye = np.random.normal(1, 0.01, size=200), np.ones(200) * 0.01 49 | for bin_method in ["downsample", "bin"]: 50 | s = 200 / res + 1 if bin_method == "downsample" else 200 / res 51 | p = PerturbationMatrix( 52 | time=time, focus=False, resolution=res, bin_method=bin_method 53 | ) 54 | assert len(p.bin_func(y)) == s 55 | assert len(p.bin_func(ye, quad=True)) == s 56 | with pytest.raises(ValueError): 57 | p.bin_func(y[:-4]) 58 | 59 | p.fit(y, ye) 60 | w = p.weights 61 | model = p.model() 62 | assert model.shape[0] == 200 63 | chi = np.sum((y - model) ** 2 / (ye ** 2)) / (p.shape[0] - p.shape[1] - 1) 64 | assert chi < 3 65 | 66 | # Test PCA 67 | flux = np.random.normal(1, 0.1, size=(200, 100)) 68 | p = PerturbationMatrix(time=time, focus=False) 69 | assert p.matrix.shape == (20, 8) 70 | p.pca(flux, ncomponents=2) 71 | assert p.matrix.shape == (20, 12) 72 | # assert np.allclose((p.vectors.sum(axis=0) / (p.vectors != 0).sum(axis=0))[8:], 0) 73 | p.fit(y, ye) 74 | p = PerturbationMatrix(time=time, focus=False, segments=False) 75 | assert p.matrix.shape == (20, 4) 76 | p.pca(flux, ncomponents=2) 77 | assert p.matrix.shape == (20, 6) 78 | # assert np.allclose((p.vectors.sum(axis=0) / (p.vectors != 0).sum(axis=0))[8:], 0) 79 | p.fit(y, ye) 80 | 81 | 82 | def test_perturbation_matrix3d(): 83 | time = np.arange(0, 10, 1) 84 | # 13 x 13 pixels, evenly spaced in x and y 85 | dx, dy = np.mgrid[:13, :13] - 6 + 0.01 86 | dx, dy = dx.ravel(), dy.ravel() 87 | 88 | # ntime x npixels 89 | flux = np.random.normal(1, 0.01, size=(10, 169)) + dx[None, :] * dy[None, :] 90 | # the perturbation model assumes the perturbation is around 1 91 | flux_err = np.ones((10, 169)) * 0.01 92 | 93 | p3 = PerturbationMatrix3D( 94 | time=time, dx=dx, dy=dy, nknots=4, radius=5, resolution=5, poly_order=1 95 | ) 96 | assert p3.cartesian_matrix.shape == (169, 81) 97 | assert p3.vectors.shape == (10, 2) 98 | assert p3.shape == ( 99 | p3.cartesian_matrix.shape[0] * p3.ntime, 100 | p3.cartesian_matrix.shape[1] * p3.nvec, 101 | ) 102 | assert p3.matrix.shape == ( 103 | p3.cartesian_matrix.shape[0] * p3.nbins, 104 | p3.cartesian_matrix.shape[1] * p3.nvec, 105 | ) 106 | p3.fit(flux, flux_err) 107 | w = p3.weights 108 | assert w.shape[0] == p3.cartesian_matrix.shape[1] * p3.nvec 109 | model = p3.model() 110 | assert model.shape == flux.shape 111 | 112 | chi = np.sum((flux - model) ** 2 / (flux_err ** 2)) / ( 113 | p3.shape[0] - p3.shape[1] - 1 114 | ) 115 | assert chi < 1.5 116 | 117 | time = np.arange(0, 100, 1) 118 | flux = np.random.normal(1, 0.01, size=(100, 169)) + dx[None, :] * dy[None, :] 119 | # the perturbation model assumes the perturbation is around 1 120 | flux_err = np.ones((100, 169)) * 0.01 121 | 122 | for bin_method in ["downsample", "bin"]: 123 | p3 = PerturbationMatrix3D( 124 | time=time, 125 | dx=dx, 126 | dy=dy, 127 | nknots=4, 128 | radius=5, 129 | poly_order=2, 130 | bin_method=bin_method, 131 | ) 132 | p3.fit(flux, flux_err) 133 | w = p3.weights 134 | model = p3.model() 135 | assert model.shape == flux.shape 136 | chi = np.sum((flux - model) ** 2 / (flux_err ** 2)) / ( 137 | p3.shape[0] - p3.shape[1] - 1 138 | ) 139 | assert chi < 3 140 | 141 | p3 = PerturbationMatrix3D( 142 | time=time, 143 | dx=dx, 144 | dy=dy, 145 | nknots=4, 146 | radius=5, 147 | poly_order=2, 148 | bin_method=bin_method, 149 | ) 150 | p3.pca(flux, ncomponents=5) 151 | p3.fit(flux, flux_err) 152 | 153 | # Add in one bad pixel 154 | flux[:, 100] = 1e5 155 | pixel_mask = np.ones(169, bool) 156 | pixel_mask[100] = False 157 | 158 | for bin_method in ["downsample", "bin"]: 159 | p3 = PerturbationMatrix3D( 160 | time=time, 161 | dx=dx, 162 | dy=dy, 163 | nknots=4, 164 | radius=5, 165 | poly_order=2, 166 | bin_method=bin_method, 167 | ) 168 | p3.fit(flux, flux_err) 169 | w = p3.weights 170 | model = p3.model() 171 | chi = np.sum( 172 | (flux[:, pixel_mask] - model[:, pixel_mask]) ** 2 173 | / (flux_err[:, pixel_mask] ** 2) 174 | ) / (p3.shape[0] - p3.shape[1] - 1) 175 | # Without the pixel masking the model doesn't fit 176 | assert chi > 3 177 | p3.fit(flux, flux_err, pixel_mask=pixel_mask) 178 | w = p3.weights 179 | model = p3.model() 180 | chi = np.sum( 181 | (flux[:, pixel_mask] - model[:, pixel_mask]) ** 2 182 | / (flux_err[:, pixel_mask] ** 2) 183 | ) / (p3.shape[0] - p3.shape[1] - 1) 184 | # with pixel masking, it should fit 185 | assert chi < 3 186 | -------------------------------------------------------------------------------- /tests/test_tpfmachine.py: -------------------------------------------------------------------------------- 1 | """ 2 | This contains a collection of functions to test the Machine API 3 | """ 4 | import os 5 | import numpy as np 6 | import pandas as pd 7 | import pytest 8 | import lightkurve as lk 9 | from astropy.utils.data import get_pkg_data_filename 10 | from astropy.time import Time 11 | 12 | from psfmachine import PACKAGEDIR 13 | from psfmachine import Machine, TPFMachine 14 | from psfmachine.tpf import ( 15 | _parse_TPFs, 16 | _wcs_from_tpfs, 17 | _preprocess, 18 | _get_coord_and_query_gaia, 19 | _clean_source_list, 20 | ) 21 | 22 | from psfmachine.utils import do_tiled_query 23 | 24 | tpfs = [] 25 | for idx in range(10): 26 | tpfs.append(lk.read(get_pkg_data_filename(f"data/tpf_test_{idx:02}.fits"))) 27 | tpfs = lk.collections.TargetPixelFileCollection(tpfs) 28 | 29 | 30 | @pytest.mark.remote_data 31 | def test_parse_TPFs(): 32 | ( 33 | times, 34 | flux, 35 | flux_err, 36 | pos_corr1, 37 | pos_corr2, 38 | column, 39 | row, 40 | unw, 41 | focus_mask, 42 | qual_mask, 43 | sat_mask, 44 | ) = _parse_TPFs(tpfs) 45 | 46 | assert times.shape == (10,) 47 | assert flux.shape == (10, 345) 48 | assert flux_err.shape == (10, 345) 49 | assert unw.shape == (345,) 50 | assert column.shape == (345,) 51 | assert row.shape == (345,) 52 | assert focus_mask.shape == (10,) 53 | assert qual_mask.shape == (10,) 54 | assert sat_mask.shape == (345,) 55 | assert pos_corr1.shape == (10, 10) 56 | assert pos_corr2.shape == (10, 10) 57 | 58 | locs, ra, dec = _wcs_from_tpfs(tpfs) 59 | assert locs.shape == (2, 345) 60 | assert ra.shape == (345,) 61 | assert dec.shape == (345,) 62 | 63 | # save-keep valid pixels for later testing 64 | mask = np.isfinite(flux).all(axis=0) & ~sat_mask 65 | locs_s = [f"{x[0]}_{x[1]}" for x in locs.T[mask]] 66 | 67 | flux, flux_err, unw, locs, ra, dec, column, row = _preprocess( 68 | flux, 69 | flux_err, 70 | unw, 71 | locs, 72 | ra, 73 | dec, 74 | column, 75 | row, 76 | tpfs, 77 | sat_mask, 78 | ) 79 | 80 | # check all original pixels show only once after preprocess 81 | # cheking no duplicated pixels remain 82 | _, unique_idx = np.unique(locs, axis=1, return_index=True) 83 | unique_idx = np.in1d(np.arange(len(ra)), unique_idx) 84 | assert (~unique_idx).sum() == 0 85 | 86 | # valid non-duplicated pixels 87 | locs_su = [f"{x[0]}_{x[1]}" for x in locs.T] 88 | # check we don't loose valid pixels 89 | assert np.isin(locs_s, locs_su).all() 90 | 91 | assert np.isfinite(flux).all() 92 | assert np.isfinite(flux_err).all() 93 | assert np.isfinite(locs).all() 94 | assert np.isfinite(ra).all() 95 | assert np.isfinite(dec).all() 96 | 97 | assert locs.shape == (2, 287) 98 | assert ra.shape == (287,) 99 | assert dec.shape == (287,) 100 | assert flux.shape == (10, 287) 101 | assert flux_err.shape == (10, 287) 102 | 103 | sources = _get_coord_and_query_gaia(tpfs) 104 | 105 | assert isinstance(sources, pd.DataFrame) 106 | assert set(["ra", "dec", "phot_g_mean_mag"]).issubset(sources.columns) 107 | assert sources.shape == (21, 14) 108 | 109 | 110 | @pytest.mark.remote_data 111 | def test_from_TPFs(): 112 | c = TPFMachine.from_TPFs(tpfs) 113 | assert isinstance(c, Machine) 114 | 115 | 116 | @pytest.mark.remote_data 117 | def test_do_tiled_query(): 118 | # unit test for TPF stack 119 | # test that the tiled query get the same results as the original psfmachine query 120 | epoch = Time(tpfs[0].time[len(tpfs[0]) // 2], format="jd").jyear 121 | sources_org = _get_coord_and_query_gaia(tpfs, magnitude_limit=18, dr=3) 122 | _, ra, dec = _wcs_from_tpfs(tpfs) 123 | sources_tiled = do_tiled_query( 124 | ra, 125 | dec, 126 | ngrid=(2, 2), 127 | magnitude_limit=18, 128 | dr=3, 129 | epoch=epoch, 130 | ) 131 | assert isinstance(sources_tiled, pd.DataFrame) 132 | assert set(["ra", "dec", "phot_g_mean_mag"]).issubset(sources_tiled.columns) 133 | # check that the tiled query contain all sources from the non-tiled query. 134 | # tiled query is always bigger that the other for TPF stacks. 135 | assert set(sources_org.designation).issubset(sources_tiled.designation) 136 | 137 | # get ra,dec values for pixels then clean source list 138 | ras, decs = [], [] 139 | for tpf in tpfs: 140 | r, d = np.hstack(tpf.get_coordinates(0)).T.reshape( 141 | [2, np.product(tpf.shape[1:])] 142 | ) 143 | ras.append(r) 144 | decs.append(d) 145 | ras, decs = np.hstack(ras), np.hstack(decs) 146 | sources_tiled, _ = _clean_source_list(sources_tiled, ras, decs) 147 | # clean source lists must match between query versions 148 | assert sources_tiled.shape == sources_org.shape 149 | assert set(sources_org.designation) == set(sources_tiled.designation) 150 | 151 | # Unit test for 360->0 deg boundary. we use a smaller sky patch now. 152 | row = np.arange(100) 153 | column = np.arange(100) 154 | column_grid, row_grid = np.meshgrid(column, row) 155 | # I subtract pix position to move the grid into the 360->0 ra boundary 156 | column_grid -= 44200 157 | ra, dec = ( 158 | tpfs[0] 159 | .wcs.wcs_pix2world( 160 | np.vstack([column_grid.ravel(), row_grid.ravel()]).T, 161 | 0.0, 162 | ) 163 | .T 164 | ) 165 | # check that ra values are in the boundary 166 | assert not ((ra < 359) & (ra > 1)).all() 167 | boundary_sources = do_tiled_query( 168 | ra, dec, ngrid=(2, 2), magnitude_limit=16, epoch=epoch, dr=3 169 | ) 170 | assert boundary_sources.shape == (85, 13) 171 | # check that no result objects are outside the boundary for ra 172 | assert not ((boundary_sources.ra < 359) & (boundary_sources.ra > 1)).all() 173 | 174 | 175 | @pytest.mark.remote_data 176 | def test_load_save_shape_model(): 177 | # instantiate a machine object 178 | machine = TPFMachine.from_TPFs(tpfs, apply_focus_mask=False) 179 | # build a shape model from TPF data 180 | machine.build_shape_model(plot=False) 181 | # save object state 182 | org_state = machine.__dict__ 183 | # save shape model to disk 184 | file_name = "%s/data/test_shape_model.fits" % os.path.abspath( 185 | os.path.dirname(__file__) 186 | ) 187 | machine.save_shape_model(output=file_name) 188 | 189 | # instantiate a new machine object with same data but load shape model from disk 190 | machine = TPFMachine.from_TPFs(tpfs, apply_focus_mask=False) 191 | machine.load_shape_model(plot=False, input=file_name) 192 | new_state = machine.__dict__ 193 | 194 | # check that critical attributes match 195 | assert org_state["n_r_knots"] == new_state["n_r_knots"] 196 | assert org_state["n_phi_knots"] == new_state["n_phi_knots"] 197 | assert org_state["rmin"] == new_state["rmin"] 198 | assert org_state["rmax"] == new_state["rmax"] 199 | assert org_state["cut_r"] == new_state["cut_r"] 200 | assert (org_state["psf_w"] == new_state["psf_w"]).all() 201 | assert ((org_state["mean_model"] == new_state["mean_model"]).data).all() 202 | # remove shape model from disk 203 | os.remove(file_name) 204 | 205 | 206 | @pytest.mark.remote_data 207 | def test_load_shape_model_from_zenodo(): 208 | # instantiate a machine object 209 | machine = TPFMachine.from_TPFs(tpfs, apply_focus_mask=False) 210 | 211 | # load a FFI-PRF model from zenodo repo 212 | machine.load_shape_model(plot=False, input=None) 213 | state = machine.__dict__ 214 | 215 | # check that critical attributes match 216 | assert state["n_r_knots"] == 10 217 | assert state["n_phi_knots"] == 15 218 | assert state["rmin"] == 1 219 | assert state["rmax"] == 16 220 | assert state["cut_r"] == 6 221 | assert state["psf_w"].shape == (169,) 222 | 223 | # check file is in directory for future use 224 | assert os.path.isfile( 225 | f"{PACKAGEDIR}/data/" 226 | f"{machine.tpf_meta['mission'][0]}_FFI_PRFmodels_v1.0.tar.gz" 227 | ) 228 | 229 | 230 | @pytest.mark.remote_data 231 | def test_get_source_centroids(): 232 | # test centroid fx 233 | machine = TPFMachine.from_TPFs(tpfs, apply_focus_mask=False) 234 | machine._get_source_mask() 235 | machine.get_source_centroids(method="poscor") 236 | # check centroid arrays have the right shape 237 | assert machine.source_centroids_column_poscor.shape == (19, 10) 238 | assert machine.source_centroids_row_poscor.shape == (19, 10) 239 | machine.get_source_centroids(method="scene") 240 | assert machine.source_centroids_column_scene.shape == (19, 10) 241 | assert machine.source_centroids_row_scene.shape == (19, 10) 242 | 243 | tpf_idx = [] 244 | for i in range(len(machine.sources)): 245 | tpf_idx.append( 246 | [k for k, ss in enumerate(machine.tpf_meta["sources"]) if i in ss] 247 | ) 248 | # check centroids agree within a pixel for target sources when comparing vs 249 | # pipeline aperture moments aperture TPF.estimate_centroids() 250 | for i in range(machine.nsources): 251 | if machine.sources.tpf_id[i]: 252 | if len(tpf_idx[i]) > 1: 253 | continue 254 | col_cent, row_cent = tpfs[tpf_idx[i][0]].estimate_centroids() 255 | assert ( 256 | np.abs(col_cent - machine.source_centroids_column_scene[i]).value < 1 257 | ).all() 258 | assert ( 259 | np.abs(row_cent - machine.source_centroids_row_scene[i]).value < 1 260 | ).all() 261 | assert ( 262 | np.abs(col_cent - machine.source_centroids_column_poscor[i]).value < 1 263 | ).all() 264 | assert ( 265 | np.abs(row_cent - machine.source_centroids_row_poscor[i]).value < 1 266 | ).all() 267 | -------------------------------------------------------------------------------- /tests/test_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from psfmachine.utils import threshold_bin 3 | 4 | 5 | def test_threshold_bin(): 6 | 7 | # Pass a bunch of nans 8 | x, y, z = ( 9 | np.random.normal(size=1000), 10 | np.random.normal(size=1000), 11 | np.ones(1000) * np.nan, 12 | ) 13 | hist, xbin, ybin, zbin, zwbin = threshold_bin(x, y, z) 14 | assert xbin.shape[0] <= x.shape[0] 15 | assert np.isnan(zbin).all() 16 | assert hist.min() >= 1.0 17 | 18 | # Pass something that should actually pass 19 | x = np.random.normal(0, 1, 5000) 20 | y = np.random.normal(0, 1, 5000) 21 | z = np.power(10, -(((x ** 2 + y ** 2) / 0.5) ** 0.5) / 2) 22 | hist, xbin, ybin, zbin, zwbin = threshold_bin(x, y, z, abs_thresh=5, bins=50) 23 | # check that some points are binned, some points are not 24 | assert hist.shape == xbin.shape 25 | assert hist.shape == ybin.shape 26 | assert hist.shape == zbin.shape 27 | assert hist.shape == zwbin.shape 28 | assert xbin.shape[0] <= x.shape[0] 29 | assert np.isin(zbin, z).sum() > 0 30 | assert hist.min() >= 1.0 31 | 32 | # add random nan values to Z 33 | z[np.random.randint(0, 5000, 50)] = np.nan 34 | hist2, xbin2, ybin2, zbin2, zwbin = threshold_bin(x, y, z, abs_thresh=5, bins=50) 35 | # check that some points are binned, some points are not 36 | assert xbin2.shape[0] <= x.shape[0] 37 | assert np.isin(zbin2, z).sum() > 0 38 | assert hist2.min() >= 1.0 39 | # We should get the same outut for (hist, xbin, ybin) and same shape for zbin 40 | assert (hist == hist2).all() 41 | assert (xbin == xbin2).all() 42 | assert (ybin == ybin2).all() 43 | assert zbin.shape == hist2.shape 44 | --------------------------------------------------------------------------------