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 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.
@@ -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
@@ -0,0 +1,68 @@
1
+ ![pysersic logo featuring a spiral s-galaxy as the S](misc/pysersic.png)
2
+
3
+ ![License](https://img.shields.io/badge/license-MIT-blue)
4
+ [![Documentation Status](https://readthedocs.org/projects/pysersic/badge/?version=latest)](https://pysersic.readthedocs.io/en/latest/?badge=latest)
5
+ ![tests](https://github.com/pysersic/pysersic/actions/workflows/pytest.yml/badge.svg)
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,4 @@
1
+ from .priors import (PySersicMultiPrior, PySersicSourcePrior, SourceProperties,autoprior)
2
+ from .pysersic import FitMulti, FitSingle
3
+ from .rendering import FourierRenderer, HybridRenderer, PixelRenderer
4
+ from .results import PySersicResults, parse_multi_results
@@ -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