periapsis 0.1.0__py3-none-any.whl

periapsis/__init__.py

@@ -0,0 +1,5 @@
+"""
+periapsis - Orbit modeling and fitting with flexible parameters
+"""
+
+__all__ = ["data", "fitting", "model", "prior"]

periapsis/data/__init__.py

@@ -0,0 +1,6 @@
+from .data import Data
+from .common import AstrometryData, RadialVelocityData, AstroRVData
+# from .gaia import GaiaData
+from .joint_data import JointData
+
+__all__ = ["Data", "AstrometryData", "RadialVelocityData", "AstroRVData", "JointData"]

periapsis/data/common.py

@@ -0,0 +1,71 @@
+from .data import Data
+from periapsis.model.orbit import Orbit
+import numpy as np
+
+class AstrometryData(Data):
+    def __init__(self, t, x, y, x_err, y_err,ref_epoch=None,mu_x=None, mu_y=None):
+        self.t = t
+        self.x = x
+        self.y = y
+        self.x_err = x_err
+        self.y_err = y_err
+
+        if ref_epoch is None:
+            self.ref_epoch = np.mean(t)
+        else:
+            self.ref_epoch = ref_epoch
+
+        if mu_x is not None and mu_y is not None:
+            self.mu_x = mu_x
+            self.mu_y = mu_y
+        else:
+            self.mu_x = None
+            self.mu_y = None
+
+    def chi2(self, orbit: Orbit):
+        x, y = orbit.astrometry(self.t)
+        chi2_x = np.sum(((self.x - x) / self.x_err) ** 2)
+        chi2_y = np.sum(((self.y - y) / self.y_err) ** 2)
+        return chi2_x + chi2_y
+    
+    def t_series(self):
+        return self.x, self.y,None, self.t
+    
+
+class RadialVelocityData(Data):
+    def __init__(self, t, rv, rv_err):
+        self.t = t
+        self.rv = rv
+        self.rv_err = rv_err
+
+    def chi2(self, orbit: Orbit):
+        vz = orbit.radial_velocity(self.t)
+        chi2_rv = np.sum(((self.rv - vz) / self.rv_err) ** 2)
+        return chi2_rv
+    
+    def t_series(self):
+        return None, None,self.rv, self.t
+
+
+class AstroRVData(Data):
+    def __init__(self, t, x, x_err, y, y_err, rv, rv_err):
+        self.t = t
+        self.x = x
+        self.x_err = x_err
+        self.y = y
+        self.y_err = y_err
+        self.rv = rv
+        self.rv_err = rv_err
+
+    def chi2(self, orbit: Orbit):
+        x, y = orbit.astrometry(self.t)
+        vz = orbit.radial_velocity(self.t)
+        chi2_x = np.sum(((self.x - x) / self.x_err) ** 2)
+        chi2_y = np.sum(((self.y - y) / self.y_err) ** 2)
+        chi2_rv = np.sum(((self.rv - vz) / self.rv_err) ** 2)
+        return chi2_x + chi2_y + chi2_rv
+    
+    def t_series(self):
+        return self.x, self.y,self.rv, self.t
+    
+    

periapsis/data/data.py

@@ -0,0 +1,40 @@
+from abc import ABC, abstractmethod
+
+import numpy as np
+
+from periapsis.model.orbit import Orbit
+
+class Data(ABC):
+    """
+    A Data object represents the observational data that we want to fit an orbit to. 
+    It can be extended to include different types of data, such as astrometry, radial velocities, etc.
+
+    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.
+    """
+
+    t: np.ndarray
+
+    @abstractmethod
+    def chi2(self, orbit: Orbit):
+        """
+        Computes the chi-squared value of the given orbit parameters compared to the data. 
+        
+        Parameters
+        ----------
+        orbit: Orbit
+            The orbit for which to compute the chi-squared value.
+
+        Returns
+        -------
+        chi2 : float
+            The chi-squared value of the given orbit parameters compared to the data.
+        """
+
+        
+        pass
+
+    def t_series(self):
+        """
+        Returns x_obs,y_obs,rv_obs, t_obs"""
+        
+        pass

periapsis/data/gaia.py

@@ -0,0 +1,10 @@
+from .data import Data
+from periapsis.model.orbit import Orbit
+
+
+class GaiaData(Data):
+    def __init__(self, t, *args):
+        raise NotImplementedError("GaiaData is not implemented yet. This is a placeholder for the actual implementation of Gaia data handling.")
+
+    def chi2(self, orbit: Orbit):
+        raise NotImplementedError("GaiaData chi2 method is not implemented yet")

periapsis/data/joint_data.py

@@ -0,0 +1,14 @@
+from typing import List
+
+from .data import Data
+from periapsis.model.orbit import Orbit
+
+class JointData(Data):
+    def __init__(self, datas: List[Data]):
+        self.datas = datas
+
+    def chi2(self, orbit: Orbit):
+        total_chi2 = 0
+        for data in self.datas:
+            total_chi2 += data.chi2(orbit)
+        return total_chi2

periapsis/fitting/__init__.py

@@ -0,0 +1,4 @@
+from .fitter import Fitter
+from .results import FitResults
+
+__all__ = ["Fitter", "FitResults"]

periapsis/fitting/fitter.py

@@ -0,0 +1,88 @@
+from abc import ABC, abstractmethod
+import numpy as np
+from periapsis.data.data import Data
+from periapsis.fitting.results import FitResults
+from periapsis.utils.helpers import _match_param_keys
+
+class Fitter(ABC):
+    def __init__(self, m1=None, **prior_kwargs):
+        """
+        A Fitter defines the configuration for fitting an orbit to data, including the priors on the orbital parameters.
+        """
+        self.m1 = m1
+        self.prior_kwargs = _match_param_keys(prior_kwargs)
+
+    @abstractmethod
+    def fit(self, data: Data) -> FitResults:
+        """
+        Fits the orbit to the given data.
+
+        Parameters
+        ----------
+        data : Data
+            Data to fit the orbit to.
+
+        Returns
+        -------
+        fit_results : FitResults
+            The results of the fit
+        """
+        pass
+
+    def _proper_motion_fit(self, data: Data):
+        """
+        Fits a proper motion model to the given data.
+
+        Parameters
+        ----------
+        data : Data
+            Data to fit the proper motion model to.
+
+        Returns
+        -------
+        results : dict
+            The results of the proper motion fit
+        """
+        ref_epoch = getattr(data, 'ref_epoch', np.mean(data.t))
+        dt = data.t - ref_epoch
+
+        if getattr(data, 'mu_x', None) is not None and getattr(data, 'mu_y', None) is not None:
+            
+            
+            x0 = np.sum((data.x - data.mu_x * dt) / data.x_err**2) / np.sum(1 / data.x_err**2)
+            y0 = np.sum((data.y - data.mu_y * dt) / data.y_err**2) / np.sum(1 / data.y_err**2)
+
+            mu_x = data.mu_x
+            mu_y = data.mu_y
+    
+            dof = 2 * len(data.t) - 2
+        else:
+            
+            A_x = np.vstack([np.ones_like(dt)/data.x_err,dt/data.x_err]).T
+            b_x = data.x/data.x_err
+            x0,mu_x = np.linalg.lstsq(A_x, b_x,rcond=None)[0]
+        
+            A_y = np.vstack([np.ones_like(dt)/data.y_err,dt/data.y_err]).T
+            b_y = data.y/data.y_err
+            y0,mu_y = np.linalg.lstsq(A_y, b_y,rcond=None)[0]
+            dof = 2*len(data.t)-4
+
+        chi2_x = np.sum((data.x-(x0+mu_x*dt))**2/data.x_err**2)
+        chi2_y = np.sum((data.y-(y0+mu_y*dt))**2/data.y_err**2)
+        chi2 = chi2_x + chi2_y
+        
+
+        return {'params':{'x0':x0,'mu_x':mu_x,'y0':y0,'mu_y':mu_y},
+                'chi2':chi2,'dof':dof}
+
+    def _astrometric_offset_seeds(self, data: Data):
+        """Return sensible starting values for optional astrometric offsets."""
+        pm_fit = self._proper_motion_fit(data)
+        return {
+            'dx': pm_fit['params']['x0'],
+            'dy': pm_fit['params']['y0'],
+            'dpmra': pm_fit['params']['mu_x'],
+            'dpmdec': pm_fit['params']['mu_y'],
+        }
+
+    

periapsis/fitting/mcmcCampbell.py

@@ -0,0 +1,121 @@
+from .fitter import Fitter
+from periapsis.data.data import Data
+from periapsis.fitting.results import FitResults
+from periapsis.model.campbell import CampbellOrbit
+from periapsis.initial.initial import InitialFit
+from periapsis.utils.helpers import _match_param_keys
+from periapsis.utils.solvers import solve_mass
+import numpy as np
+import emcee
+from typing import cast
+
+
+def _log_prior(params, param_order,prior_kwargs,m1=None,m2_max=None):
+    lp = 0.0
+    for name, val in zip(param_order, params):
+        prior = prior_kwargs.get(name)
+        if prior is not None:
+            lp += prior.logpdf(val)
+            if not np.isfinite(lp):
+                return -np.inf
+        else:
+            print(f"Warning:Missing prior for {name}.")
+
+    if m1 is not None and m2_max is not None:
+        params_dict = _match_param_keys(dict(zip(param_order, params)))
+        m2 = solve_mass(params_dict['a'], params_dict['P'], m1)
+        if not np.isfinite(m2) or m2 > m2_max:
+            return -np.inf
+    return lp
+
+
+def _campbell_log_like(params, data, param_order):
+    params_dict = _match_param_keys(dict(zip(param_order, params)))
+    model = CampbellOrbit(ref_epoch=getattr(data, 'ref_epoch', None), **params_dict)
+    chi2 = data.chi2(model)
+    if not np.isfinite(chi2):
+        return -np.inf
+    return -0.5 * chi2
+
+
+def _campbell_log_posterior(params, data, param_order,prior_kwargs,m1=None,m2_max=None):
+    lp = _log_prior(params, param_order, prior_kwargs,m1=m1,m2_max=m2_max)
+    if not np.isfinite(lp):
+        return -np.inf
+    return lp + _campbell_log_like(params, data, param_order)
+
+class MCMCCampbell(Fitter):
+    def __init__(self, nwalkers, niter, m1=None,m2_max=None, pool=None, **priors):
+        super().__init__(m1=m1, **priors)
+        self.nwalkers = nwalkers
+        self.niter = niter
+        self.pool = pool
+        self.m2_max = m2_max
+        
+
+    def fit(self, data: Data) -> FitResults:
+
+        pm_fit = self._proper_motion_fit(data)
+
+        param_order = list(self.prior_kwargs.keys())
+        ndim = len(param_order)
+        
+        initial_dict = InitialFit(data,method='Campbell', **self.prior_kwargs).get_intial()
+        initial_guess = np.array([initial_dict[name] for name in param_order])
+        bounds = np.array(
+            [[self.prior_kwargs[name].min, self.prior_kwargs[name].max] for name in param_order],
+            dtype=float,
+        )
+        lower = bounds[:, 0]
+        upper = bounds[:, 1]
+        initial_guess = np.clip(initial_guess, lower, upper)
+
+        pos = np.clip(initial_guess + 1e-4 * np.random.randn(self.nwalkers, ndim), lower, upper)
+
+        sampler = emcee.EnsembleSampler(
+            self.nwalkers,
+            ndim,
+            _campbell_log_posterior,
+            args=(data, param_order, self.prior_kwargs, self.m1, self.m2_max),
+            pool=self.pool,
+        )
+        sampler.run_mcmc(pos, self.niter, progress=True)
+
+        chain = cast(np.ndarray, sampler.get_chain())
+        tau = emcee.autocorr.integrated_time(chain, quiet=True)
+        Ess = (self.niter*self.nwalkers)/tau
+        mean_acceptance_fraction = np.mean(sampler.acceptance_fraction)
+
+        burn = int(np.nanmax(tau) * 2)
+        thin = int(np.nanmin(tau) * 2)
+
+        samples = cast(np.ndarray, sampler.get_chain(discard=burn, thin=thin, flat=True))
+        lnprobs = cast(np.ndarray, sampler.get_log_prob(discard=burn, thin=thin, flat=True))
+
+        
+        param_means = chain.mean(axis=1)
+
+        best_i = np.argmax(lnprobs)
+        best_params = dict(zip(param_order, samples[best_i]))
+        median_params = dict(zip(param_order, np.median(samples, axis=0)))
+        
+        results_dict = {}
+        for i, name in enumerate(param_order):
+            results_dict[name] = samples[:, i]
+
+        results_dict['lnprob'] = lnprobs
+        results_dict['Ess'] = Ess
+        results_dict['mean_acceptance_fraction'] = mean_acceptance_fraction
+        results_dict['tau'] = tau
+        results_dict['param_means'] = param_means
+        results_dict['param_names'] = param_order
+        results_dict['MAP_params'] = best_params
+        results_dict['median_params'] = median_params
+        results_dict['PM_fit'] = pm_fit
+        results_dict['ref_epoch'] = getattr(data, 'ref_epoch', None)
+        results_dict['raw_sampler'] = None
+        results_dict['backend'] = 'emcee'
+        results_dict['fit_method'] = 'Campbell'
+        fit_results = FitResults(**results_dict)
+        fit_results.add_mass_samples(m1=self.m1)
+        return fit_results

periapsis/fitting/mcmcThieleInnes.py

@@ -0,0 +1,107 @@
+from .fitter import Fitter
+from periapsis.data.data import Data
+from periapsis.fitting.results import FitResults
+from periapsis.model import thieleinnes
+from periapsis.initial.initial import InitialFit
+from periapsis.utils.solvers import transform_theile
+from periapsis.utils.helpers import _match_param_keys
+import numpy as np
+import emcee
+
+class MCMCThieleInnes(Fitter):
+    def __init__(self, nwalkers, niter, m1=None, **priors):
+        super().__init__(m1=m1, **priors)
+        self.nwalkers = nwalkers
+        self.niter = niter
+
+    def fit(self, data: Data) -> FitResults:
+        pm_fit = self._proper_motion_fit(data)
+
+        param_order = list(self.prior_kwargs.keys())
+        ndim = len(param_order)
+
+        def ln_prior(params):
+            #this is a placeholder, we will need to figure out transformations
+            lp = 0 
+            for name,val in zip(param_order, params):
+                prior = self.prior_kwargs.get(name)
+                if prior is not None:
+                    lp += prior.logpdf(val)
+                    if not np.isfinite(lp):
+                        return -np.inf
+                else:
+                    print(f"Warning:Missing prior for {name}.")
+            return lp
+
+        def ln_like(params, data):
+            params_dict = _match_param_keys(dict(zip(param_order, params)))
+            model = thieleinnes.ThieleInnesOrbit(ref_epoch=getattr(data, 'ref_epoch', None), **params_dict)
+            chi2 = data.chi2(model)
+            if not np.isfinite(chi2):
+                return -np.inf
+            return -0.5 * chi2
+
+        def ln_prob(params, data):
+            lp = ln_prior(params)
+            if not np.isfinite(lp):
+                return -np.inf
+            return lp + ln_like(params, data)
+            
+        initial_dict = InitialFit(data,method='ThieleInnes', **self.prior_kwargs).get_intial()
+        initial_guess = np.array([initial_dict[name] for name in param_order])
+        bounds = np.array(
+            [[self.prior_kwargs[name].min, self.prior_kwargs[name].max] for name in param_order],
+            dtype=float,
+        )
+        lower = bounds[:, 0]
+        upper = bounds[:, 1]
+        initial_guess = np.clip(initial_guess, lower, upper)
+        
+        pos = np.clip(initial_guess + 1e-4 * np.random.randn(self.nwalkers, ndim), lower, upper)
+
+        sampler = emcee.EnsembleSampler(self.nwalkers, ndim, ln_prob, args=(data,))
+        sampler.run_mcmc(pos, self.niter, progress=True)
+
+        chain = sampler.get_chain()
+        param_means = chain.mean(axis=1)
+
+        tau = emcee.autocorr.integrated_time(chain,quiet=True)
+
+        Ess = (self.niter*self.nwalkers)/tau
+
+        maf = np.mean(sampler.acceptance_fraction)
+
+        burn = int(np.nanmax(tau) * 2)
+        thin = int(np.nanmin(tau) * 2)
+
+        samples = sampler.get_chain(discard=burn,thin=thin,flat=True)
+        lnprobs = sampler.get_log_prob(discard=burn,thin=thin,flat=True)
+
+        best_i = np.argmax(lnprobs)
+        best_params = dict(zip(param_order, samples[best_i]))
+
+        median_params = dict(zip(param_order, np.median(samples, axis=0)))
+
+        results_dict = {}
+        for i, name in enumerate(param_order):
+            results_dict[name] = samples[:, i]
+        results_dict['lnprob'] = lnprobs
+        results_dict['Ess'] = Ess
+        results_dict['mean_acceptance_fraction'] = maf
+        results_dict['tau'] = tau
+        results_dict['param_means'] = param_means
+        results_dict['param_names'] = param_order
+        results_dict['MAP_params'] = best_params
+        results_dict['median_params'] = median_params
+        results_dict['PM_fit'] = pm_fit
+        results_dict['ref_epoch'] = getattr(data, 'ref_epoch', None)
+
+        results_dict['backend'] = 'emcee'
+        results_dict['fit_method'] = 'ThieleInnes'
+        results_dict['raw_sampler'] = None
+
+        fit_results = FitResults(**results_dict)
+        fit_results.add_mass_samples(m1=self.m1)
+        return fit_results
+            
+