pysersic 0.1.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- pysersic-0.1.1/LICENSE +21 -0
- pysersic-0.1.1/PKG-INFO +8 -0
- pysersic-0.1.1/README.md +68 -0
- pysersic-0.1.1/pysersic/__init__.py +4 -0
- pysersic-0.1.1/pysersic/galfit.py +190 -0
- pysersic-0.1.1/pysersic/loss.py +328 -0
- pysersic-0.1.1/pysersic/priors.py +810 -0
- pysersic-0.1.1/pysersic/pysersic.py +537 -0
- pysersic-0.1.1/pysersic/rendering.py +1183 -0
- pysersic-0.1.1/pysersic/results.py +544 -0
- pysersic-0.1.1/pysersic.egg-info/PKG-INFO +8 -0
- pysersic-0.1.1/pysersic.egg-info/SOURCES.txt +16 -0
- pysersic-0.1.1/pysersic.egg-info/dependency_links.txt +1 -0
- pysersic-0.1.1/pysersic.egg-info/requires.txt +12 -0
- pysersic-0.1.1/pysersic.egg-info/top_level.txt +1 -0
- pysersic-0.1.1/setup.cfg +7 -0
- pysersic-0.1.1/setup.py +24 -0
pysersic-0.1.1/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2023 pysersic
|
|
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.
|
pysersic-0.1.1/PKG-INFO
ADDED
|
@@ -0,0 +1,8 @@
|
|
|
1
|
+
Metadata-Version: 2.1
|
|
2
|
+
Name: pysersic
|
|
3
|
+
Version: 0.1.1
|
|
4
|
+
Summary: A Tool for fitting sersic profiles in python
|
|
5
|
+
Download-URL: https://github.com/pysersic/pysersic/archive/refs/tags/v0.1.1.tar.gz
|
|
6
|
+
Author: Imad Pasha & Tim Miller
|
|
7
|
+
Author-email: imad.pasha@yale.edu
|
|
8
|
+
License-File: LICENSE
|
pysersic-0.1.1/README.md
ADDED
|
@@ -0,0 +1,68 @@
|
|
|
1
|
+

|
|
2
|
+
|
|
3
|
+

|
|
4
|
+
[](https://pysersic.readthedocs.io/en/latest/?badge=latest)
|
|
5
|
+

|
|
6
|
+
|
|
7
|
+
pysersic is a code for fitting sersic profiles to galaxy images using Bayesian Inference. It is written in python using [jax](https://github.com/google/jax) with inference performed using [numpyro](https://github.com/pyro-ppl/numpyro)
|
|
8
|
+
|
|
9
|
+
## Installation
|
|
10
|
+
|
|
11
|
+
First you should install jax. The easiest way to do this is to use ``` pip install 'jax[cpu]' ```, however if you are on windows or planning on running it on a GPU, the installation can be a little more complicated, please see the guide [here](https://github.com/google/jax#installation). The other package dependencies for ```pysersic``` are:
|
|
12
|
+
- arviz
|
|
13
|
+
- asdf
|
|
14
|
+
- astropy
|
|
15
|
+
- corner
|
|
16
|
+
- jax
|
|
17
|
+
- matplotlib
|
|
18
|
+
- numpy
|
|
19
|
+
- numpyro
|
|
20
|
+
- pandas
|
|
21
|
+
- photutils
|
|
22
|
+
- scipy
|
|
23
|
+
- tqdm
|
|
24
|
+
|
|
25
|
+
Next you can install pysersic! We will be uploading it to pypi very soon, but for now you should install it locally by cloning the github repo.
|
|
26
|
+
|
|
27
|
+
```
|
|
28
|
+
git clone https://github.com/pysersic/pysersic.git
|
|
29
|
+
cd pysersic
|
|
30
|
+
pip install -e .
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
## Basic usage
|
|
34
|
+
|
|
35
|
+
All you need to run pysersic is a cutout of you galaxy of interest, and cutout of the error map, pixelized version of the PSF, and optionally a mask specifying bad pixels or nearby sources. The most basic setup is shown below for many many more details and in depth examples please see the [documentation](https://pysersic.readthedocs.io/en/latest/). First we set up a prior.
|
|
36
|
+
|
|
37
|
+
```
|
|
38
|
+
from pysersic.priors import SourceProperties
|
|
39
|
+
|
|
40
|
+
props = SourceProperties(img_cutout,mask=mask) # Optional mask
|
|
41
|
+
prior = props.generate_prior('sersic', # Other profiles inclues 'exp', 'dev' and 'pointsource'
|
|
42
|
+
sky_type='none') # Can also use 'flat' or 'tilted-plane' to simultaneously fit a background
|
|
43
|
+
```
|
|
44
|
+
Then we initialize the fitting object:
|
|
45
|
+
|
|
46
|
+
```
|
|
47
|
+
from pysersic import FitSingle
|
|
48
|
+
from pysersic.loss import gaussian_loss
|
|
49
|
+
|
|
50
|
+
fitter = FitSingle(data=img_cutout, # Cutout of galaxy
|
|
51
|
+
rms=sig_cutout, # Cutout of pixel errors
|
|
52
|
+
psf=psf, #Point spread function
|
|
53
|
+
prior=prior, # Prior that we initialized above
|
|
54
|
+
mask=mask, #Optional mask
|
|
55
|
+
loss_func=gaussian_loss) # Can specify loss function! See loss.py for many options or write your own!
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
Now we can fit!
|
|
59
|
+
|
|
60
|
+
```
|
|
61
|
+
map_dict = fitter.find_MAP() # Find the 'best-fit' parameters as the maximum-a-posteriori. Returns a dictionary containing the MAP parameters and model image
|
|
62
|
+
|
|
63
|
+
svi_res = fitter.estimate_posterior('svi-flow') #Train a normalizing flow to estimate the posterior, retruns a PysersicResults object containing the posterior and other helper methods
|
|
64
|
+
nuts_res = fitter.sample() #Sample the posterior using the No U-turn sampler (NUTS), retruns a PysersicResults object
|
|
65
|
+
```
|
|
66
|
+
|
|
67
|
+
## Citation
|
|
68
|
+
We will be submitting a paper soon describing pysersic for now please reference this webpage and cite 'Pasha & Miller in prep.'
|
|
@@ -0,0 +1,190 @@
|
|
|
1
|
+
from typing import Iterable, Union
|
|
2
|
+
|
|
3
|
+
import jax.numpy as jnp
|
|
4
|
+
from astropy.io import fits
|
|
5
|
+
|
|
6
|
+
from pysersic import FitMulti, FitSingle, priors
|
|
7
|
+
|
|
8
|
+
from .loss import gaussian_loss
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def search_start(list_of_strings: Iterable, to_match: str)-> list:
|
|
12
|
+
"""Search a list of strings and match those which start with to_match
|
|
13
|
+
|
|
14
|
+
Parameters
|
|
15
|
+
----------
|
|
16
|
+
list_of_strings : Iterable
|
|
17
|
+
to_match : str
|
|
18
|
+
pattern to match
|
|
19
|
+
|
|
20
|
+
Returns
|
|
21
|
+
-------
|
|
22
|
+
list
|
|
23
|
+
List of results contianing the index and full string
|
|
24
|
+
"""
|
|
25
|
+
result = []
|
|
26
|
+
for j,string in enumerate(list_of_strings):
|
|
27
|
+
if string.split()[0] == to_match:
|
|
28
|
+
result.append([j,string])
|
|
29
|
+
return result
|
|
30
|
+
|
|
31
|
+
def load_fits_from_galfit_line(line:str) -> jnp.array:
|
|
32
|
+
"""Load fits file from galfit config line
|
|
33
|
+
|
|
34
|
+
Parameters
|
|
35
|
+
----------
|
|
36
|
+
line : str
|
|
37
|
+
|
|
38
|
+
Returns
|
|
39
|
+
-------
|
|
40
|
+
jnp.array
|
|
41
|
+
Image in first HDU
|
|
42
|
+
"""
|
|
43
|
+
fname = line.split()[1]
|
|
44
|
+
return jnp.array( fits.getdata(fname) )
|
|
45
|
+
|
|
46
|
+
def match_line_and_load_galfit(list_of_strings: Iterable, to_match: str, assert_message: str = "")-> jnp.array:
|
|
47
|
+
"""Combine the above two functions to search and load fits files from a galfit config
|
|
48
|
+
|
|
49
|
+
Parameters
|
|
50
|
+
----------
|
|
51
|
+
list_of_strings : Iterable
|
|
52
|
+
to_match : str
|
|
53
|
+
assert_message : str, optional
|
|
54
|
+
Message to raise if line file is not present
|
|
55
|
+
|
|
56
|
+
Returns
|
|
57
|
+
-------
|
|
58
|
+
jnp.array
|
|
59
|
+
_description_
|
|
60
|
+
|
|
61
|
+
Raises
|
|
62
|
+
------
|
|
63
|
+
UserWarning
|
|
64
|
+
_description_
|
|
65
|
+
"""
|
|
66
|
+
matches = search_start(list_of_strings, to_match)
|
|
67
|
+
line = matches[0][1]
|
|
68
|
+
if line.split()[1] == 'none':
|
|
69
|
+
raise UserWarning(assert_message)
|
|
70
|
+
return load_fits_from_galfit_line(line)
|
|
71
|
+
|
|
72
|
+
def generate_fitter_from_galfit_config(config_loc: str, loss: callable = gaussian_loss ) -> Union[FitSingle,FitMulti]:
|
|
73
|
+
"""Function to generate a Fitter instance from a galfit config file. This function is only meant to be a starting place to easily try PySersic on the many galfit config files we know you all have lying around. As such this does not follow all of the rules and constraints the galfit takes into account. For more fine grained control we recommend initializing a pysersic Fitter instance yourself, please see the `examples/` folder or the documentation for some help with this.
|
|
74
|
+
|
|
75
|
+
Parameters
|
|
76
|
+
----------
|
|
77
|
+
config_loc : str
|
|
78
|
+
Config file locations
|
|
79
|
+
loss : callable, optional
|
|
80
|
+
Loss function to use, see loss.py for more details, by default gaussian_loss
|
|
81
|
+
|
|
82
|
+
Returns
|
|
83
|
+
-------
|
|
84
|
+
Union[FitSingle,FitMulti]
|
|
85
|
+
Fitter based on config file. If there is only one source then returns a FitSingle object, otherwise returns a fit_multi
|
|
86
|
+
"""
|
|
87
|
+
|
|
88
|
+
with open(config_loc,'r') as f:
|
|
89
|
+
lines = []
|
|
90
|
+
for line in f.readlines():
|
|
91
|
+
if line[0] == '#' or line[:1] == '\n' or line[0] == '=':
|
|
92
|
+
continue
|
|
93
|
+
else:
|
|
94
|
+
lines.append(line)
|
|
95
|
+
|
|
96
|
+
im = match_line_and_load_galfit(lines, 'A)', assert_message='Missing science image, labelled D) in config file')
|
|
97
|
+
rms = match_line_and_load_galfit(lines, 'C)', assert_message='Missing rms image, labelled C) in config file')
|
|
98
|
+
psf = match_line_and_load_galfit(lines,'D)', assert_message='Missing psf image, labelled D) in config file')
|
|
99
|
+
|
|
100
|
+
try:
|
|
101
|
+
mask = match_line_and_load_galfit(lines,'F)')
|
|
102
|
+
except:
|
|
103
|
+
mask = jnp.zeros_like(im)
|
|
104
|
+
zpt = float ( search_start(lines, 'J)')[0][1].split()[1] )
|
|
105
|
+
source_list = search_start(lines, '0)')
|
|
106
|
+
|
|
107
|
+
sky_type = 'none'
|
|
108
|
+
source_dict = {}
|
|
109
|
+
source_dict['type'] = []
|
|
110
|
+
source_dict['flux'] = []
|
|
111
|
+
source_dict['r'] = []
|
|
112
|
+
source_dict['x'] = []
|
|
113
|
+
source_dict['y'] = []
|
|
114
|
+
|
|
115
|
+
for i,source in enumerate(source_list):
|
|
116
|
+
ind = source[0]
|
|
117
|
+
|
|
118
|
+
if i == len(source_list) -1:
|
|
119
|
+
next_ind = None
|
|
120
|
+
else:
|
|
121
|
+
next_ind = source_list[i+1][0]
|
|
122
|
+
source_string_list = lines[ind:next_ind]
|
|
123
|
+
|
|
124
|
+
source_type = source[1].split()[1]
|
|
125
|
+
if source_type == 'sersic' or 'dev' in source_type or 'exp' in source_type or source_type == 'psf':
|
|
126
|
+
if 'dev' in source_type:
|
|
127
|
+
source_type = 'dev'
|
|
128
|
+
elif 'exp' in source_type:
|
|
129
|
+
source_type = 'exp'
|
|
130
|
+
elif source_type == 'psf':
|
|
131
|
+
source_type = 'pointsource'
|
|
132
|
+
|
|
133
|
+
mag = float( search_start(source_string_list, '3)')[0][1].split()[1] )
|
|
134
|
+
flux = 10.**( (mag - zpt)/-2.5)
|
|
135
|
+
|
|
136
|
+
x,y = search_start(source_string_list, '1)')[0][1].split()[1:3]
|
|
137
|
+
x = float(x)
|
|
138
|
+
y = float(y)
|
|
139
|
+
|
|
140
|
+
source_dict['type'].append(source_type)
|
|
141
|
+
source_dict['x'].append(x)
|
|
142
|
+
source_dict['y'].append(y)
|
|
143
|
+
source_dict['flux'].append(flux)
|
|
144
|
+
|
|
145
|
+
if source_type != 'pointsource':
|
|
146
|
+
r = float( search_start(source_string_list, '4)')[0][1].split()[1] )
|
|
147
|
+
else:
|
|
148
|
+
r = -99
|
|
149
|
+
source_dict['r'].append(r)
|
|
150
|
+
|
|
151
|
+
elif source_type == 'sky':
|
|
152
|
+
print('sky')
|
|
153
|
+
fit_back = search_start(source_string_list, '1)')[0][1].split()[2]
|
|
154
|
+
back_init = float(search_start(source_string_list, '1)')[0][1].split()[1])
|
|
155
|
+
fit_x_sl = search_start(source_string_list,'2)')[0][1].split()[2]
|
|
156
|
+
fit_y_sl = search_start(source_string_list,'2)')[0][1].split()[2]
|
|
157
|
+
if fit_y_sl == '1' or fit_x_sl == '1':
|
|
158
|
+
sky_type = 'tilted-plane'
|
|
159
|
+
elif fit_back == '1':
|
|
160
|
+
sky_type = 'flat'
|
|
161
|
+
else:
|
|
162
|
+
sky_type = 'none'
|
|
163
|
+
else:
|
|
164
|
+
print (f"Skipping source {i}, profile type of {source_type} currently not supported")
|
|
165
|
+
|
|
166
|
+
N_sources = len(source_dict['type'])
|
|
167
|
+
types = source_dict['type']
|
|
168
|
+
print(f"Found {N_sources} source(s) to fit of type(s) {types}")
|
|
169
|
+
|
|
170
|
+
if len(source_dict['type']) == 1:
|
|
171
|
+
prof_type = source_dict['type'][0]
|
|
172
|
+
sp = priors.SourceProperties(im, mask = mask)
|
|
173
|
+
sp.set_sky_guess(sky_guess = 0, sky_guess_err = 1)
|
|
174
|
+
sp.set_flux_guess(flux_guess= source_dict['flux'][0])
|
|
175
|
+
sp.set_position_guess( [source_dict['x'][0],source_dict['y'][0] ])
|
|
176
|
+
|
|
177
|
+
if prof_type != 'psf':
|
|
178
|
+
sp.set_r_eff_guess(r_eff_guess=source_dict['r'][0])
|
|
179
|
+
prior = sp.generate_prior(prof_type, sky_type = sky_type)
|
|
180
|
+
ps_fitter = FitSingle(data = im, rms = rms, psf = psf, prior = prior, loss_func = loss )
|
|
181
|
+
|
|
182
|
+
elif len(source_dict['type']) > 1:
|
|
183
|
+
multi_prior = priors.PySersicMultiPrior(catalog=source_dict, sky_type = sky_type,sky_guess=back_init, sky_guess_err = back_init + 1e-5)
|
|
184
|
+
ps_fitter = FitMulti(data = im, rms = rms, psf = psf, prior = multi_prior, loss_func = loss)
|
|
185
|
+
|
|
186
|
+
else:
|
|
187
|
+
print ('Did not find any viable sources in galfit config')
|
|
188
|
+
ps_fitter = 0
|
|
189
|
+
|
|
190
|
+
return ps_fitter
|
|
@@ -0,0 +1,328 @@
|
|
|
1
|
+
import jax.numpy as jnp
|
|
2
|
+
from numpyro import distributions as dist, sample, handlers, factor,deterministic
|
|
3
|
+
from typing import Optional
|
|
4
|
+
|
|
5
|
+
def gaussian_loss(mod: jnp.array,
|
|
6
|
+
data: jnp.array,
|
|
7
|
+
rms: jnp.array,
|
|
8
|
+
mask: jnp.array)-> float:
|
|
9
|
+
"""Basic Gaussian loss function using given uncertainties
|
|
10
|
+
|
|
11
|
+
Parameters
|
|
12
|
+
----------
|
|
13
|
+
mod : jnp.array
|
|
14
|
+
Model image
|
|
15
|
+
data : jnp.array
|
|
16
|
+
data to be fit
|
|
17
|
+
rms : jnp.array
|
|
18
|
+
per pixel 1-sigma uncertainties
|
|
19
|
+
|
|
20
|
+
Returns
|
|
21
|
+
-------
|
|
22
|
+
float
|
|
23
|
+
Sampled loss function
|
|
24
|
+
"""
|
|
25
|
+
|
|
26
|
+
with handlers.mask(mask = mask):
|
|
27
|
+
loss = sample("Loss", dist.Normal(mod, rms), obs=data)
|
|
28
|
+
return loss
|
|
29
|
+
|
|
30
|
+
def cash_loss(mod: jnp.array,
|
|
31
|
+
data: jnp.array,
|
|
32
|
+
rms: jnp.array,
|
|
33
|
+
mask: jnp.array)-> float:
|
|
34
|
+
"""
|
|
35
|
+
Cash statistic based on Poisson statistics derived in Cash (1979) (DOI 10.1086/156922) and advocated for in Erwin (2015) (https://arxiv.org/abs/1408.1097) for use in Sersic fitting. Since the is based on Poisson statistics, scaling of the image will produce different confidence intervals. Additionally, since a logarithm is taken of the model image, negative values associated with different sky models will cause issues.
|
|
36
|
+
|
|
37
|
+
Parameters
|
|
38
|
+
----------
|
|
39
|
+
mod : jnp.array
|
|
40
|
+
Model image
|
|
41
|
+
data : jnp.array
|
|
42
|
+
data to be fit
|
|
43
|
+
rms : jnp.array
|
|
44
|
+
per pixel 1-sigma uncertainties
|
|
45
|
+
|
|
46
|
+
Returns
|
|
47
|
+
-------
|
|
48
|
+
float
|
|
49
|
+
Sampled loss function
|
|
50
|
+
"""
|
|
51
|
+
|
|
52
|
+
with handlers.mask(mask = mask):
|
|
53
|
+
loss = factor('cash_loss', -1*(mod - data*jnp.log(mod)))
|
|
54
|
+
return loss
|
|
55
|
+
|
|
56
|
+
def gaussian_loss_w_frac(mod: jnp.array,
|
|
57
|
+
data: jnp.array,
|
|
58
|
+
rms: jnp.array,
|
|
59
|
+
mask: jnp.array)-> float:
|
|
60
|
+
"""Gaussian loss with and additional fractional increase to all uncertainties such that,
|
|
61
|
+
|
|
62
|
+
$$ \sigma_{new,i} = (1 + f) * \sigma_{old,i} $$
|
|
63
|
+
|
|
64
|
+
f is a free parameter varied between -0.5 and 2
|
|
65
|
+
|
|
66
|
+
Parameters
|
|
67
|
+
----------
|
|
68
|
+
mod : jnp.array
|
|
69
|
+
Model image
|
|
70
|
+
data : jnp.array
|
|
71
|
+
data to be fit
|
|
72
|
+
rms : jnp.array
|
|
73
|
+
per pixel 1-sigma uncertainties
|
|
74
|
+
|
|
75
|
+
Returns
|
|
76
|
+
-------
|
|
77
|
+
float
|
|
78
|
+
Sampled loss function
|
|
79
|
+
"""
|
|
80
|
+
|
|
81
|
+
scatter_frac = sample('frac_scatter_increase', dist.TruncatedNormal(low = -0.5, high = 2) )
|
|
82
|
+
with handlers.mask(mask = mask):
|
|
83
|
+
loss = sample("Loss", dist.Normal(mod, (1+ scatter_frac)*rms), obs=data)
|
|
84
|
+
return loss
|
|
85
|
+
|
|
86
|
+
def gaussian_loss_w_sys(mod: jnp.array,
|
|
87
|
+
data: jnp.array,
|
|
88
|
+
rms: jnp.array,
|
|
89
|
+
mask: jnp.array)-> float:
|
|
90
|
+
"""Gaussian loss with and additional systematic increase such_that
|
|
91
|
+
|
|
92
|
+
$$ \sigma_{new,i}^2 = \sigma_{old,i}^2 + \sigma_{sys}^2 $$
|
|
93
|
+
|
|
94
|
+
\sigma_{sys} is a free parameter with a Normal prior, with loc = 0 and scale = mean(rms) truncated to ensure > 0
|
|
95
|
+
|
|
96
|
+
Parameters
|
|
97
|
+
----------
|
|
98
|
+
mod : jnp.array
|
|
99
|
+
Model image
|
|
100
|
+
data : jnp.array
|
|
101
|
+
data to be fit
|
|
102
|
+
rms : jnp.array
|
|
103
|
+
per pixel 1-sigma uncertainties
|
|
104
|
+
|
|
105
|
+
Returns
|
|
106
|
+
-------
|
|
107
|
+
float
|
|
108
|
+
Sampled loss function
|
|
109
|
+
"""
|
|
110
|
+
|
|
111
|
+
sys_scatter_base = sample('sys_scatter_base', dist.TruncatedNormal(low = 0, scale = 1 ) )
|
|
112
|
+
sys_scatter = deterministic('sys_scatter', sys_scatter_base*jnp.mean(rms))
|
|
113
|
+
with handlers.mask(mask = mask):
|
|
114
|
+
loss = sample("Loss", dist.Normal(mod, jnp.sqrt(rms**2 + sys_scatter**2)), obs=data)
|
|
115
|
+
return loss
|
|
116
|
+
|
|
117
|
+
def student_t_loss(mod: jnp.array,
|
|
118
|
+
data: jnp.array,
|
|
119
|
+
rms: jnp.array,
|
|
120
|
+
mask: jnp.array,
|
|
121
|
+
nu: Optional[int] = 5)-> float:
|
|
122
|
+
"""
|
|
123
|
+
Student T loss, with a df = 5 by default. This has fatter tails than Gaussian loss (or chi squared) so is more resilient to outliers
|
|
124
|
+
|
|
125
|
+
Parameters
|
|
126
|
+
----------
|
|
127
|
+
mod : jnp.array
|
|
128
|
+
Model image
|
|
129
|
+
data : jnp.array
|
|
130
|
+
data to be fit
|
|
131
|
+
rms : jnp.array
|
|
132
|
+
per pixel 1-sigma uncertainties
|
|
133
|
+
|
|
134
|
+
Returns
|
|
135
|
+
-------
|
|
136
|
+
float
|
|
137
|
+
Sampled loss function
|
|
138
|
+
"""
|
|
139
|
+
nu = 5.
|
|
140
|
+
with handlers.mask(mask = mask):
|
|
141
|
+
loss = sample("Loss", dist.StudentT(nu,mod,jnp.sqrt((nu-2.)/2.)*rms), obs=data)
|
|
142
|
+
return loss
|
|
143
|
+
|
|
144
|
+
def student_t_loss_free_sys(mod: jnp.array,
|
|
145
|
+
data: jnp.array,
|
|
146
|
+
rms: jnp.array,
|
|
147
|
+
mask: jnp.array,
|
|
148
|
+
nu: Optional[int] = 5)-> float:
|
|
149
|
+
"""
|
|
150
|
+
Student T loss, which has fatter tails than Gaussian loss (or chi squared) so is so is more resilient to outliers. In addition, add additional systematic increase such that
|
|
151
|
+
|
|
152
|
+
$$ \sigma_{new,i}^2 = \sigma_{old,i}^2 + \sigma_{sys}^2 $$
|
|
153
|
+
|
|
154
|
+
\sigma_{sys} is a free parameter with a Normal prior, with loc = 0 and scale = mean(rms) truncated to ensure > 0
|
|
155
|
+
|
|
156
|
+
|
|
157
|
+
Parameters
|
|
158
|
+
----------
|
|
159
|
+
mod : jnp.array
|
|
160
|
+
Model image
|
|
161
|
+
data : jnp.array
|
|
162
|
+
data to be fit
|
|
163
|
+
rms : jnp.array
|
|
164
|
+
per pixel 1-sigma uncertainties
|
|
165
|
+
|
|
166
|
+
Returns
|
|
167
|
+
-------
|
|
168
|
+
float
|
|
169
|
+
Sampled loss function
|
|
170
|
+
"""
|
|
171
|
+
sys_scatter_base = sample('sys_scatter_base', dist.TruncatedNormal(low = 0, scale = 1 ) )
|
|
172
|
+
sys_scatter = deterministic('sys_scatter', sys_scatter_base*jnp.mean(rms))
|
|
173
|
+
rms_new = jnp.sqrt((nu-2.)/2.)*jnp.sqrt(rms**2 + sys_scatter**2)
|
|
174
|
+
|
|
175
|
+
with handlers.mask(mask = mask):
|
|
176
|
+
loss = sample("Loss", dist.StudentT(nu, mod, rms_new), obs=data)
|
|
177
|
+
return loss
|
|
178
|
+
|
|
179
|
+
def pseudo_huber_loss(mod: jnp.array,
|
|
180
|
+
data: jnp.array,
|
|
181
|
+
rms: jnp.array,
|
|
182
|
+
mask: jnp.array,
|
|
183
|
+
delta: Optional[int] = 3
|
|
184
|
+
)-> float:
|
|
185
|
+
"""
|
|
186
|
+
Pseudo huber loss function of the form:
|
|
187
|
+
|
|
188
|
+
$$ L = \delta^2 * ( \sqrt{1 + (a/\delta)^2} - 1) $$
|
|
189
|
+
|
|
190
|
+
where a is the residuals scaled by the rms and delta can be chosen. This loss function is more robust to outliers than the Gaussian loss function as it is meant to transition for L2 to L1 loss at residuals greater than delta. The delta parameter is 3 by default but can be varied.
|
|
191
|
+
|
|
192
|
+
Parameters
|
|
193
|
+
----------
|
|
194
|
+
mod : jnp.array
|
|
195
|
+
Model image
|
|
196
|
+
data : jnp.array
|
|
197
|
+
data to be fit
|
|
198
|
+
rms : jnp.array
|
|
199
|
+
per pixel 1-sigma uncertainties
|
|
200
|
+
|
|
201
|
+
Returns
|
|
202
|
+
-------
|
|
203
|
+
float
|
|
204
|
+
Sampled loss function
|
|
205
|
+
"""
|
|
206
|
+
|
|
207
|
+
with handlers.mask(mask = mask):
|
|
208
|
+
res = (data-mod)/rms
|
|
209
|
+
loss = factor('pseudo_huber_loss', -1.*(jnp.sqrt(1 + ( res/delta )**2) - 1) )
|
|
210
|
+
return loss
|
|
211
|
+
|
|
212
|
+
def gaussian_mixture(mod: jnp.array,
|
|
213
|
+
data: jnp.array,
|
|
214
|
+
rms: jnp.array,
|
|
215
|
+
mask: jnp.array,
|
|
216
|
+
c: Optional[float] = 5.
|
|
217
|
+
)-> float:
|
|
218
|
+
"""
|
|
219
|
+
Gaussian mixture loss function, with one representing a "contaminating" outlier distribution with standard deviation equal to c*rms where c is 5 by default. The "contaminating fraction" or fraction of outliers is a free parameter with a Uniform prior between 0 and 0.25.
|
|
220
|
+
|
|
221
|
+
Parameters
|
|
222
|
+
----------
|
|
223
|
+
Parameters
|
|
224
|
+
----------
|
|
225
|
+
mod : jnp.array
|
|
226
|
+
Model image
|
|
227
|
+
data : jnp.array
|
|
228
|
+
data to be fit
|
|
229
|
+
rms : jnp.array
|
|
230
|
+
per pixel 1-sigma uncertainties
|
|
231
|
+
c : float, optional
|
|
232
|
+
factor to increase rms for outlier distribution, by default 5
|
|
233
|
+
|
|
234
|
+
Returns
|
|
235
|
+
-------
|
|
236
|
+
float
|
|
237
|
+
_description_
|
|
238
|
+
"""
|
|
239
|
+
contam_frac_base = sample('outlier_frac_base', dist.TruncatedNormal(low = 0, scale = 1., high = 5.) )
|
|
240
|
+
contam_frac = deterministic('outlier_frac', contam_frac_base*0.05 )
|
|
241
|
+
|
|
242
|
+
mixture_dists = dist.Categorical(probs = jnp.array([1-contam_frac, contam_frac]))
|
|
243
|
+
component_dists = dist.Normal(jnp.stack([mod,mod],axis = -1), jnp.stack([rms,c*rms],axis = -1) )
|
|
244
|
+
|
|
245
|
+
with handlers.mask(mask = mask):
|
|
246
|
+
loss = sample("Loss", dist.MixtureSameFamily(mixture_dists, component_dists), obs=data)
|
|
247
|
+
return loss
|
|
248
|
+
|
|
249
|
+
def gaussian_mixture_w_sys(mod: jnp.array,
|
|
250
|
+
data: jnp.array,
|
|
251
|
+
rms: jnp.array,
|
|
252
|
+
mask: jnp.array,
|
|
253
|
+
c: Optional[float] = 5.
|
|
254
|
+
)-> float:
|
|
255
|
+
"""
|
|
256
|
+
Gaussian mixture loss function, with one representing a "contaminating" outlier distribution with standard deviation equal to c*rms where c is 5 by default. The "outlier fraction" or fraction of outliers is a free parameter with a Uniform prior between 0 and 0.25.
|
|
257
|
+
|
|
258
|
+
Parameters
|
|
259
|
+
----------
|
|
260
|
+
Parameters
|
|
261
|
+
----------
|
|
262
|
+
mod : jnp.array
|
|
263
|
+
Model image
|
|
264
|
+
data : jnp.array
|
|
265
|
+
data to be fit
|
|
266
|
+
rms : jnp.array
|
|
267
|
+
per pixel 1-sigma uncertainties
|
|
268
|
+
c : float, optional
|
|
269
|
+
factor to increase rms for outlier distribution, by default 5
|
|
270
|
+
|
|
271
|
+
Returns
|
|
272
|
+
-------
|
|
273
|
+
float
|
|
274
|
+
_description_
|
|
275
|
+
"""
|
|
276
|
+
contam_frac_base = sample('outlier_frac_base', dist.TruncatedNormal(low = 0, scale = 1., high = 5.) )
|
|
277
|
+
contam_frac = deterministic('outlier_frac', contam_frac_base*0.05 )
|
|
278
|
+
|
|
279
|
+
sys_scatter_base = sample('sys_scatter_base', dist.TruncatedNormal(low = 0, scale = 1 ) )
|
|
280
|
+
sys_scatter = deterministic('sys_scatter', sys_scatter_base*jnp.mean(rms))
|
|
281
|
+
|
|
282
|
+
rms_new = jnp.sqrt(rms**2 + sys_scatter**2)
|
|
283
|
+
mixture_dists = dist.Categorical(probs = jnp.array([1-contam_frac, contam_frac]))
|
|
284
|
+
component_dists = dist.Normal(jnp.stack([mod,mod],axis = -1), jnp.stack([rms_new,c*rms_new],axis = -1) )
|
|
285
|
+
|
|
286
|
+
with handlers.mask(mask = mask):
|
|
287
|
+
loss = sample("Loss", dist.MixtureSameFamily(mixture_dists, component_dists), obs=data)
|
|
288
|
+
return loss
|
|
289
|
+
|
|
290
|
+
def gaussian_mixture_w_frac(mod: jnp.array,
|
|
291
|
+
data: jnp.array,
|
|
292
|
+
rms: jnp.array,
|
|
293
|
+
mask: jnp.array,
|
|
294
|
+
c: Optional[float] = 5.
|
|
295
|
+
)-> float:
|
|
296
|
+
"""
|
|
297
|
+
Gaussian mixture loss function, with one representing a "contaminating" outlier distribution with standard deviation equal to c*rms where c is 5 by default. The "outlier fraction" or fraction of outliers is a free parameter with a Uniform prior between 0 and 0.25.
|
|
298
|
+
|
|
299
|
+
Parameters
|
|
300
|
+
----------
|
|
301
|
+
Parameters
|
|
302
|
+
----------
|
|
303
|
+
mod : jnp.array
|
|
304
|
+
Model image
|
|
305
|
+
data : jnp.array
|
|
306
|
+
data to be fit
|
|
307
|
+
rms : jnp.array
|
|
308
|
+
per pixel 1-sigma uncertainties
|
|
309
|
+
c : float, optional
|
|
310
|
+
factor to increase rms for outlier distribution, by default 5
|
|
311
|
+
|
|
312
|
+
Returns
|
|
313
|
+
-------
|
|
314
|
+
float
|
|
315
|
+
_description_
|
|
316
|
+
"""
|
|
317
|
+
contam_frac_base = sample('outlier_frac_base', dist.TruncatedNormal(low = 0, scale = 1., high = 5.) )
|
|
318
|
+
contam_frac = deterministic('outlier_frac', contam_frac_base*0.05 )
|
|
319
|
+
|
|
320
|
+
sig_frac = sample('sig_frac', dist.TruncatedNormal(low = -2./3., high = 2., loc = 0, scale = 0.5 ) )
|
|
321
|
+
|
|
322
|
+
rms_new = (1+sig_frac)*rms
|
|
323
|
+
mixture_dists = dist.Categorical(probs = jnp.array([1-contam_frac, contam_frac]))
|
|
324
|
+
component_dists = dist.Normal(jnp.stack([mod,mod],axis = -1), jnp.stack([rms_new,c*rms],axis = -1) )
|
|
325
|
+
|
|
326
|
+
with handlers.mask(mask = mask):
|
|
327
|
+
loss = sample("Loss", dist.MixtureSameFamily(mixture_dists, component_dists), obs=data)
|
|
328
|
+
return loss
|