<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import numpy as np
import astropy.units as u
import astropy.constants as c

"""
This module contains the massive object classes that make up a star-planet
system.
"""


# The main star class
class MainStar(object):
    """
    The main star of a given system.

    mass : ``astropy.units.Quantity``
        The mass of the star.

    name : ``str`` or ``None``
        The name of the star. Default is ``None``.
    """
    def __init__(self, mass, name=None):
        self.mass = mass
        self.name = name


# The companion sub-class
class Companion(object):
    """
    The massive companion class. It can be either a binary star, an exoplanet,
    maybe even a black hole! It can be anything that has a mass and orbits
    another massive object. General relativity effects are not implemented yet.

    Parameters
    ----------
    k : ``astropy.units.Quantity`` or ``None``, optional
        The radial velocity semi-amplitude. Default is ``None``.

    period_orb : ``astropy.units.Quantity`` or ``None``, optional
        The orbital period. Default is ``None``.

    t_0 : ``astropy.units.Quantity`` or ``None``, optional
        The time of periastron passage. Default is ``None``.

    omega : ``astropy.units.Quantity`` or ``None``, optional
        Argument of periapse. Default is ``None``.

    ecc : ``float`` or ``None``, optional
        Eccentricity of the orbit. Default is ``None``.

    msini : ``astropy.units.Quantity`` or ``None``, optional
        Mass of the companion multiplied by sine of the inclination of the
        orbital plane in relation to the line of sight. Default is ``None``.

    semi_a : ``astropy.units.Quantity`` or ``None``, optional
        Semi-major axis of the orbit. Default is ``None``.

    name : ``str`` or ``None``, optional
        Name of the companion. Default is ``None``.
    """
    def __init__(self, k=None, period_orb=None, t_0=None, omega=None, ecc=None,
                 msini=None, semi_a=None, name=None, main_star=None):
        self.main_star = main_star
        self.msini = msini
        self.semi_a = semi_a
        self.k = k
        self.period_orb = period_orb
        self.t_0 = t_0
        self.omega = omega
        self.ecc = ecc
        self.name = name


# The star-companion system class
class System(object):
    """
    The star-companions system class.

    Parameters
    ----------
    main_star : ``radial.object.MainStar``
        The main star of the system.

    companions : list
        Python list containing the all the ``radial.object.Companion`` of the
        system.

    name : ``str`` or ``None``, optional
        Name of the system. Default is ``None``.
    """
    def __init__(self, main_star, companions, name=None):
        self.name = name
        self.main_star = main_star
        self.companions = companions
        self.n_c = len(self.companions)     # Number of companions

        # Initializing useful global parameters
        self.f = []

    # Compute mass functions of the companions
    def mass_func(self):
        """
        Compute the mass functions of all the companions of the system. This
        method will also compute the msini and semi_a of the companions and
        save the values in their respective parameters.
        """
        for comp in self.companions:
            assert isinstance(comp.k, u.Quantity), 'k needs to be provided.'
            assert isinstance(comp.period_orb, u.Quantity), 'period_orb ' \
                                                            'needs to be ' \
                                                            'provided.'
            assert isinstance(comp.ecc, float), 'ecc needs to be provided.'
            f = (comp.period_orb * comp.k ** 3 * (1 - comp.ecc ** 2) ** (3 / 2)
                 / (2 * np.pi * c.G)).to(u.solMass)
            m_star = self.main_star.mass
            m_roots = np.roots([1,  -f.value,  -2 * m_star.value * f.value,
                                -m_star.value ** 2 * f.value])
            comp.msini = abs(m_roots[0]) * u.solMass
            comp.semi_a = (np.sqrt(c.G / comp.k * comp.msini * comp.period_orb /
                           (2 * np.pi * np.sqrt(1 - comp.ecc ** 2)))).to(u.AU)
            self.f.append(f)<commit_msg>Added methods to compute and plot radial velocities to System class<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import numpy as np
import astropy.units as u
import astropy.constants as c
import matplotlib.pyplot as plt
import matplotlib.markers as mrk
from radial import orbit

"""
This module contains the massive object classes that make up a star-planet
system.
"""


# The main star class
class MainStar(object):
    """
    The main star of a given system.

    mass : ``astropy.units.Quantity``
        The mass of the star.

    name : ``str`` or ``None``
        The name of the star. Default is ``None``.
    """
    def __init__(self, mass, name=None):
        self.mass = mass
        self.name = name

        # The radial velocities corresponding to each individual companion
        self.radial_v_ind = []

        # The total radial velocities
        self.radial_v = None


# The companion sub-class
class Companion(object):
    """
    The massive companion class. It can be either a binary star, an exoplanet,
    maybe even a black hole! It can be anything that has a mass and orbits
    another massive object. General relativity effects are not implemented yet.

    Parameters
    ----------
    k : ``astropy.units.Quantity`` or ``None``, optional
        The radial velocity semi-amplitude. Default is ``None``.

    period_orb : ``astropy.units.Quantity`` or ``None``, optional
        The orbital period. Default is ``None``.

    t_0 : ``astropy.units.Quantity`` or ``None``, optional
        The time of periastron passage. Default is ``None``.

    omega : ``astropy.units.Quantity`` or ``None``, optional
        Argument of periapse. Default is ``None``.

    ecc : ``float`` or ``None``, optional
        Eccentricity of the orbit. Default is ``None``.

    msini : ``astropy.units.Quantity`` or ``None``, optional
        Mass of the companion multiplied by sine of the inclination of the
        orbital plane in relation to the line of sight. Default is ``None``.

    semi_a : ``astropy.units.Quantity`` or ``None``, optional
        Semi-major axis of the orbit. Default is ``None``.

    name : ``str`` or ``None``, optional
        Name of the companion. Default is ``None``.
    """
    def __init__(self, k=None, period_orb=None, t_0=None, omega=None, ecc=None,
                 msini=None, semi_a=None, name=None, main_star=None):
        self.main_star = main_star
        self.msini = msini
        self.semi_a = semi_a
        self.k = k
        self.period_orb = period_orb
        self.t_0 = t_0
        self.omega = omega
        self.ecc = ecc
        self.name = name


# The star-companion system class
class System(object):
    """
    The star-companions system class.

    Parameters
    ----------
    main_star : ``radial.object.MainStar``
        The main star of the system.

    companions : list
        Python list containing the all the ``radial.object.Companion`` of the
        system.

    time : ``astropy.units.Quantity`` or ``None``
        A scalar or ``numpy.ndarray`` containing the times in which the radial
        velocities are measured. Default is ``None``.

    name : ``str`` or ``None``, optional
        Name of the system. Default is ``None``.
    """
    def __init__(self, main_star, companions, time=None, name=None):
        self.name = name
        self.main_star = main_star
        self.companions = companions
        self.n_c = len(self.companions)     # Number of companions
        self.time = time

        # Initializing useful global parameters
        self.f = []

    # Compute mass functions of the companions
    def mass_func(self):
        """
        Compute the mass functions of all the companions of the system. This
        method will also compute the msini and semi_a of the companions and
        save the values in their respective parameters.
        """
        for comp in self.companions:
            assert isinstance(comp.k, u.Quantity), 'k needs to be provided.'
            assert isinstance(comp.period_orb, u.Quantity), 'period_orb ' \
                                                            'needs to be ' \
                                                            'provided.'
            assert isinstance(comp.ecc, float), 'ecc needs to be provided.'
            f = (comp.period_orb * comp.k ** 3 * (1 - comp.ecc ** 2) ** (3 / 2)
                 / (2 * np.pi * c.G)).to(u.solMass)
            m_star = self.main_star.mass
            m_roots = np.roots([1,  -f.value,  -2 * m_star.value * f.value,
                                -m_star.value ** 2 * f.value])
            comp.msini = abs(m_roots[0]) * u.solMass
            comp.semi_a = (np.sqrt(c.G / comp.k * comp.msini * comp.period_orb /
                           (2 * np.pi * np.sqrt(1 - comp.ecc ** 2)))).to(u.AU)
            self.f.append(f)

    # Compute the radial velocities of the main star
    def compute_rv(self):
        """

        Parameters
        ----------
        time : ``astropy.units.Quantity``

        """
        ts = self.time.to(u.d).value
        star = self.main_star
        for comp in self.companions:
            subsystem = orbit.BinarySystem(k=comp.k.to(u.m / u.s).value,
                                           period=comp.period_orb.to(u.d).value,
                                           t0=comp.t_0.to(u.d).value,
                                           omega=comp.omega.to(u.rad).value,
                                           ecc=comp.ecc)
            star.radial_v_ind.append(subsystem.get_rvs(ts))
        star.radial_v_ind = np.array(star.radial_v_ind)

        # Combine the individual radial velocities to compose the total
        star.radial_v = np.sum(star.radial_v_ind, axis=0) * u.m / u.s
        star.radial_v_ind = star.radial_v_ind * u.m / u.s

    # Plot radial velocities of the main star
    def plot_rv(self, companion_index=None):
        """

        Parameters
        ----------
        companion_index : ``int`` or ``None``
            The companion index indicates which set of radial velocities will be
            plotted. If ``None``, then the total radial velocities are plotted.
            Default is ``None``.

        Returns
        -------

        """
        fig, ax = plt.subplots()
        star = self.main_star
        i = companion_index

        if i is None:
            ax.plot(self.time, star.radial_v)
            ax.set_xlabel('Time ({})'.format(str(self.time.unit)))
            ax.set_ylabel('Radial velocities ({})'.format(
                str(star.radial_v.unit)))
        else:
            ax.plot(self.time, star.radial_v_ind[i])
            ax.set_xlabel('Time ({})'.format(str(self.time.unit)))
            ax.set_ylabel('Radial velocities ({})'.format(
                str(star.radial_v_ind[i].unit)))

        return fig, ax
{'fraction_non_alphanumeric': 0.1147180434591361, 'fraction_numerical': 0.003430984428609132, 'mean_word_length': 1.9171157300487554, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '606990', 'n_tokens_mistral': 3352, 'n_tokens_neox': 3205, 'n_words': 1040}