11979115482

Committed 22 Nov 2024 07:45PM UTC coverage: 65.499% (-0.2%) from 65.737%

Build # 11979115482

Build Type

Pull #403

github

Committed by

web-flow

Commit Message

Merge 0d3dd536f into 86cc76cd9

Pull Request Pull Request #403: Improved version tracking and deprecated SVN

Run Details

36 of 54 new or added lines in 2 files covered. (66.67%)

43 existing lines in 5 files now uncovered.

9534 of 14556 relevant lines covered (65.5%)

0.65 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

94.53

/EXOSIMS/PlanetPopulation/SAG13.py

from EXOSIMS.PlanetPopulation.KeplerLike2 import KeplerLike2
from EXOSIMS.util.InverseTransformSampler import InverseTransformSampler
import astropy.units as u
import astropy.constants as const
import numpy as np
import scipy.integrate as integrate
from EXOSIMS.util._numpy_compat import copy_if_needed


class SAG13(KeplerLike2):
    """Planet Population module based on SAG13 occurrence rates.

    This is the current working model based on averaging multiple studies.
    These do not yet represent official scientific values.

    Attributes:
        SAG13coeffs (float 4x2 ndarray):
            Coefficients used by the SAG13 broken power law. The 4 lines
            correspond to Gamma, alpha, beta, and the minimum radius.
        Gamma (float ndarray):
            Gamma coefficients used by SAG13 broken power law.
        alpha (float ndarray):
            Alpha coefficients used by SAG13 broken power law.
        beta (float ndarray):
            Beta coefficients used by SAG13 broken power law.
        Rplim (float ndarray):
            Minimum radius used by SAG13 broken power law.
        SAG13starMass (astropy Quantity):
            Assumed stellar mass corresponding to the given set of coefficients.
        mu (astropy Quantity):
            Gravitational parameter associated with SAG13starMass.
        Ca (float 2x1 ndarray):
            Constants used for sampling.

    """

    def __init__(
        self,
        SAG13coeffs=[[0.38, -0.19, 0.26, 0.0], [0.73, -1.18, 0.59, 3.4]],
        SAG13starMass=1.0,
        Rprange=[2 / 3.0, 17.0859375],
        arange=[0.09084645, 1.45354324],
        **specs
    ):

        # first initialize with KeplerLike constructor
        specs["Rprange"] = Rprange
        specs["arange"] = arange
        # load SAG13 star mass in solMass: 1.3 (F), 1 (G), 0.70 (K), 0.35 (M)
        self.SAG13starMass = float(SAG13starMass) * u.solMass
        self.mu = const.G * self.SAG13starMass

        # load SAG13 coefficients (Gamma, alpha, beta, Rplim)
        self.SAG13coeffs = np.array(SAG13coeffs, dtype=float)
        assert self.SAG13coeffs.ndim <= 2, "SAG13coeffs array dimension must be <= 2."
        # if only one row of coefficients, make sure the forth element
        # (minimum radius) is set to zero
        if self.SAG13coeffs.ndim == 1:
            self.SAG13coeffs = np.array(np.append(self.SAG13coeffs[:3], 0.0), ndmin=2)
        # make sure the array is of shape (4, n) where the forth row
        # contains the minimum radius values (broken power law)
        if len(self.SAG13coeffs) != 4:
            self.SAG13coeffs = self.SAG13coeffs.T
        assert len(self.SAG13coeffs) == 4, "SAG13coeffs array must have 4 rows."
        # sort by minimum radius
        self.SAG13coeffs = self.SAG13coeffs[:, np.argsort(self.SAG13coeffs[3, :])]

        # split out SAG13 coeffs
        self.Gamma = self.SAG13coeffs[0, :]
        self.alpha = self.SAG13coeffs[1, :]
        self.beta = self.SAG13coeffs[2, :]
        self.Rplim = np.append(self.SAG13coeffs[3, :], np.inf)

        KeplerLike2.__init__(self, **specs)

        # intermediate function
        m = self.mu.to("AU3/year2").value
        ftmp = (
            lambda x, b, m=m, ak=self.smaknee: (2.0 * np.pi * np.sqrt(x**3 / m))
            ** (b - 1.0)
            * (3.0 * np.pi * np.sqrt(x / m))
            * np.exp(-((x / ak) ** 3))
        )
        # intermediate constants used elsewhere
        self.Ca = np.zeros((2,))
        for i in range(2):
            self.Ca[i] = integrate.quad(
                ftmp,
                self.arange[0].to("AU").value,
                self.arange[1].to("AU").value,
                args=(self.beta[i],),
            )[0]

        # set up samplers for sma and Rp
        # probability density function of sma given Rp < Rplim[1]
        f_sma_given_Rp1 = lambda a, beta=self.beta[0], m=m, C=self.Ca[
            0
        ], smaknee=self.smaknee: self.dist_sma_given_radius(a, beta, m, C, smaknee)
        # sampler for Rp < Rplim:
        # unitless sma range
        ar = self.arange.to("AU").value
        self.sma_sampler1 = InverseTransformSampler(f_sma_given_Rp1, ar[0], ar[1])
        # probability density function of sma given Rp > Rplim[1]
        f_sma_given_Rp2 = lambda a, beta=self.beta[1], m=m, C=self.Ca[
            1
        ], smaknee=self.smaknee: self.dist_sma_given_radius(a, beta, m, C, smaknee)
        self.sma_sampler2 = InverseTransformSampler(f_sma_given_Rp2, ar[0], ar[1])

        self.Rp_sampler = InverseTransformSampler(
            self.dist_radius,
            self.Rprange[0].to("earthRad").value,
            self.Rprange[1].to("earthRad").value,
        )

        # determine eta
        if self.Rprange[1].to("earthRad").value < self.Rplim[1]:
            self.eta = (
                self.Gamma[0]
                * (
                    self.Rprange[1].to("earthRad").value ** self.alpha[0]
                    - self.Rprange[0].to("earthRad").value ** self.alpha[0]
                )
                / self.alpha[0]
                * self.Ca[0]
            )
        elif self.Rprange[0].to("earthRad").value > self.Rplim[1]:
            self.eta = (
                self.Gamma[1]
                * (
                    self.Rprange[1].to("earthRad").value ** self.alpha[1]
                    - self.Rprange[0].to("earthRad").value ** self.alpha[1]
                )
                / self.alpha[1]
                * self.Ca[1]
            )
        else:
            self.eta = (
                self.Gamma[0]
                * (
                    self.Rplim[1] ** self.alpha[0]
                    - self.Rprange[0].to("earthRad").value ** self.alpha[0]
                )
                / self.alpha[0]
                * self.Ca[0]
            )
            self.eta += (
                self.Gamma[1]
                * (
                    self.Rprange[1].to("earthRad").value ** self.alpha[1]
                    - self.Rplim[1] ** self.alpha[1]
                )
                / self.alpha[1]
                * self.Ca[1]
            )

        self._outspec["eta"] = self.eta

        # populate _outspec with SAG13 specific attributes
        self._outspec["SAG13starMass"] = self.SAG13starMass.to("solMass").value
        self._outspec["SAG13coeffs"] = self.SAG13coeffs

    def gen_radius_sma(self, n):
        """Generate radius values in earth radius and semi-major axis values in AU.

        This method is called by gen_radius and gen_sma.

        Args:
            n (integer):
                Number of samples to generate

        Returns:
            tuple:
            Rp (astropy Quantity array):
                Planet radius values in units of Earth radius
            a (astropy Quantity array):
                Semi-major axis values in units of AU

        """

        Rp = self.Rp_sampler(n)
        a = np.zeros(Rp.shape)
        a[Rp < self.Rplim[1]] = self.sma_sampler1(len(Rp[Rp < self.Rplim[1]]))
        a[Rp >= self.Rplim[1]] = self.sma_sampler2(len(Rp[Rp >= self.Rplim[1]]))
        Rp = Rp * u.earthRad
        a = a * u.AU

        return Rp, a

    def gen_plan_params(self, n):
        """Generate semi-major axis (AU), eccentricity, geometric albedo, and
        planetary radius (earthRad)

        Semi-major axis and planetary radius are jointly distributed.
        Eccentricity is a Rayleigh distribution. Albedo is dependent on the
        PlanetPhysicalModel but is calculated such that it is independent of
        other parameters.

        Args:
            n (integer):
                Number of samples to generate

        Returns:
            tuple:
            a (astropy Quantity array):
                Semi-major axis in units of AU
            e (float ndarray):
                Eccentricity
            p (float ndarray):
                Geometric albedo
            Rp (astropy Quantity array):
                Planetary radius in units of earthRad

        """
        n = self.gen_input_check(n)
        # generate semi-major axis and planetary radius samples
        Rp, a = self.gen_radius_sma(n)

        # check for constrainOrbits == True for eccentricity samples
        # constants
        C1 = np.exp(-self.erange[0] ** 2 / (2.0 * self.esigma**2))
        ar = self.arange.to("AU").value
        if self.constrainOrbits:
            # restrict semi-major axis limits
            arcon = np.array(
                [ar[0] / (1.0 - self.erange[0]), ar[1] / (1.0 + self.erange[0])]
            )
            # clip sma values to sma range
            sma = np.clip(a.to("AU").value, arcon[0], arcon[1])
            # upper limit for eccentricity given sma
            elim = np.zeros(len(sma))
            amean = np.mean(ar)
            elim[sma <= amean] = 1.0 - ar[0] / sma[sma <= amean]
            elim[sma > amean] = ar[1] / sma[sma > amean] - 1.0
            elim[elim > self.erange[1]] = self.erange[1]
            elim[elim < self.erange[0]] = self.erange[0]
            # additional constant
            C2 = C1 - np.exp(-(elim**2) / (2.0 * self.esigma**2))
            a = sma * u.AU
        else:
            C2 = self.enorm
        e = self.esigma * np.sqrt(-2.0 * np.log(C1 - C2 * np.random.uniform(size=n)))
        # generate albedo from semi-major axis
        p = self.PlanetPhysicalModel.calc_albedo_from_sma(a, self.prange)

        return a, e, p, Rp

    def dist_sma_radius(self, a, R):
        """Joint probability density function for semi-major axis (AU) and
        planetary radius in Earth radius.

        This method performs a change of variables on the SAG13 broken power
        law (originally in planetary radius and period).

        Args:
            a (float ndarray):
                Semi-major axis values in AU. Not an astropy quantity
            R (float ndarray):
                Planetary radius values in Earth radius. Not an astropy quantity

        Returns:
            float ndarray:
                Joint (semi-major axis and planetary radius) probability density
                matrix of shape (len(R),len(a))

        """
        # cast to arrays
        a = np.array(a, ndmin=1, copy=copy_if_needed)
        R = np.array(R, ndmin=1, copy=copy_if_needed)

        assert (
            a.shape == R.shape
        ), "input semi-major axis and planetary radius must have same shape"

        mu = self.mu.to("AU3/year2").value

        f = np.zeros(a.shape)
        mask1 = (
            (R < self.Rplim[1])
            & (R > self.Rprange[0].value)
            & (R < self.Rprange[1].value)
            & (a > self.arange[0].value)
            & (a < self.arange[1].value)
        )
        mask2 = (
            (R > self.Rplim[1])
            & (R > self.Rprange[0].value)
            & (R < self.Rprange[1].value)
            & (a > self.arange[0].value)
            & (a < self.arange[1].value)
        )

        # for R < boundary radius
        f[mask1] = self.Gamma[0] * R[mask1] ** (self.alpha[0] - 1.0)
        f[mask1] *= (2.0 * np.pi * np.sqrt(a[mask1] ** 3 / mu)) ** (self.beta[0] - 1.0)
        f[mask1] *= (3.0 * np.pi * np.sqrt(a[mask1] / mu)) * np.exp(
            -((a[mask1] / self.smaknee) ** 3)
        )
        f[mask1] /= self.eta

        # for R > boundary radius
        f[mask2] = self.Gamma[1] * R[mask2] ** (self.alpha[1] - 1.0)
        f[mask2] *= (2.0 * np.pi * np.sqrt(a[mask2] ** 3 / mu)) ** (self.beta[1] - 1.0)
        f[mask2] *= (3.0 * np.pi * np.sqrt(a[mask2] / mu)) * np.exp(
            -((a[mask2] / self.smaknee) ** 3)
        )
        f[mask2] /= self.eta

        return f

    def dist_sma(self, a):
        """Marginalized probability density function for semi-major axis in AU.

        Args:
            a (float ndarray):
                Semi-major axis value(s) in AU. Not an astropy quantity.

        Returns:
            float ndarray:
                Semi-major axis probability density

        """
        # cast to array
        a = np.array(a, ndmin=1, copy=copy_if_needed)
        # unitless sma range
        ar = self.arange.to("AU").value
        mu = self.mu.to("AU3/year2").value
        f = np.zeros(a.shape)
        mask = np.array((a >= ar[0]) & (a <= ar[1]), ndmin=1)

        Rmin = self.Rprange[0].to("earthRad").value
        Rmax = self.Rprange[1].to("earthRad").value

        f[mask] = (3.0 * np.pi * np.sqrt(a[mask] / mu)) * np.exp(
            -((a[mask] / self.smaknee) ** 3)
        )

        if Rmin < self.Rplim[1] and Rmax < self.Rplim[1]:
            C1 = (
                self.Gamma[0]
                * (Rmax ** self.alpha[0] - Rmin ** self.alpha[0])
                / self.alpha[0]
            )
            f[mask] *= C1 * (2.0 * np.pi * np.sqrt(a[mask] ** 3 / mu)) ** (
                self.beta[0] - 1.0
            )
        elif Rmin > self.Rplim[1] and Rmax > self.Rplim[1]:
            C2 = (
                self.Gamma[1]
                * (Rmax ** self.alpha[1] - Rmin ** self.alpha[1])
                / self.alpha[1]
            )
            f[mask] *= C2 * (2.0 * np.pi * np.sqrt(a[mask] ** 3 / mu)) ** (
                self.beta[1] - 1.0
            )
        else:
            C1 = (
                self.Gamma[0]
                * (self.Rplim[1] ** self.alpha[0] - Rmin ** self.alpha[0])
                / self.alpha[0]
            )
            C2 = (
                self.Gamma[1]
                * (Rmax ** self.alpha[1] - self.Rplim[1] ** self.alpha[1])
                / self.alpha[1]
            )
            f[mask] *= C1 * (2.0 * np.pi * np.sqrt(a[mask] ** 3 / mu)) ** (
                self.beta[0] - 1.0
            ) + C2 * (2.0 * np.pi * np.sqrt(a[mask] ** 3 / mu)) ** (self.beta[1] - 1.0)

        f /= self.eta

        return f

    def dist_radius(self, Rp):
        """Marginalized probability density function for planetary radius in
        Earth radius.

        Args:
            Rp (float ndarray):
                Planetary radius value(s) in Earth radius. Not an astropy quantity.

        Returns:
            float ndarray:
                Planetary radius probability density

        """

        # cast Rp to array
        Rp = np.array(Rp, ndmin=1, copy=copy_if_needed)
        f = np.zeros(Rp.shape)
        # unitless Rp range
        Rr = self.Rprange.to("earthRad").value

        mask1 = np.array((Rp >= Rr[0]) & (Rp <= self.Rplim[1]), ndmin=1)
        mask2 = np.array((Rp >= self.Rplim[1]) & (Rp <= Rr[1]), ndmin=1)

        masks = [mask1, mask2]
        for i in range(2):
            f[masks[i]] = (
                self.Gamma[i]
                * Rp[masks[i]] ** (self.alpha[i] - 1.0)
                * self.Ca[i]
                / self.eta
            )

        return f

    def dist_sma_given_radius(self, a, beta, m, C, smaknee):
        """Conditional probability density function of semi-major axis given
        planetary radius.

        Args:
            a (float ndarray):
                Semi-major axis value(s) in AU. Not an astropy quantity
            beta (float):
                Exponent for distribution
            m (float):
                Gravitational parameter (AU3/year2)
            C (float):
                Normalization for distribution
            smaknee (float):
                Coefficient for decay

        Returns:
            float ndarray:
                Probability density

        """
        # cast a to array
        a = np.array(a, ndmin=1, copy=copy_if_needed)
        ar = self.arange.to("AU").value
        mask = np.array((a >= ar[0]) & (a <= ar[1]), ndmin=1)
        f = np.zeros(a.shape)
        f[mask] = (
            (2.0 * np.pi * np.sqrt(a[mask] ** 3 / m)) ** (beta - 1.0)
            * (3.0 * np.pi * np.sqrt(a[mask] / m))
            * np.exp(-((a / smaknee) ** 3))
            / C
        )

        return f

1	from EXOSIMS.PlanetPopulation.KeplerLike2 import KeplerLike2	1✔
2	from EXOSIMS.util.InverseTransformSampler import InverseTransformSampler	1✔
3	import astropy.units as u	1✔
4	import astropy.constants as const	1✔
5	import numpy as np	1✔
6	import scipy.integrate as integrate	1✔
7	from EXOSIMS.util._numpy_compat import copy_if_needed	1✔
8
9
10	class SAG13(KeplerLike2):	1✔
11	"""Planet Population module based on SAG13 occurrence rates.
12
13	This is the current working model based on averaging multiple studies.
14	These do not yet represent official scientific values.
15
16	Attributes:
17	SAG13coeffs (float 4x2 ndarray):
18	Coefficients used by the SAG13 broken power law. The 4 lines
19	correspond to Gamma, alpha, beta, and the minimum radius.
20	Gamma (float ndarray):
21	Gamma coefficients used by SAG13 broken power law.
22	alpha (float ndarray):
23	Alpha coefficients used by SAG13 broken power law.
24	beta (float ndarray):
25	Beta coefficients used by SAG13 broken power law.
26	Rplim (float ndarray):
27	Minimum radius used by SAG13 broken power law.
28	SAG13starMass (astropy Quantity):
29	Assumed stellar mass corresponding to the given set of coefficients.
30	mu (astropy Quantity):
31	Gravitational parameter associated with SAG13starMass.
32	Ca (float 2x1 ndarray):
33	Constants used for sampling.
34
35	"""
36
37	def __init__(	1✔
38	self,
39	SAG13coeffs=[[0.38, -0.19, 0.26, 0.0], [0.73, -1.18, 0.59, 3.4]],
40	SAG13starMass=1.0,
41	Rprange=[2 / 3.0, 17.0859375],
42	arange=[0.09084645, 1.45354324],
43	**specs
44	):
45
46	# first initialize with KeplerLike constructor
47	specs["Rprange"] = Rprange	1✔
48	specs["arange"] = arange	1✔
49	# load SAG13 star mass in solMass: 1.3 (F), 1 (G), 0.70 (K), 0.35 (M)
50	self.SAG13starMass = float(SAG13starMass) * u.solMass	1✔
51	self.mu = const.G * self.SAG13starMass	1✔
52
53	# load SAG13 coefficients (Gamma, alpha, beta, Rplim)
54	self.SAG13coeffs = np.array(SAG13coeffs, dtype=float)	1✔
55	assert self.SAG13coeffs.ndim <= 2, "SAG13coeffs array dimension must be <= 2."	1✔
56	# if only one row of coefficients, make sure the forth element
57	# (minimum radius) is set to zero
58	if self.SAG13coeffs.ndim == 1:	1✔
59	self.SAG13coeffs = np.array(np.append(self.SAG13coeffs[:3], 0.0), ndmin=2)	×
60	# make sure the array is of shape (4, n) where the forth row
61	# contains the minimum radius values (broken power law)
62	if len(self.SAG13coeffs) != 4:	1✔
63	self.SAG13coeffs = self.SAG13coeffs.T	1✔
64	assert len(self.SAG13coeffs) == 4, "SAG13coeffs array must have 4 rows."	1✔
65	# sort by minimum radius
66	self.SAG13coeffs = self.SAG13coeffs[:, np.argsort(self.SAG13coeffs[3, :])]	1✔
67
68	# split out SAG13 coeffs
69	self.Gamma = self.SAG13coeffs[0, :]	1✔
70	self.alpha = self.SAG13coeffs[1, :]	1✔
71	self.beta = self.SAG13coeffs[2, :]	1✔
72	self.Rplim = np.append(self.SAG13coeffs[3, :], np.inf)	1✔
73
74	KeplerLike2.__init__(self, **specs)	1✔
75
76	# intermediate function
77	m = self.mu.to("AU3/year2").value	1✔
78	ftmp = (	1✔
79	lambda x, b, m=m, ak=self.smaknee: (2.0 * np.pi * np.sqrt(x**3 / m))
80	** (b - 1.0)
81	* (3.0 * np.pi * np.sqrt(x / m))
82	* np.exp(-((x / ak) ** 3))
83	)
84	# intermediate constants used elsewhere
85	self.Ca = np.zeros((2,))	1✔
86	for i in range(2):	1✔
87	self.Ca[i] = integrate.quad(	1✔
88	ftmp,
89	self.arange[0].to("AU").value,
90	self.arange[1].to("AU").value,
91	args=(self.beta[i],),
92	)[0]
93
94	# set up samplers for sma and Rp
95	# probability density function of sma given Rp < Rplim[1]
96	f_sma_given_Rp1 = lambda a, beta=self.beta[0], m=m, C=self.Ca[	1✔
97	0
98	], smaknee=self.smaknee: self.dist_sma_given_radius(a, beta, m, C, smaknee)
99	# sampler for Rp < Rplim:
100	# unitless sma range
101	ar = self.arange.to("AU").value	1✔
102	self.sma_sampler1 = InverseTransformSampler(f_sma_given_Rp1, ar[0], ar[1])	1✔
103	# probability density function of sma given Rp > Rplim[1]
104	f_sma_given_Rp2 = lambda a, beta=self.beta[1], m=m, C=self.Ca[	1✔
105	1
106	], smaknee=self.smaknee: self.dist_sma_given_radius(a, beta, m, C, smaknee)
107	self.sma_sampler2 = InverseTransformSampler(f_sma_given_Rp2, ar[0], ar[1])	1✔
108
109	self.Rp_sampler = InverseTransformSampler(	1✔
110	self.dist_radius,
111	self.Rprange[0].to("earthRad").value,
112	self.Rprange[1].to("earthRad").value,
113	)
114
115	# determine eta
116	if self.Rprange[1].to("earthRad").value < self.Rplim[1]:	1✔
117	self.eta = (	×
118	self.Gamma[0]
119	* (
120	self.Rprange[1].to("earthRad").value ** self.alpha[0]
121	- self.Rprange[0].to("earthRad").value ** self.alpha[0]
122	)
123	/ self.alpha[0]
124	* self.Ca[0]
125	)
126	elif self.Rprange[0].to("earthRad").value > self.Rplim[1]:	1✔
UNCOV 127	self.eta = (	×
128	self.Gamma[1]
129	* (
130	self.Rprange[1].to("earthRad").value ** self.alpha[1]
131	- self.Rprange[0].to("earthRad").value ** self.alpha[1]
132	)
133	/ self.alpha[1]
134	* self.Ca[1]
135	)
136	else:
137	self.eta = (	1✔
138	self.Gamma[0]
139	* (
140	self.Rplim[1] ** self.alpha[0]
141	- self.Rprange[0].to("earthRad").value ** self.alpha[0]
142	)
143	/ self.alpha[0]
144	* self.Ca[0]
145	)
146	self.eta += (	1✔
147	self.Gamma[1]
148	* (
149	self.Rprange[1].to("earthRad").value ** self.alpha[1]
150	- self.Rplim[1] ** self.alpha[1]
151	)
152	/ self.alpha[1]
153	* self.Ca[1]
154	)
155
156	self._outspec["eta"] = self.eta	1✔
157
158	# populate _outspec with SAG13 specific attributes
159	self._outspec["SAG13starMass"] = self.SAG13starMass.to("solMass").value	1✔
160	self._outspec["SAG13coeffs"] = self.SAG13coeffs	1✔
161
162	def gen_radius_sma(self, n):	1✔
163	"""Generate radius values in earth radius and semi-major axis values in AU.
164
165	This method is called by gen_radius and gen_sma.
166
167	Args:
168	n (integer):
169	Number of samples to generate
170
171	Returns:
172	tuple:
173	Rp (astropy Quantity array):
174	Planet radius values in units of Earth radius
175	a (astropy Quantity array):
176	Semi-major axis values in units of AU
177
178	"""
179
180	Rp = self.Rp_sampler(n)	1✔
181	a = np.zeros(Rp.shape)	1✔
182	a[Rp < self.Rplim[1]] = self.sma_sampler1(len(Rp[Rp < self.Rplim[1]]))	1✔
183	a[Rp >= self.Rplim[1]] = self.sma_sampler2(len(Rp[Rp >= self.Rplim[1]]))	1✔
184	Rp = Rp * u.earthRad	1✔
185	a = a * u.AU	1✔
186
187	return Rp, a	1✔
188
189	def gen_plan_params(self, n):	1✔
190	"""Generate semi-major axis (AU), eccentricity, geometric albedo, and
191	planetary radius (earthRad)
192
193	Semi-major axis and planetary radius are jointly distributed.
194	Eccentricity is a Rayleigh distribution. Albedo is dependent on the
195	PlanetPhysicalModel but is calculated such that it is independent of
196	other parameters.
197
198	Args:
199	n (integer):
200	Number of samples to generate
201
202	Returns:
203	tuple:
204	a (astropy Quantity array):
205	Semi-major axis in units of AU
206	e (float ndarray):
207	Eccentricity
208	p (float ndarray):
209	Geometric albedo
210	Rp (astropy Quantity array):
211	Planetary radius in units of earthRad
212
213	"""
214	n = self.gen_input_check(n)	1✔
215	# generate semi-major axis and planetary radius samples
216	Rp, a = self.gen_radius_sma(n)	1✔
217
218	# check for constrainOrbits == True for eccentricity samples
219	# constants
220	C1 = np.exp(-self.erange[0] ** 2 / (2.0 * self.esigma**2))	1✔
221	ar = self.arange.to("AU").value	1✔
222	if self.constrainOrbits:	1✔
223	# restrict semi-major axis limits
224	arcon = np.array(	1✔
225	[ar[0] / (1.0 - self.erange[0]), ar[1] / (1.0 + self.erange[0])]
226	)
227	# clip sma values to sma range
228	sma = np.clip(a.to("AU").value, arcon[0], arcon[1])	1✔
229	# upper limit for eccentricity given sma
230	elim = np.zeros(len(sma))	1✔
231	amean = np.mean(ar)	1✔
232	elim[sma <= amean] = 1.0 - ar[0] / sma[sma <= amean]	1✔
233	elim[sma > amean] = ar[1] / sma[sma > amean] - 1.0	1✔
234	elim[elim > self.erange[1]] = self.erange[1]	1✔
235	elim[elim < self.erange[0]] = self.erange[0]	1✔
236	# additional constant
237	C2 = C1 - np.exp(-(elim*2) / (2.0 self.esigma**2))	1✔
238	a = sma * u.AU	1✔
239	else:
240	C2 = self.enorm	1✔
241	e = self.esigma * np.sqrt(-2.0 * np.log(C1 - C2 * np.random.uniform(size=n)))	1✔
242	# generate albedo from semi-major axis
243	p = self.PlanetPhysicalModel.calc_albedo_from_sma(a, self.prange)	1✔
244
245	return a, e, p, Rp	1✔
246
247	def dist_sma_radius(self, a, R):	1✔
248	"""Joint probability density function for semi-major axis (AU) and
249	planetary radius in Earth radius.
250
251	This method performs a change of variables on the SAG13 broken power
252	law (originally in planetary radius and period).
253
254	Args:
255	a (float ndarray):
256	Semi-major axis values in AU. Not an astropy quantity
257	R (float ndarray):
258	Planetary radius values in Earth radius. Not an astropy quantity
259
260	Returns:
261	float ndarray:
262	Joint (semi-major axis and planetary radius) probability density
263	matrix of shape (len(R),len(a))
264
265	"""
266	# cast to arrays
267	a = np.array(a, ndmin=1, copy=copy_if_needed)	1✔
268	R = np.array(R, ndmin=1, copy=copy_if_needed)	1✔
269
270	assert (	1✔
271	a.shape == R.shape
272	), "input semi-major axis and planetary radius must have same shape"
273
274	mu = self.mu.to("AU3/year2").value	1✔
275
276	f = np.zeros(a.shape)	1✔
277	mask1 = (	1✔
278	(R < self.Rplim[1])
279	& (R > self.Rprange[0].value)
280	& (R < self.Rprange[1].value)
281	& (a > self.arange[0].value)
282	& (a < self.arange[1].value)
283	)
284	mask2 = (	1✔
285	(R > self.Rplim[1])
286	& (R > self.Rprange[0].value)
287	& (R < self.Rprange[1].value)
288	& (a > self.arange[0].value)
289	& (a < self.arange[1].value)
290	)
291
292	# for R < boundary radius
293	f[mask1] = self.Gamma[0] * R[mask1] ** (self.alpha[0] - 1.0)	1✔
294	f[mask1] = (2.0 np.pi * np.sqrt(a[mask1] 3 / mu)) (self.beta[0] - 1.0)	1✔
295	f[mask1] = (3.0 np.pi * np.sqrt(a[mask1] / mu)) * np.exp(	1✔
296	-((a[mask1] / self.smaknee) ** 3)
297	)
298	f[mask1] /= self.eta	1✔
299
300	# for R > boundary radius
301	f[mask2] = self.Gamma[1] * R[mask2] ** (self.alpha[1] - 1.0)	1✔
302	f[mask2] = (2.0 np.pi * np.sqrt(a[mask2] 3 / mu)) (self.beta[1] - 1.0)	1✔
303	f[mask2] = (3.0 np.pi * np.sqrt(a[mask2] / mu)) * np.exp(	1✔
304	-((a[mask2] / self.smaknee) ** 3)
305	)
306	f[mask2] /= self.eta	1✔
307
308	return f	1✔
309
310	def dist_sma(self, a):	1✔
311	"""Marginalized probability density function for semi-major axis in AU.
312
313	Args:
314	a (float ndarray):
315	Semi-major axis value(s) in AU. Not an astropy quantity.
316
317	Returns:
318	float ndarray:
319	Semi-major axis probability density
320
321	"""
322	# cast to array
323	a = np.array(a, ndmin=1, copy=copy_if_needed)	1✔
324	# unitless sma range
325	ar = self.arange.to("AU").value	1✔
326	mu = self.mu.to("AU3/year2").value	1✔
327	f = np.zeros(a.shape)	1✔
328	mask = np.array((a >= ar[0]) & (a <= ar[1]), ndmin=1)	1✔
329
330	Rmin = self.Rprange[0].to("earthRad").value	1✔
331	Rmax = self.Rprange[1].to("earthRad").value	1✔
332
333	f[mask] = (3.0 * np.pi * np.sqrt(a[mask] / mu)) * np.exp(	1✔
334	-((a[mask] / self.smaknee) ** 3)
335	)
336
337	if Rmin < self.Rplim[1] and Rmax < self.Rplim[1]:	1✔
338	C1 = (	×
339	self.Gamma[0]
340	* (Rmax self.alpha[0] - Rmin self.alpha[0])
341	/ self.alpha[0]
342	)
343	f[mask] = C1 (2.0 * np.pi * np.sqrt(a[mask] 3 / mu)) (	×
344	self.beta[0] - 1.0
345	)
346	elif Rmin > self.Rplim[1] and Rmax > self.Rplim[1]:	1✔
UNCOV 347	C2 = (	×
348	self.Gamma[1]
349	* (Rmax self.alpha[1] - Rmin self.alpha[1])
350	/ self.alpha[1]
351	)
UNCOV 352	f[mask] = C2 (2.0 * np.pi * np.sqrt(a[mask] 3 / mu)) (	×
353	self.beta[1] - 1.0
354	)
355	else:
356	C1 = (	1✔
357	self.Gamma[0]
358	* (self.Rplim[1] self.alpha[0] - Rmin self.alpha[0])
359	/ self.alpha[0]
360	)
361	C2 = (	1✔
362	self.Gamma[1]
363	* (Rmax self.alpha[1] - self.Rplim[1] self.alpha[1])
364	/ self.alpha[1]
365	)
366	f[mask] = C1 (2.0 * np.pi * np.sqrt(a[mask] 3 / mu)) (	1✔
367	self.beta[0] - 1.0
368	) + C2 * (2.0 * np.pi * np.sqrt(a[mask] 3 / mu)) (self.beta[1] - 1.0)
369
370	f /= self.eta	1✔
371
372	return f	1✔
373
374	def dist_radius(self, Rp):	1✔
375	"""Marginalized probability density function for planetary radius in
376	Earth radius.
377
378	Args:
379	Rp (float ndarray):
380	Planetary radius value(s) in Earth radius. Not an astropy quantity.
381
382	Returns:
383	float ndarray:
384	Planetary radius probability density
385
386	"""
387
388	# cast Rp to array
389	Rp = np.array(Rp, ndmin=1, copy=copy_if_needed)	1✔
390	f = np.zeros(Rp.shape)	1✔
391	# unitless Rp range
392	Rr = self.Rprange.to("earthRad").value	1✔
393
394	mask1 = np.array((Rp >= Rr[0]) & (Rp <= self.Rplim[1]), ndmin=1)	1✔
395	mask2 = np.array((Rp >= self.Rplim[1]) & (Rp <= Rr[1]), ndmin=1)	1✔
396
397	masks = [mask1, mask2]	1✔
398	for i in range(2):	1✔
399	f[masks[i]] = (	1✔
400	self.Gamma[i]
401	* Rp[masks[i]] ** (self.alpha[i] - 1.0)
402	* self.Ca[i]
403	/ self.eta
404	)
405
406	return f	1✔
407
408	def dist_sma_given_radius(self, a, beta, m, C, smaknee):	1✔
409	"""Conditional probability density function of semi-major axis given
410	planetary radius.
411
412	Args:
413	a (float ndarray):
414	Semi-major axis value(s) in AU. Not an astropy quantity
415	beta (float):
416	Exponent for distribution
417	m (float):
418	Gravitational parameter (AU3/year2)
419	C (float):
420	Normalization for distribution
421	smaknee (float):
422	Coefficient for decay
423
424	Returns:
425	float ndarray:
426	Probability density
427
428	"""
429	# cast a to array
430	a = np.array(a, ndmin=1, copy=copy_if_needed)	1✔
431	ar = self.arange.to("AU").value	1✔
432	mask = np.array((a >= ar[0]) & (a <= ar[1]), ndmin=1)	1✔
433	f = np.zeros(a.shape)	1✔
434	f[mask] = (	1✔
435	(2.0 * np.pi * np.sqrt(a[mask] 3 / m)) (beta - 1.0)
436	* (3.0 * np.pi * np.sqrt(a[mask] / m))
437	* np.exp(-((a / smaknee) ** 3))
438	/ C
439	)
440
441	return f	1✔

dsavransky / EXOSIMS / 11979115482

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous