<commit_before># Standard library
from collections import OrderedDict
import copy
import os

# Third-party
import astropy.units as u
from astropy.table import Table, QTable
import numpy as np
from twobody import KeplerOrbit, PolynomialRVTrend

# Project
from .prior_helpers import (_validate_polytrend, _get_nonlinear_equiv_units,
                            _get_linear_equiv_units)
from .samples_helpers import write_table_hdf5

__all__ = ['JokerSamples']


class JokerSamples:
    _hdf5_path = 'samples'

    def __init__(self, samples=None, poly_trend=1, t0=None, **kwargs):
        """
        A dictionary-like object for storing prior or posterior samples from
        The Joker, with some extra functionality.

        Parameters
        ----------
        samples : `~astropy.table.QTable`, table-like (optional)
            The samples data as an Astropy table object, or something
            convertable to an Astropy table (e.g., a dictionary with
            `~astropy.units.Quantity` object values). This is optional because
            the samples data can be added later by setting keys on the resulting
            instance.
        poly_trend : int (optional)
            Specifies the number of coefficients in an additional polynomial
            velocity trend, meant to capture long-term trends in the data. See
            the docstring for `thejoker.JokerPrior` for more details.
        t0 : `astropy.time.Time`, numeric (optional)
            The reference time for the orbital parameters.
        **kwargs
            Additional keyword arguments are stored internally as metadata.
        """
        self.tbl = QTable()

        if isinstance(samples, Table):
            poly_trend = samples.meta.pop('poly_trend', poly_trend)
            t0 = samples.meta.pop('t0', t0)
            kwargs.update(samples.meta)

        poly_trend, v_trend_names = _validate_polytrend(poly_trend)

        self._valid_units = {**_get_nonlinear_equiv_units(),
                             **_get_linear_equiv_units(v_trend_names)}

        # log-prior and log-likelihood values are also valid:
        self._valid_units['ln_prior'] = u.one
        self._valid_units['ln_likelihood'] = u.one

        self.tbl.meta['poly_trend'] = poly_trend
        self.tbl.meta['t0'] = t0
        for k, v in kwargs.items():
            self.tbl.meta[k] = v

        # Doing this here validates the input:
        if samples is not None:
            _tbl = QTable(samples)
            for colname in _tbl.colnames:
                self[colname] = _tbl[colname]

        # used for speed-ups below
        self._cache = dict()

    def __getitem__(self, key):
        if isinstance(key, int):
            key = slice(key, key+1)
            return self.__class__(samples=self.tbl[key])

        elif isinstance(key, str) and key in self.par_names:
            return self.tbl[key]

        else:
            return self.__class__(samples=self.tbl[key])

    def __setitem__(self, key, val):
        if key not in self._valid_units:
            raise ValueError(f"Invalid parameter name '{key}'. Must be one "
                             "of: {0}".format(list(self._valid_units.keys())))

        if not hasattr(val, 'unit'):
            val = val * u.one  # eccentricity

        expected_unit = self._valid_units[key]
        if not val.unit.is_equivalent(expected_unit):
            raise u.UnitsError(f"Units of '{key}' must be convertable to "
                               f"{expected_unit}")

        self.tbl[key] = val

    @property
    def poly_trend(self):
        return self.tbl.meta['poly_trend']

    @property
    def t0(self):
        return self.tbl.meta['t0']

    @property
    def par_names(self):
        return self.tbl.colnames

    def __len__(self):
        return len(self.tbl)

    def __repr__(self):
        return (f'<JokerSamples [{", ".join(self.par_names)}] '
                f'({len(self)} samples)>')

    def __str__(self):
        return self.__repr__()

    ##########################################################################
    # Interaction with TwoBody

    def get_orbit(self, index=None, **kwargs):
        """
        Get a `twobody.KeplerOrbit` object for the samples at the specified
        index.

        Parameters
        ----------
        index : int (optional)
            The index of the samples to turn into a `twobody.KeplerOrbit`
            instance. If the samples object is scalar, no index is necessary.
        **kwargs
            Other keyword arguments are passed to the `twobody.KeplerOrbit`
            initializer. For example, you can specify the inclination by passing
            ``i=...`, or  longitude of the ascending node by passing
            ``Omega=...``.

        Returns
        -------
        orbit : `twobody.KeplerOrbit`
            The samples converted to an orbit object. The barycenter position
            and distance are set to arbitrary values.
        """
        if 'orbit' not in self._cache:
            self._cache['orbit'] = KeplerOrbit(P=1*u.yr, e=0., omega=0*u.deg,
                                               Omega=0*u.deg, i=90*u.deg,
                                               a=1*u.au, t0=self.t0)

        # all of this to avoid the __init__ of KeplerOrbit / KeplerElements
        orbit = copy.copy(self._cache['orbit'])

        P = self['P']
        e = self['e']
        K = self['K']
        omega = self['omega']
        M0 = self['M0']
        a = kwargs.pop('a', P * K / (2*np.pi) * np.sqrt(1 - e**2))

        _, v_trend_names = _validate_polytrend(self.poly_trend)
        names = list(_get_linear_equiv_units(v_trend_names).keys())
        if len(self) == 1:
            if index is not None and index > 0:
                raise ValueError('Samples are scalar-valued!')

            trend_coeffs = [self[x] for x in names[1:]]  # skip K

        else:
            P = P[index]
            e = e[index]
            a = a[index]
            omega = omega[index]
            M0 = M0[index]

            trend_coeffs = [self[x][index] for x in names[1:]]  # skip K

        orbit.elements._P = P
        orbit.elements._e = e * u.dimensionless_unscaled
        orbit.elements._a = a
        orbit.elements._omega = omega
        orbit.elements._M0 = M0
        orbit.elements._Omega = kwargs.pop('Omega', 0*u.deg)
        orbit.elements._i = kwargs.pop('i', 90*u.deg)
        orbit._vtrend = PolynomialRVTrend(trend_coeffs, t0=self.t0)
        orbit._barycenter = kwargs.pop('barycenter', None)<commit_msg>fix orbit for scalar samples<commit_after># Standard library
from collections import OrderedDict
import copy
import os

# Third-party
import astropy.units as u
from astropy.table import Table, QTable
import numpy as np
from twobody import KeplerOrbit, PolynomialRVTrend

# Project
from .prior_helpers import (_validate_polytrend, _get_nonlinear_equiv_units,
                            _get_linear_equiv_units)
from .samples_helpers import write_table_hdf5

__all__ = ['JokerSamples']


class JokerSamples:
    _hdf5_path = 'samples'

    def __init__(self, samples=None, poly_trend=1, t0=None, **kwargs):
        """
        A dictionary-like object for storing prior or posterior samples from
        The Joker, with some extra functionality.

        Parameters
        ----------
        samples : `~astropy.table.QTable`, table-like (optional)
            The samples data as an Astropy table object, or something
            convertable to an Astropy table (e.g., a dictionary with
            `~astropy.units.Quantity` object values). This is optional because
            the samples data can be added later by setting keys on the resulting
            instance.
        poly_trend : int (optional)
            Specifies the number of coefficients in an additional polynomial
            velocity trend, meant to capture long-term trends in the data. See
            the docstring for `thejoker.JokerPrior` for more details.
        t0 : `astropy.time.Time`, numeric (optional)
            The reference time for the orbital parameters.
        **kwargs
            Additional keyword arguments are stored internally as metadata.
        """
        self.tbl = QTable()

        if isinstance(samples, Table):
            poly_trend = samples.meta.pop('poly_trend', poly_trend)
            t0 = samples.meta.pop('t0', t0)
            kwargs.update(samples.meta)

        poly_trend, v_trend_names = _validate_polytrend(poly_trend)

        self._valid_units = {**_get_nonlinear_equiv_units(),
                             **_get_linear_equiv_units(v_trend_names)}

        # log-prior and log-likelihood values are also valid:
        self._valid_units['ln_prior'] = u.one
        self._valid_units['ln_likelihood'] = u.one

        self.tbl.meta['poly_trend'] = poly_trend
        self.tbl.meta['t0'] = t0
        for k, v in kwargs.items():
            self.tbl.meta[k] = v

        # Doing this here validates the input:
        if samples is not None:
            _tbl = QTable(samples)
            for colname in _tbl.colnames:
                self[colname] = _tbl[colname]

        # used for speed-ups below
        self._cache = dict()

    def __getitem__(self, key):
        if isinstance(key, int):
            key = slice(key, key+1)
            return self.__class__(samples=self.tbl[key])

        elif isinstance(key, str) and key in self.par_names:
            return self.tbl[key]

        else:
            return self.__class__(samples=self.tbl[key])

    def __setitem__(self, key, val):
        if key not in self._valid_units:
            raise ValueError(f"Invalid parameter name '{key}'. Must be one "
                             "of: {0}".format(list(self._valid_units.keys())))

        if not hasattr(val, 'unit'):
            val = val * u.one  # eccentricity

        expected_unit = self._valid_units[key]
        if not val.unit.is_equivalent(expected_unit):
            raise u.UnitsError(f"Units of '{key}' must be convertable to "
                               f"{expected_unit}")

        self.tbl[key] = val

    @property
    def poly_trend(self):
        return self.tbl.meta['poly_trend']

    @property
    def t0(self):
        return self.tbl.meta['t0']

    @property
    def par_names(self):
        return self.tbl.colnames

    def __len__(self):
        return len(self.tbl)

    def __repr__(self):
        return (f'<JokerSamples [{", ".join(self.par_names)}] '
                f'({len(self)} samples)>')

    def __str__(self):
        return self.__repr__()

    ##########################################################################
    # Interaction with TwoBody

    def get_orbit(self, index=None, **kwargs):
        """
        Get a `twobody.KeplerOrbit` object for the samples at the specified
        index.

        Parameters
        ----------
        index : int (optional)
            The index of the samples to turn into a `twobody.KeplerOrbit`
            instance. If the samples object is scalar, no index is necessary.
        **kwargs
            Other keyword arguments are passed to the `twobody.KeplerOrbit`
            initializer. For example, you can specify the inclination by passing
            ``i=...`, or  longitude of the ascending node by passing
            ``Omega=...``.

        Returns
        -------
        orbit : `twobody.KeplerOrbit`
            The samples converted to an orbit object. The barycenter position
            and distance are set to arbitrary values.
        """
        if 'orbit' not in self._cache:
            self._cache['orbit'] = KeplerOrbit(P=1*u.yr, e=0., omega=0*u.deg,
                                               Omega=0*u.deg, i=90*u.deg,
                                               a=1*u.au, t0=self.t0)

        # all of this to avoid the __init__ of KeplerOrbit / KeplerElements
        orbit = copy.copy(self._cache['orbit'])

        P = self['P']
        e = self['e']
        K = self['K']
        omega = self['omega']
        M0 = self['M0']
        a = kwargs.pop('a', P * K / (2*np.pi) * np.sqrt(1 - e**2))

        _, v_trend_names = _validate_polytrend(self.poly_trend)
        names = list(_get_linear_equiv_units(v_trend_names).keys())

        if index is None:
            if len(self) == 1:
                index = 0
            else:
                raise ValueError("You must specify an index when the number "
                                 f"of samples is >1 (here, it's {len(self)})")

        P = P[index]
        e = e[index]
        a = a[index]
        omega = omega[index]
        M0 = M0[index]

        trend_coeffs = [self[x][index] for x in names[1:]]  # skip K

        orbit.elements._P = P
        orbit.elements._e = e * u.dimensionless_unscaled
        orbit.elements._a = a
        orbit.elements._omega = omega
        orbit.elements._M0 = M0
        orbit.elements._Omega = kwargs.pop('Omega', 0*u.deg)
        orbit.elements._i = kwargs.pop('i', 90*u.deg)
        orbit._vtrend = PolynomialRVTrend(trend_coeffs, t0=self.t0)
        orbit._barycenter = kwargs.pop('barycenter', None)
{'fraction_non_alphanumeric': 0.11203541173776997, 'fraction_numerical': 0.006258108830038922, 'mean_word_length': 1.985645933014354, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 9, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1262730', 'n_tokens_mistral': 3868, 'n_tokens_neox': 3714, 'n_words': 1061}