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.
Files changed (47) hide show
  1. periapsis-0.1.0/LICENSE +21 -0
  2. periapsis-0.1.0/PKG-INFO +61 -0
  3. periapsis-0.1.0/README.md +40 -0
  4. periapsis-0.1.0/pyproject.toml +30 -0
  5. periapsis-0.1.0/setup.cfg +4 -0
  6. periapsis-0.1.0/src/periapsis/__init__.py +5 -0
  7. periapsis-0.1.0/src/periapsis/data/__init__.py +6 -0
  8. periapsis-0.1.0/src/periapsis/data/common.py +71 -0
  9. periapsis-0.1.0/src/periapsis/data/data.py +40 -0
  10. periapsis-0.1.0/src/periapsis/data/gaia.py +10 -0
  11. periapsis-0.1.0/src/periapsis/data/joint_data.py +14 -0
  12. periapsis-0.1.0/src/periapsis/fitting/__init__.py +4 -0
  13. periapsis-0.1.0/src/periapsis/fitting/fitter.py +88 -0
  14. periapsis-0.1.0/src/periapsis/fitting/mcmcCampbell.py +121 -0
  15. periapsis-0.1.0/src/periapsis/fitting/mcmcThieleInnes.py +107 -0
  16. periapsis-0.1.0/src/periapsis/fitting/mcmclinear.py +249 -0
  17. periapsis-0.1.0/src/periapsis/fitting/results.py +65 -0
  18. periapsis-0.1.0/src/periapsis/fitting/ultranestCampbell.py +90 -0
  19. periapsis-0.1.0/src/periapsis/fitting/ultranestThieleInnes.py +71 -0
  20. periapsis-0.1.0/src/periapsis/fitting/ultranestlinear.py +227 -0
  21. periapsis-0.1.0/src/periapsis/initial/__init__.py +0 -0
  22. periapsis-0.1.0/src/periapsis/initial/initial.py +144 -0
  23. periapsis-0.1.0/src/periapsis/model/__init__.py +3 -0
  24. periapsis-0.1.0/src/periapsis/model/campbell.py +55 -0
  25. periapsis-0.1.0/src/periapsis/model/orbit.py +89 -0
  26. periapsis-0.1.0/src/periapsis/model/thieleinnes.py +48 -0
  27. periapsis-0.1.0/src/periapsis/plotting/__init__.py +0 -0
  28. periapsis-0.1.0/src/periapsis/plotting/plots.py +506 -0
  29. periapsis-0.1.0/src/periapsis/prior/__init__.py +5 -0
  30. periapsis-0.1.0/src/periapsis/prior/normal_prior.py +26 -0
  31. periapsis-0.1.0/src/periapsis/prior/prior.py +10 -0
  32. periapsis-0.1.0/src/periapsis/prior/uniform_prior.py +22 -0
  33. periapsis-0.1.0/src/periapsis/stats/__init__.py +0 -0
  34. periapsis-0.1.0/src/periapsis/stats/stat_funcs.py +163 -0
  35. periapsis-0.1.0/src/periapsis/utils/helpers.py +66 -0
  36. periapsis-0.1.0/src/periapsis/utils/solvers.py +163 -0
  37. periapsis-0.1.0/src/periapsis.egg-info/PKG-INFO +61 -0
  38. periapsis-0.1.0/src/periapsis.egg-info/SOURCES.txt +45 -0
  39. periapsis-0.1.0/src/periapsis.egg-info/dependency_links.txt +1 -0
  40. periapsis-0.1.0/src/periapsis.egg-info/requires.txt +7 -0
  41. periapsis-0.1.0/src/periapsis.egg-info/top_level.txt +1 -0
  42. periapsis-0.1.0/tests/test_data.py +85 -0
  43. periapsis-0.1.0/tests/test_fitting.py +41 -0
  44. periapsis-0.1.0/tests/test_model.py +24 -0
  45. periapsis-0.1.0/tests/test_prior.py +29 -0
  46. periapsis-0.1.0/tests/test_results.py +18 -0
  47. periapsis-0.1.0/tests/test_stats.py +96 -0
@@ -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.
@@ -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,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,5 @@
1
+ """
2
+ periapsis - Orbit modeling and fitting with flexible parameters
3
+ """
4
+
5
+ __all__ = ["data", "fitting", "model", "prior"]
@@ -0,0 +1,6 @@
1
+ from .data import Data
2
+ from .common import AstrometryData, RadialVelocityData, AstroRVData
3
+ # from .gaia import GaiaData
4
+ from .joint_data import JointData
5
+
6
+ __all__ = ["Data", "AstrometryData", "RadialVelocityData", "AstroRVData", "JointData"]
@@ -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,4 @@
1
+ from .fitter import Fitter
2
+ from .results import FitResults
3
+
4
+ __all__ = ["Fitter", "FitResults"]
@@ -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
+