periapsis 0.1.0__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.
- periapsis-0.1.0/LICENSE +21 -0
- periapsis-0.1.0/PKG-INFO +61 -0
- periapsis-0.1.0/README.md +40 -0
- periapsis-0.1.0/pyproject.toml +30 -0
- periapsis-0.1.0/setup.cfg +4 -0
- periapsis-0.1.0/src/periapsis/__init__.py +5 -0
- periapsis-0.1.0/src/periapsis/data/__init__.py +6 -0
- periapsis-0.1.0/src/periapsis/data/common.py +71 -0
- periapsis-0.1.0/src/periapsis/data/data.py +40 -0
- periapsis-0.1.0/src/periapsis/data/gaia.py +10 -0
- periapsis-0.1.0/src/periapsis/data/joint_data.py +14 -0
- periapsis-0.1.0/src/periapsis/fitting/__init__.py +4 -0
- periapsis-0.1.0/src/periapsis/fitting/fitter.py +88 -0
- periapsis-0.1.0/src/periapsis/fitting/mcmcCampbell.py +121 -0
- periapsis-0.1.0/src/periapsis/fitting/mcmcThieleInnes.py +107 -0
- periapsis-0.1.0/src/periapsis/fitting/mcmclinear.py +249 -0
- periapsis-0.1.0/src/periapsis/fitting/results.py +65 -0
- periapsis-0.1.0/src/periapsis/fitting/ultranestCampbell.py +90 -0
- periapsis-0.1.0/src/periapsis/fitting/ultranestThieleInnes.py +71 -0
- periapsis-0.1.0/src/periapsis/fitting/ultranestlinear.py +227 -0
- periapsis-0.1.0/src/periapsis/initial/__init__.py +0 -0
- periapsis-0.1.0/src/periapsis/initial/initial.py +144 -0
- periapsis-0.1.0/src/periapsis/model/__init__.py +3 -0
- periapsis-0.1.0/src/periapsis/model/campbell.py +55 -0
- periapsis-0.1.0/src/periapsis/model/orbit.py +89 -0
- periapsis-0.1.0/src/periapsis/model/thieleinnes.py +48 -0
- periapsis-0.1.0/src/periapsis/plotting/__init__.py +0 -0
- periapsis-0.1.0/src/periapsis/plotting/plots.py +506 -0
- periapsis-0.1.0/src/periapsis/prior/__init__.py +5 -0
- periapsis-0.1.0/src/periapsis/prior/normal_prior.py +26 -0
- periapsis-0.1.0/src/periapsis/prior/prior.py +10 -0
- periapsis-0.1.0/src/periapsis/prior/uniform_prior.py +22 -0
- periapsis-0.1.0/src/periapsis/stats/__init__.py +0 -0
- periapsis-0.1.0/src/periapsis/stats/stat_funcs.py +163 -0
- periapsis-0.1.0/src/periapsis/utils/helpers.py +66 -0
- periapsis-0.1.0/src/periapsis/utils/solvers.py +163 -0
- periapsis-0.1.0/src/periapsis.egg-info/PKG-INFO +61 -0
- periapsis-0.1.0/src/periapsis.egg-info/SOURCES.txt +45 -0
- periapsis-0.1.0/src/periapsis.egg-info/dependency_links.txt +1 -0
- periapsis-0.1.0/src/periapsis.egg-info/requires.txt +7 -0
- periapsis-0.1.0/src/periapsis.egg-info/top_level.txt +1 -0
- periapsis-0.1.0/tests/test_data.py +85 -0
- periapsis-0.1.0/tests/test_fitting.py +41 -0
- periapsis-0.1.0/tests/test_model.py +24 -0
- periapsis-0.1.0/tests/test_prior.py +29 -0
- periapsis-0.1.0/tests/test_results.py +18 -0
- periapsis-0.1.0/tests/test_stats.py +96 -0
periapsis-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Evan Kerr and Matthew Whitaker
|
|
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.
|
periapsis-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,61 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: periapsis
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: An orbital fitting package
|
|
5
|
+
Author-email: Evan Kerr <evan.kerr@utah.edu>, Matthew Whitaker <matthew.whitaker@utah.edu>
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Keywords: binaries,astronomy,regression,model_fitting
|
|
8
|
+
Classifier: Programming Language :: Python :: 3
|
|
9
|
+
Classifier: Operating System :: OS Independent
|
|
10
|
+
Requires-Python: >=3.9
|
|
11
|
+
Description-Content-Type: text/markdown
|
|
12
|
+
License-File: LICENSE
|
|
13
|
+
Requires-Dist: numpy
|
|
14
|
+
Requires-Dist: scipy
|
|
15
|
+
Requires-Dist: astropy
|
|
16
|
+
Requires-Dist: matplotlib
|
|
17
|
+
Requires-Dist: corner
|
|
18
|
+
Requires-Dist: emcee
|
|
19
|
+
Requires-Dist: ultranest
|
|
20
|
+
Dynamic: license-file
|
|
21
|
+
|
|
22
|
+
# periapsis
|
|
23
|
+
|
|
24
|
+
Package for efficiently modeling and fitting orbits with various parameterizations and priors with support for data from a variety of sources.
|
|
25
|
+
|
|
26
|
+
## Installing
|
|
27
|
+
|
|
28
|
+
### pip
|
|
29
|
+
|
|
30
|
+
```bash
|
|
31
|
+
pip install periapsis
|
|
32
|
+
```
|
|
33
|
+
|
|
34
|
+
## Usage Example
|
|
35
|
+
|
|
36
|
+
```python
|
|
37
|
+
from periapsis.data import AstrometryData, RadialVelocityData, JointData
|
|
38
|
+
from periapsis.fitting import MCMCFitter
|
|
39
|
+
from periapsis.prior import UniformPrior
|
|
40
|
+
import numpy as np
|
|
41
|
+
|
|
42
|
+
fit_data = JointData([
|
|
43
|
+
AstrometryData(t_astro, x, y),
|
|
44
|
+
RadialVelocityData(t_rv, rv),
|
|
45
|
+
])
|
|
46
|
+
|
|
47
|
+
fitter = MCMCFitter(
|
|
48
|
+
nwalkers=32,
|
|
49
|
+
niter=10000,
|
|
50
|
+
P=UniformPrior(10, 20000), # orbital period, days
|
|
51
|
+
t0=UniformPrior(1990, 2050), # time of periapsis passage
|
|
52
|
+
a=UniformPrior(0.01, 1000), # semi-major axis, AU
|
|
53
|
+
e=UniformPrior(0, 1), # eccentricity
|
|
54
|
+
cosi=UniformPrior(-1, 1), # cos(inclination)
|
|
55
|
+
omega=UniformPrior(0, 2*np.pi), # argument of periapsis
|
|
56
|
+
bigomega=UniformPrior(0, 2*np.pi), # longitude of ascending node
|
|
57
|
+
)
|
|
58
|
+
|
|
59
|
+
result = fitter.fit(fit_data)
|
|
60
|
+
```
|
|
61
|
+
|
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
# periapsis
|
|
2
|
+
|
|
3
|
+
Package for efficiently modeling and fitting orbits with various parameterizations and priors with support for data from a variety of sources.
|
|
4
|
+
|
|
5
|
+
## Installing
|
|
6
|
+
|
|
7
|
+
### pip
|
|
8
|
+
|
|
9
|
+
```bash
|
|
10
|
+
pip install periapsis
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
## Usage Example
|
|
14
|
+
|
|
15
|
+
```python
|
|
16
|
+
from periapsis.data import AstrometryData, RadialVelocityData, JointData
|
|
17
|
+
from periapsis.fitting import MCMCFitter
|
|
18
|
+
from periapsis.prior import UniformPrior
|
|
19
|
+
import numpy as np
|
|
20
|
+
|
|
21
|
+
fit_data = JointData([
|
|
22
|
+
AstrometryData(t_astro, x, y),
|
|
23
|
+
RadialVelocityData(t_rv, rv),
|
|
24
|
+
])
|
|
25
|
+
|
|
26
|
+
fitter = MCMCFitter(
|
|
27
|
+
nwalkers=32,
|
|
28
|
+
niter=10000,
|
|
29
|
+
P=UniformPrior(10, 20000), # orbital period, days
|
|
30
|
+
t0=UniformPrior(1990, 2050), # time of periapsis passage
|
|
31
|
+
a=UniformPrior(0.01, 1000), # semi-major axis, AU
|
|
32
|
+
e=UniformPrior(0, 1), # eccentricity
|
|
33
|
+
cosi=UniformPrior(-1, 1), # cos(inclination)
|
|
34
|
+
omega=UniformPrior(0, 2*np.pi), # argument of periapsis
|
|
35
|
+
bigomega=UniformPrior(0, 2*np.pi), # longitude of ascending node
|
|
36
|
+
)
|
|
37
|
+
|
|
38
|
+
result = fitter.fit(fit_data)
|
|
39
|
+
```
|
|
40
|
+
|
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools >= 77.0.3"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "periapsis"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
authors = [
|
|
9
|
+
{name = "Evan Kerr", email ="evan.kerr@utah.edu"},
|
|
10
|
+
{name = "Matthew Whitaker", email ="matthew.whitaker@utah.edu"},
|
|
11
|
+
]
|
|
12
|
+
description = "An orbital fitting package"
|
|
13
|
+
readme = "README.md"
|
|
14
|
+
requires-python = ">=3.9"
|
|
15
|
+
classifiers = [
|
|
16
|
+
"Programming Language :: Python :: 3",
|
|
17
|
+
"Operating System :: OS Independent",
|
|
18
|
+
]
|
|
19
|
+
dependencies = [
|
|
20
|
+
"numpy",
|
|
21
|
+
"scipy",
|
|
22
|
+
"astropy",
|
|
23
|
+
"matplotlib",
|
|
24
|
+
"corner",
|
|
25
|
+
"emcee",
|
|
26
|
+
"ultranest"
|
|
27
|
+
]
|
|
28
|
+
license = "MIT"
|
|
29
|
+
license-files = ["LICENSE"]
|
|
30
|
+
keywords = ["binaries","astronomy","regression","model_fitting"]
|
|
@@ -0,0 +1,71 @@
|
|
|
1
|
+
from .data import Data
|
|
2
|
+
from periapsis.model.orbit import Orbit
|
|
3
|
+
import numpy as np
|
|
4
|
+
|
|
5
|
+
class AstrometryData(Data):
|
|
6
|
+
def __init__(self, t, x, y, x_err, y_err,ref_epoch=None,mu_x=None, mu_y=None):
|
|
7
|
+
self.t = t
|
|
8
|
+
self.x = x
|
|
9
|
+
self.y = y
|
|
10
|
+
self.x_err = x_err
|
|
11
|
+
self.y_err = y_err
|
|
12
|
+
|
|
13
|
+
if ref_epoch is None:
|
|
14
|
+
self.ref_epoch = np.mean(t)
|
|
15
|
+
else:
|
|
16
|
+
self.ref_epoch = ref_epoch
|
|
17
|
+
|
|
18
|
+
if mu_x is not None and mu_y is not None:
|
|
19
|
+
self.mu_x = mu_x
|
|
20
|
+
self.mu_y = mu_y
|
|
21
|
+
else:
|
|
22
|
+
self.mu_x = None
|
|
23
|
+
self.mu_y = None
|
|
24
|
+
|
|
25
|
+
def chi2(self, orbit: Orbit):
|
|
26
|
+
x, y = orbit.astrometry(self.t)
|
|
27
|
+
chi2_x = np.sum(((self.x - x) / self.x_err) ** 2)
|
|
28
|
+
chi2_y = np.sum(((self.y - y) / self.y_err) ** 2)
|
|
29
|
+
return chi2_x + chi2_y
|
|
30
|
+
|
|
31
|
+
def t_series(self):
|
|
32
|
+
return self.x, self.y,None, self.t
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
class RadialVelocityData(Data):
|
|
36
|
+
def __init__(self, t, rv, rv_err):
|
|
37
|
+
self.t = t
|
|
38
|
+
self.rv = rv
|
|
39
|
+
self.rv_err = rv_err
|
|
40
|
+
|
|
41
|
+
def chi2(self, orbit: Orbit):
|
|
42
|
+
vz = orbit.radial_velocity(self.t)
|
|
43
|
+
chi2_rv = np.sum(((self.rv - vz) / self.rv_err) ** 2)
|
|
44
|
+
return chi2_rv
|
|
45
|
+
|
|
46
|
+
def t_series(self):
|
|
47
|
+
return None, None,self.rv, self.t
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
class AstroRVData(Data):
|
|
51
|
+
def __init__(self, t, x, x_err, y, y_err, rv, rv_err):
|
|
52
|
+
self.t = t
|
|
53
|
+
self.x = x
|
|
54
|
+
self.x_err = x_err
|
|
55
|
+
self.y = y
|
|
56
|
+
self.y_err = y_err
|
|
57
|
+
self.rv = rv
|
|
58
|
+
self.rv_err = rv_err
|
|
59
|
+
|
|
60
|
+
def chi2(self, orbit: Orbit):
|
|
61
|
+
x, y = orbit.astrometry(self.t)
|
|
62
|
+
vz = orbit.radial_velocity(self.t)
|
|
63
|
+
chi2_x = np.sum(((self.x - x) / self.x_err) ** 2)
|
|
64
|
+
chi2_y = np.sum(((self.y - y) / self.y_err) ** 2)
|
|
65
|
+
chi2_rv = np.sum(((self.rv - vz) / self.rv_err) ** 2)
|
|
66
|
+
return chi2_x + chi2_y + chi2_rv
|
|
67
|
+
|
|
68
|
+
def t_series(self):
|
|
69
|
+
return self.x, self.y,self.rv, self.t
|
|
70
|
+
|
|
71
|
+
|
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
from abc import ABC, abstractmethod
|
|
2
|
+
|
|
3
|
+
import numpy as np
|
|
4
|
+
|
|
5
|
+
from periapsis.model.orbit import Orbit
|
|
6
|
+
|
|
7
|
+
class Data(ABC):
|
|
8
|
+
"""
|
|
9
|
+
A Data object represents the observational data that we want to fit an orbit to.
|
|
10
|
+
It can be extended to include different types of data, such as astrometry, radial velocities, etc.
|
|
11
|
+
|
|
12
|
+
Classes that extend this class will map alternate representations to an absolute 7-dimensional t, x, y, z, vx, vy, vz format that can be used to fit orbits. At least one dimension besides time needs to be available for the given data.
|
|
13
|
+
"""
|
|
14
|
+
|
|
15
|
+
t: np.ndarray
|
|
16
|
+
|
|
17
|
+
@abstractmethod
|
|
18
|
+
def chi2(self, orbit: Orbit):
|
|
19
|
+
"""
|
|
20
|
+
Computes the chi-squared value of the given orbit parameters compared to the data.
|
|
21
|
+
|
|
22
|
+
Parameters
|
|
23
|
+
----------
|
|
24
|
+
orbit: Orbit
|
|
25
|
+
The orbit for which to compute the chi-squared value.
|
|
26
|
+
|
|
27
|
+
Returns
|
|
28
|
+
-------
|
|
29
|
+
chi2 : float
|
|
30
|
+
The chi-squared value of the given orbit parameters compared to the data.
|
|
31
|
+
"""
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
pass
|
|
35
|
+
|
|
36
|
+
def t_series(self):
|
|
37
|
+
"""
|
|
38
|
+
Returns x_obs,y_obs,rv_obs, t_obs"""
|
|
39
|
+
|
|
40
|
+
pass
|
|
@@ -0,0 +1,10 @@
|
|
|
1
|
+
from .data import Data
|
|
2
|
+
from periapsis.model.orbit import Orbit
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
class GaiaData(Data):
|
|
6
|
+
def __init__(self, t, *args):
|
|
7
|
+
raise NotImplementedError("GaiaData is not implemented yet. This is a placeholder for the actual implementation of Gaia data handling.")
|
|
8
|
+
|
|
9
|
+
def chi2(self, orbit: Orbit):
|
|
10
|
+
raise NotImplementedError("GaiaData chi2 method is not implemented yet")
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
from typing import List
|
|
2
|
+
|
|
3
|
+
from .data import Data
|
|
4
|
+
from periapsis.model.orbit import Orbit
|
|
5
|
+
|
|
6
|
+
class JointData(Data):
|
|
7
|
+
def __init__(self, datas: List[Data]):
|
|
8
|
+
self.datas = datas
|
|
9
|
+
|
|
10
|
+
def chi2(self, orbit: Orbit):
|
|
11
|
+
total_chi2 = 0
|
|
12
|
+
for data in self.datas:
|
|
13
|
+
total_chi2 += data.chi2(orbit)
|
|
14
|
+
return total_chi2
|
|
@@ -0,0 +1,88 @@
|
|
|
1
|
+
from abc import ABC, abstractmethod
|
|
2
|
+
import numpy as np
|
|
3
|
+
from periapsis.data.data import Data
|
|
4
|
+
from periapsis.fitting.results import FitResults
|
|
5
|
+
from periapsis.utils.helpers import _match_param_keys
|
|
6
|
+
|
|
7
|
+
class Fitter(ABC):
|
|
8
|
+
def __init__(self, m1=None, **prior_kwargs):
|
|
9
|
+
"""
|
|
10
|
+
A Fitter defines the configuration for fitting an orbit to data, including the priors on the orbital parameters.
|
|
11
|
+
"""
|
|
12
|
+
self.m1 = m1
|
|
13
|
+
self.prior_kwargs = _match_param_keys(prior_kwargs)
|
|
14
|
+
|
|
15
|
+
@abstractmethod
|
|
16
|
+
def fit(self, data: Data) -> FitResults:
|
|
17
|
+
"""
|
|
18
|
+
Fits the orbit to the given data.
|
|
19
|
+
|
|
20
|
+
Parameters
|
|
21
|
+
----------
|
|
22
|
+
data : Data
|
|
23
|
+
Data to fit the orbit to.
|
|
24
|
+
|
|
25
|
+
Returns
|
|
26
|
+
-------
|
|
27
|
+
fit_results : FitResults
|
|
28
|
+
The results of the fit
|
|
29
|
+
"""
|
|
30
|
+
pass
|
|
31
|
+
|
|
32
|
+
def _proper_motion_fit(self, data: Data):
|
|
33
|
+
"""
|
|
34
|
+
Fits a proper motion model to the given data.
|
|
35
|
+
|
|
36
|
+
Parameters
|
|
37
|
+
----------
|
|
38
|
+
data : Data
|
|
39
|
+
Data to fit the proper motion model to.
|
|
40
|
+
|
|
41
|
+
Returns
|
|
42
|
+
-------
|
|
43
|
+
results : dict
|
|
44
|
+
The results of the proper motion fit
|
|
45
|
+
"""
|
|
46
|
+
ref_epoch = getattr(data, 'ref_epoch', np.mean(data.t))
|
|
47
|
+
dt = data.t - ref_epoch
|
|
48
|
+
|
|
49
|
+
if getattr(data, 'mu_x', None) is not None and getattr(data, 'mu_y', None) is not None:
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
x0 = np.sum((data.x - data.mu_x * dt) / data.x_err**2) / np.sum(1 / data.x_err**2)
|
|
53
|
+
y0 = np.sum((data.y - data.mu_y * dt) / data.y_err**2) / np.sum(1 / data.y_err**2)
|
|
54
|
+
|
|
55
|
+
mu_x = data.mu_x
|
|
56
|
+
mu_y = data.mu_y
|
|
57
|
+
|
|
58
|
+
dof = 2 * len(data.t) - 2
|
|
59
|
+
else:
|
|
60
|
+
|
|
61
|
+
A_x = np.vstack([np.ones_like(dt)/data.x_err,dt/data.x_err]).T
|
|
62
|
+
b_x = data.x/data.x_err
|
|
63
|
+
x0,mu_x = np.linalg.lstsq(A_x, b_x,rcond=None)[0]
|
|
64
|
+
|
|
65
|
+
A_y = np.vstack([np.ones_like(dt)/data.y_err,dt/data.y_err]).T
|
|
66
|
+
b_y = data.y/data.y_err
|
|
67
|
+
y0,mu_y = np.linalg.lstsq(A_y, b_y,rcond=None)[0]
|
|
68
|
+
dof = 2*len(data.t)-4
|
|
69
|
+
|
|
70
|
+
chi2_x = np.sum((data.x-(x0+mu_x*dt))**2/data.x_err**2)
|
|
71
|
+
chi2_y = np.sum((data.y-(y0+mu_y*dt))**2/data.y_err**2)
|
|
72
|
+
chi2 = chi2_x + chi2_y
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
return {'params':{'x0':x0,'mu_x':mu_x,'y0':y0,'mu_y':mu_y},
|
|
76
|
+
'chi2':chi2,'dof':dof}
|
|
77
|
+
|
|
78
|
+
def _astrometric_offset_seeds(self, data: Data):
|
|
79
|
+
"""Return sensible starting values for optional astrometric offsets."""
|
|
80
|
+
pm_fit = self._proper_motion_fit(data)
|
|
81
|
+
return {
|
|
82
|
+
'dx': pm_fit['params']['x0'],
|
|
83
|
+
'dy': pm_fit['params']['y0'],
|
|
84
|
+
'dpmra': pm_fit['params']['mu_x'],
|
|
85
|
+
'dpmdec': pm_fit['params']['mu_y'],
|
|
86
|
+
}
|
|
87
|
+
|
|
88
|
+
|
|
@@ -0,0 +1,121 @@
|
|
|
1
|
+
from .fitter import Fitter
|
|
2
|
+
from periapsis.data.data import Data
|
|
3
|
+
from periapsis.fitting.results import FitResults
|
|
4
|
+
from periapsis.model.campbell import CampbellOrbit
|
|
5
|
+
from periapsis.initial.initial import InitialFit
|
|
6
|
+
from periapsis.utils.helpers import _match_param_keys
|
|
7
|
+
from periapsis.utils.solvers import solve_mass
|
|
8
|
+
import numpy as np
|
|
9
|
+
import emcee
|
|
10
|
+
from typing import cast
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def _log_prior(params, param_order,prior_kwargs,m1=None,m2_max=None):
|
|
14
|
+
lp = 0.0
|
|
15
|
+
for name, val in zip(param_order, params):
|
|
16
|
+
prior = prior_kwargs.get(name)
|
|
17
|
+
if prior is not None:
|
|
18
|
+
lp += prior.logpdf(val)
|
|
19
|
+
if not np.isfinite(lp):
|
|
20
|
+
return -np.inf
|
|
21
|
+
else:
|
|
22
|
+
print(f"Warning:Missing prior for {name}.")
|
|
23
|
+
|
|
24
|
+
if m1 is not None and m2_max is not None:
|
|
25
|
+
params_dict = _match_param_keys(dict(zip(param_order, params)))
|
|
26
|
+
m2 = solve_mass(params_dict['a'], params_dict['P'], m1)
|
|
27
|
+
if not np.isfinite(m2) or m2 > m2_max:
|
|
28
|
+
return -np.inf
|
|
29
|
+
return lp
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def _campbell_log_like(params, data, param_order):
|
|
33
|
+
params_dict = _match_param_keys(dict(zip(param_order, params)))
|
|
34
|
+
model = CampbellOrbit(ref_epoch=getattr(data, 'ref_epoch', None), **params_dict)
|
|
35
|
+
chi2 = data.chi2(model)
|
|
36
|
+
if not np.isfinite(chi2):
|
|
37
|
+
return -np.inf
|
|
38
|
+
return -0.5 * chi2
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
def _campbell_log_posterior(params, data, param_order,prior_kwargs,m1=None,m2_max=None):
|
|
42
|
+
lp = _log_prior(params, param_order, prior_kwargs,m1=m1,m2_max=m2_max)
|
|
43
|
+
if not np.isfinite(lp):
|
|
44
|
+
return -np.inf
|
|
45
|
+
return lp + _campbell_log_like(params, data, param_order)
|
|
46
|
+
|
|
47
|
+
class MCMCCampbell(Fitter):
|
|
48
|
+
def __init__(self, nwalkers, niter, m1=None,m2_max=None, pool=None, **priors):
|
|
49
|
+
super().__init__(m1=m1, **priors)
|
|
50
|
+
self.nwalkers = nwalkers
|
|
51
|
+
self.niter = niter
|
|
52
|
+
self.pool = pool
|
|
53
|
+
self.m2_max = m2_max
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
def fit(self, data: Data) -> FitResults:
|
|
57
|
+
|
|
58
|
+
pm_fit = self._proper_motion_fit(data)
|
|
59
|
+
|
|
60
|
+
param_order = list(self.prior_kwargs.keys())
|
|
61
|
+
ndim = len(param_order)
|
|
62
|
+
|
|
63
|
+
initial_dict = InitialFit(data,method='Campbell', **self.prior_kwargs).get_intial()
|
|
64
|
+
initial_guess = np.array([initial_dict[name] for name in param_order])
|
|
65
|
+
bounds = np.array(
|
|
66
|
+
[[self.prior_kwargs[name].min, self.prior_kwargs[name].max] for name in param_order],
|
|
67
|
+
dtype=float,
|
|
68
|
+
)
|
|
69
|
+
lower = bounds[:, 0]
|
|
70
|
+
upper = bounds[:, 1]
|
|
71
|
+
initial_guess = np.clip(initial_guess, lower, upper)
|
|
72
|
+
|
|
73
|
+
pos = np.clip(initial_guess + 1e-4 * np.random.randn(self.nwalkers, ndim), lower, upper)
|
|
74
|
+
|
|
75
|
+
sampler = emcee.EnsembleSampler(
|
|
76
|
+
self.nwalkers,
|
|
77
|
+
ndim,
|
|
78
|
+
_campbell_log_posterior,
|
|
79
|
+
args=(data, param_order, self.prior_kwargs, self.m1, self.m2_max),
|
|
80
|
+
pool=self.pool,
|
|
81
|
+
)
|
|
82
|
+
sampler.run_mcmc(pos, self.niter, progress=True)
|
|
83
|
+
|
|
84
|
+
chain = cast(np.ndarray, sampler.get_chain())
|
|
85
|
+
tau = emcee.autocorr.integrated_time(chain, quiet=True)
|
|
86
|
+
Ess = (self.niter*self.nwalkers)/tau
|
|
87
|
+
mean_acceptance_fraction = np.mean(sampler.acceptance_fraction)
|
|
88
|
+
|
|
89
|
+
burn = int(np.nanmax(tau) * 2)
|
|
90
|
+
thin = int(np.nanmin(tau) * 2)
|
|
91
|
+
|
|
92
|
+
samples = cast(np.ndarray, sampler.get_chain(discard=burn, thin=thin, flat=True))
|
|
93
|
+
lnprobs = cast(np.ndarray, sampler.get_log_prob(discard=burn, thin=thin, flat=True))
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
param_means = chain.mean(axis=1)
|
|
97
|
+
|
|
98
|
+
best_i = np.argmax(lnprobs)
|
|
99
|
+
best_params = dict(zip(param_order, samples[best_i]))
|
|
100
|
+
median_params = dict(zip(param_order, np.median(samples, axis=0)))
|
|
101
|
+
|
|
102
|
+
results_dict = {}
|
|
103
|
+
for i, name in enumerate(param_order):
|
|
104
|
+
results_dict[name] = samples[:, i]
|
|
105
|
+
|
|
106
|
+
results_dict['lnprob'] = lnprobs
|
|
107
|
+
results_dict['Ess'] = Ess
|
|
108
|
+
results_dict['mean_acceptance_fraction'] = mean_acceptance_fraction
|
|
109
|
+
results_dict['tau'] = tau
|
|
110
|
+
results_dict['param_means'] = param_means
|
|
111
|
+
results_dict['param_names'] = param_order
|
|
112
|
+
results_dict['MAP_params'] = best_params
|
|
113
|
+
results_dict['median_params'] = median_params
|
|
114
|
+
results_dict['PM_fit'] = pm_fit
|
|
115
|
+
results_dict['ref_epoch'] = getattr(data, 'ref_epoch', None)
|
|
116
|
+
results_dict['raw_sampler'] = None
|
|
117
|
+
results_dict['backend'] = 'emcee'
|
|
118
|
+
results_dict['fit_method'] = 'Campbell'
|
|
119
|
+
fit_results = FitResults(**results_dict)
|
|
120
|
+
fit_results.add_mass_samples(m1=self.m1)
|
|
121
|
+
return fit_results
|
|
@@ -0,0 +1,107 @@
|
|
|
1
|
+
from .fitter import Fitter
|
|
2
|
+
from periapsis.data.data import Data
|
|
3
|
+
from periapsis.fitting.results import FitResults
|
|
4
|
+
from periapsis.model import thieleinnes
|
|
5
|
+
from periapsis.initial.initial import InitialFit
|
|
6
|
+
from periapsis.utils.solvers import transform_theile
|
|
7
|
+
from periapsis.utils.helpers import _match_param_keys
|
|
8
|
+
import numpy as np
|
|
9
|
+
import emcee
|
|
10
|
+
|
|
11
|
+
class MCMCThieleInnes(Fitter):
|
|
12
|
+
def __init__(self, nwalkers, niter, m1=None, **priors):
|
|
13
|
+
super().__init__(m1=m1, **priors)
|
|
14
|
+
self.nwalkers = nwalkers
|
|
15
|
+
self.niter = niter
|
|
16
|
+
|
|
17
|
+
def fit(self, data: Data) -> FitResults:
|
|
18
|
+
pm_fit = self._proper_motion_fit(data)
|
|
19
|
+
|
|
20
|
+
param_order = list(self.prior_kwargs.keys())
|
|
21
|
+
ndim = len(param_order)
|
|
22
|
+
|
|
23
|
+
def ln_prior(params):
|
|
24
|
+
#this is a placeholder, we will need to figure out transformations
|
|
25
|
+
lp = 0
|
|
26
|
+
for name,val in zip(param_order, params):
|
|
27
|
+
prior = self.prior_kwargs.get(name)
|
|
28
|
+
if prior is not None:
|
|
29
|
+
lp += prior.logpdf(val)
|
|
30
|
+
if not np.isfinite(lp):
|
|
31
|
+
return -np.inf
|
|
32
|
+
else:
|
|
33
|
+
print(f"Warning:Missing prior for {name}.")
|
|
34
|
+
return lp
|
|
35
|
+
|
|
36
|
+
def ln_like(params, data):
|
|
37
|
+
params_dict = _match_param_keys(dict(zip(param_order, params)))
|
|
38
|
+
model = thieleinnes.ThieleInnesOrbit(ref_epoch=getattr(data, 'ref_epoch', None), **params_dict)
|
|
39
|
+
chi2 = data.chi2(model)
|
|
40
|
+
if not np.isfinite(chi2):
|
|
41
|
+
return -np.inf
|
|
42
|
+
return -0.5 * chi2
|
|
43
|
+
|
|
44
|
+
def ln_prob(params, data):
|
|
45
|
+
lp = ln_prior(params)
|
|
46
|
+
if not np.isfinite(lp):
|
|
47
|
+
return -np.inf
|
|
48
|
+
return lp + ln_like(params, data)
|
|
49
|
+
|
|
50
|
+
initial_dict = InitialFit(data,method='ThieleInnes', **self.prior_kwargs).get_intial()
|
|
51
|
+
initial_guess = np.array([initial_dict[name] for name in param_order])
|
|
52
|
+
bounds = np.array(
|
|
53
|
+
[[self.prior_kwargs[name].min, self.prior_kwargs[name].max] for name in param_order],
|
|
54
|
+
dtype=float,
|
|
55
|
+
)
|
|
56
|
+
lower = bounds[:, 0]
|
|
57
|
+
upper = bounds[:, 1]
|
|
58
|
+
initial_guess = np.clip(initial_guess, lower, upper)
|
|
59
|
+
|
|
60
|
+
pos = np.clip(initial_guess + 1e-4 * np.random.randn(self.nwalkers, ndim), lower, upper)
|
|
61
|
+
|
|
62
|
+
sampler = emcee.EnsembleSampler(self.nwalkers, ndim, ln_prob, args=(data,))
|
|
63
|
+
sampler.run_mcmc(pos, self.niter, progress=True)
|
|
64
|
+
|
|
65
|
+
chain = sampler.get_chain()
|
|
66
|
+
param_means = chain.mean(axis=1)
|
|
67
|
+
|
|
68
|
+
tau = emcee.autocorr.integrated_time(chain,quiet=True)
|
|
69
|
+
|
|
70
|
+
Ess = (self.niter*self.nwalkers)/tau
|
|
71
|
+
|
|
72
|
+
maf = np.mean(sampler.acceptance_fraction)
|
|
73
|
+
|
|
74
|
+
burn = int(np.nanmax(tau) * 2)
|
|
75
|
+
thin = int(np.nanmin(tau) * 2)
|
|
76
|
+
|
|
77
|
+
samples = sampler.get_chain(discard=burn,thin=thin,flat=True)
|
|
78
|
+
lnprobs = sampler.get_log_prob(discard=burn,thin=thin,flat=True)
|
|
79
|
+
|
|
80
|
+
best_i = np.argmax(lnprobs)
|
|
81
|
+
best_params = dict(zip(param_order, samples[best_i]))
|
|
82
|
+
|
|
83
|
+
median_params = dict(zip(param_order, np.median(samples, axis=0)))
|
|
84
|
+
|
|
85
|
+
results_dict = {}
|
|
86
|
+
for i, name in enumerate(param_order):
|
|
87
|
+
results_dict[name] = samples[:, i]
|
|
88
|
+
results_dict['lnprob'] = lnprobs
|
|
89
|
+
results_dict['Ess'] = Ess
|
|
90
|
+
results_dict['mean_acceptance_fraction'] = maf
|
|
91
|
+
results_dict['tau'] = tau
|
|
92
|
+
results_dict['param_means'] = param_means
|
|
93
|
+
results_dict['param_names'] = param_order
|
|
94
|
+
results_dict['MAP_params'] = best_params
|
|
95
|
+
results_dict['median_params'] = median_params
|
|
96
|
+
results_dict['PM_fit'] = pm_fit
|
|
97
|
+
results_dict['ref_epoch'] = getattr(data, 'ref_epoch', None)
|
|
98
|
+
|
|
99
|
+
results_dict['backend'] = 'emcee'
|
|
100
|
+
results_dict['fit_method'] = 'ThieleInnes'
|
|
101
|
+
results_dict['raw_sampler'] = None
|
|
102
|
+
|
|
103
|
+
fit_results = FitResults(**results_dict)
|
|
104
|
+
fit_results.add_mass_samples(m1=self.m1)
|
|
105
|
+
return fit_results
|
|
106
|
+
|
|
107
|
+
|