17447393926

Committed 03 Sep 2025 10:11PM UTC coverage: 87.165%. First build

Build # 17447393926

Build Type

Pull #395

github

Committed by

bjfultn

Commit Message

Fix Python 3.11 compatibility: replace old Python 2 raise syntax and bare except clauses

Pull Request Pull Request #395: testing new ci

Run Details

170 of 241 new or added lines in 8 files covered. (70.54%)

3735 of 4285 relevant lines covered (87.16%)

0.87 hits per line

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

85.87

/radvel/kepler.py


import numpy as np
import radvel

# Try to import Kepler's equation solver written in C
def _check_cext():
    """Check if C extension is available - call this dynamically"""
    try:
        from . import _kepler
        return True
    except ImportError:
        return False

# Set default value, but we'll check dynamically in functions
cext = _check_cext()


def rv_drive(t, orbel, use_c_kepler_solver=cext):
    """RV Drive
    Args:
        t (array of floats): times of observations
        orbel (array of floats): [per, tp, e, om, K].\
            Omega is expected to be\
            in radians
        use_c_kepler_solver (bool): (default: True) If \
            True use the Kepler solver written in C, else \
            use the Python/NumPy version.
    Returns:
        rv: (array of floats): radial velocity model
    """

    # unpack array of parameters
    per, tp, e, om, k = orbel

    # Performance boost for circular orbits
    if e == 0.0:
        m = 2 * np.pi * (((t - tp) / per) - np.floor((t - tp) / per))
        return k * np.cos(m + om)

    if per < 0:
        per = 1e-4
    if e < 0:
        e = 0
    if e > 0.99:
        e = 0.99

    # Calculate the approximate eccentric anomaly, E1, via the mean anomaly  M.
    if use_c_kepler_solver and _check_cext():
        try:
            from . import _kepler
            rv = _kepler.rv_drive_array(t, per, tp, e, om, k)
        except (ImportError, NameError):
            # Fall back to pure NumPy implementation when C extension is unavailable
            nu = radvel.orbit.true_anomaly(t, tp, per, e)
            rv = k * (np.cos(nu + om) + e * np.cos(om))
    else:
        # Fall back to pure NumPy implementation when C extension is unavailable
        nu = radvel.orbit.true_anomaly(t, tp, per, e)
        rv = k * (np.cos(nu + om) + e * np.cos(om))

    return rv


def kepler(Marr, eccarr):
    """Solve Kepler's Equation
    Args:
        Marr (array): input Mean anomaly
        eccarr (array): eccentricity
    Returns:
        array: eccentric anomaly
    """

    conv = 1.0e-12  # convergence criterion
    k = 0.85

    Earr = Marr + np.sign(np.sin(Marr)) * k * eccarr  # first guess at E
    # fiarr should go to zero when converges
    fiarr = ( Earr - eccarr * np.sin(Earr) - Marr)
    convd = np.where(np.abs(fiarr) > conv)[0]  # which indices have not converged
    nd = len(convd)  # number of unconverged elements
    count = 0

    while nd > 0:  # while unconverged elements exist
        count += 1

        M = Marr[convd]  # just the unconverged elements ...
        ecc = eccarr[convd]
        E = Earr[convd]

        fi = fiarr[convd]  # fi = E - e*np.sin(E)-M    ; should go to 0
        fip = 1 - ecc * np.cos(E)  # d/dE(fi) ;i.e.,  fi^(prime)
        fipp = ecc * np.sin(E)  # d/dE(d/dE(fi)) ;i.e.,  fi^(\prime\prime)
        fippp = 1 - fip  # d/dE(d/dE(d/dE(fi))) ;i.e.,  fi^(\prime\prime\prime)

        # first, second, and third order corrections to E
        d1 = -fi / fip
        d2 = -fi / (fip + d1 * fipp / 2.0)
        d3 = -fi / (fip + d2 * fipp / 2.0 + d2 * d2 * fippp / 6.0)
        E = E + d3
        Earr[convd] = E
        fiarr = ( Earr - eccarr * np.sin( Earr ) - Marr) # how well did we do?
        convd = np.abs(fiarr) > conv  # test for convergence
        nd = np.sum(convd is True)

    if Earr.size > 1:
        return Earr
    else:
        return Earr[0]


def profile():
    # Profile and compare C-based Kepler solver with
    # Python/Numpy implementation

    import timeit

    ecc = 0.1
    numloops = 5000
    print("\nECCENTRICITY = {}".format(ecc))

    for size in [10, 30, 100, 300, 1000]:

        setup = """\
from radvel.kepler import rv_drive
import numpy as np
gc.enable()
ecc = %f
orbel = [32.468, 2456000, ecc, np.pi/2, 10.0]
t = np.linspace(2455000, 2457000, %d)
""" % (ecc, size)

        if _check_cext():
            print("\nProfiling pure C code for an RV time series with {} "
                  "observations".format(size))
            tc = timeit.timeit('rv_drive(t, orbel, use_c_kepler_solver=True)',
                               setup=setup, number=numloops)
            print("Ran %d model calculations in %5.3f seconds" % (numloops, tc))
        else:
            print("\nC extension not available; skipping C profiling for {} observations".format(size))
            tc = None

        print("Profiling Python code for an RV time series with {} "
              "observations".format(size))
        tp = timeit.timeit('rv_drive(t, orbel, use_c_kepler_solver=False)',
                           setup=setup, number=numloops)
        print("Ran %d model calculations in %5.3f seconds" % (numloops, tp))
        if tc is not None and tc > 0:
            print("The C version runs %5.2f times faster" % (tp/tc))

    ecc = 0.7
    numloops = 5000
    print("\nECCENTRICITY = {}".format(ecc))

    for size in [30]:
            setup = """\
from radvel.kepler import rv_drive
import numpy as np
gc.enable()
ecc = %f
orbel = [32.468, 2456000, ecc, np.pi/2, 10.0]
t = np.linspace(2455000, 2457000, %d)
    """ % (ecc, size)

            if _check_cext():
                print("\nProfiling pure C code for an RV time series with {} "
                      "observations".format(size))
                tc = timeit.timeit('rv_drive(t, orbel, use_c_kepler_solver=True)',
                                   setup=setup, number=numloops)
                print("Ran %d model calculations in %5.3f seconds" % (numloops, tc))
            else:
                print("\nC extension not available; skipping C profiling for {} observations".format(size))
                tc = None

            print("Profiling Python code for an RV time series with {} "
                  "observations".format(size))
            tp = timeit.timeit('rv_drive(t, orbel, use_c_kepler_solver=False)',
                               setup=setup, number=numloops)
            print("Ran %d model calculations in %5.3f seconds" % (numloops, tp))
            if tc is not None and tc > 0:
                print("The C version runs %5.2f times faster" % (tp / tc))

1
2	import numpy as np	1✔
3	import radvel	1✔
4
5	# Try to import Kepler's equation solver written in C
6	def _check_cext():	1✔
7	"""Check if C extension is available - call this dynamically"""
8	try:	1✔
9	from . import _kepler	1✔
10	return True	1✔
NEW 11	except ImportError:	×
NEW 12	return False	×
13
14	# Set default value, but we'll check dynamically in functions
15	cext = _check_cext()	1✔
16
17
18	def rv_drive(t, orbel, use_c_kepler_solver=cext):	1✔
19	"""RV Drive
20	Args:
21	t (array of floats): times of observations
22	orbel (array of floats): [per, tp, e, om, K].\
23	Omega is expected to be\
24	in radians
25	use_c_kepler_solver (bool): (default: True) If \
26	True use the Kepler solver written in C, else \
27	use the Python/NumPy version.
28	Returns:
29	rv: (array of floats): radial velocity model
30	"""
31
32	# unpack array of parameters
33	per, tp, e, om, k = orbel	1✔
34
35	# Performance boost for circular orbits
36	if e == 0.0:	1✔
37	m = 2 * np.pi * (((t - tp) / per) - np.floor((t - tp) / per))	1✔
38	return k * np.cos(m + om)	1✔
39
40	if per < 0:	1✔
41	per = 1e-4	×
42	if e < 0:	1✔
43	e = 0	×
44	if e > 0.99:	1✔
45	e = 0.99	×
46
47	# Calculate the approximate eccentric anomaly, E1, via the mean anomaly M.
48	if use_c_kepler_solver and _check_cext():	1✔
49	try:	1✔
50	from . import _kepler	1✔
51	rv = _kepler.rv_drive_array(t, per, tp, e, om, k)	1✔
NEW 52	except (ImportError, NameError):	×
53	# Fall back to pure NumPy implementation when C extension is unavailable
NEW 54	nu = radvel.orbit.true_anomaly(t, tp, per, e)	×
NEW 55	rv = k * (np.cos(nu + om) + e * np.cos(om))	×
56	else:
57	# Fall back to pure NumPy implementation when C extension is unavailable
58	nu = radvel.orbit.true_anomaly(t, tp, per, e)	1✔
59	rv = k * (np.cos(nu + om) + e * np.cos(om))	1✔
60
61	return rv	1✔
62
63
64	def kepler(Marr, eccarr):	1✔
65	"""Solve Kepler's Equation
66	Args:
67	Marr (array): input Mean anomaly
68	eccarr (array): eccentricity
69	Returns:
70	array: eccentric anomaly
71	"""
72
73	conv = 1.0e-12 # convergence criterion	1✔
74	k = 0.85	1✔
75
76	Earr = Marr + np.sign(np.sin(Marr)) * k * eccarr # first guess at E	1✔
77	# fiarr should go to zero when converges
78	fiarr = ( Earr - eccarr * np.sin(Earr) - Marr)	1✔
79	convd = np.where(np.abs(fiarr) > conv)[0] # which indices have not converged	1✔
80	nd = len(convd) # number of unconverged elements	1✔
81	count = 0	1✔
82
83	while nd > 0: # while unconverged elements exist	1✔
84	count += 1	1✔
85
86	M = Marr[convd] # just the unconverged elements ...	1✔
87	ecc = eccarr[convd]	1✔
88	E = Earr[convd]	1✔
89
90	fi = fiarr[convd] # fi = E - e*np.sin(E)-M ; should go to 0	1✔
91	fip = 1 - ecc * np.cos(E) # d/dE(fi) ;i.e., fi^(prime)	1✔
92	fipp = ecc * np.sin(E) # d/dE(d/dE(fi)) ;i.e., fi^(\prime\prime)	1✔
93	fippp = 1 - fip # d/dE(d/dE(d/dE(fi))) ;i.e., fi^(\prime\prime\prime)	1✔
94
95	# first, second, and third order corrections to E
96	d1 = -fi / fip	1✔
97	d2 = -fi / (fip + d1 * fipp / 2.0)	1✔
98	d3 = -fi / (fip + d2 * fipp / 2.0 + d2 * d2 * fippp / 6.0)	1✔
99	E = E + d3	1✔
100	Earr[convd] = E	1✔
101	fiarr = ( Earr - eccarr * np.sin( Earr ) - Marr) # how well did we do?	1✔
102	convd = np.abs(fiarr) > conv # test for convergence	1✔
103	nd = np.sum(convd is True)	1✔
104
105	if Earr.size > 1:	1✔
106	return Earr	1✔
107	else:
108	return Earr[0]	×
109
110
111	def profile():	1✔
112	# Profile and compare C-based Kepler solver with
113	# Python/Numpy implementation
114
115	import timeit	1✔
116
117	ecc = 0.1	1✔
118	numloops = 5000	1✔
119	print("\nECCENTRICITY = {}".format(ecc))	1✔
120
121	for size in [10, 30, 100, 300, 1000]:	1✔
122
123	setup = """\	1✔
124	from radvel.kepler import rv_drive
125	import numpy as np
126	gc.enable()
127	ecc = %f
128	orbel = [32.468, 2456000, ecc, np.pi/2, 10.0]
129	t = np.linspace(2455000, 2457000, %d)
130	""" % (ecc, size)
131
132	if _check_cext():	1✔
133	print("\nProfiling pure C code for an RV time series with {} "	1✔
134	"observations".format(size))
135	tc = timeit.timeit('rv_drive(t, orbel, use_c_kepler_solver=True)',	1✔
136	setup=setup, number=numloops)
137	print("Ran %d model calculations in %5.3f seconds" % (numloops, tc))	1✔
138	else:
NEW 139	print("\nC extension not available; skipping C profiling for {} observations".format(size))	×
NEW 140	tc = None	×
141
142	print("Profiling Python code for an RV time series with {} "	1✔
143	"observations".format(size))
144	tp = timeit.timeit('rv_drive(t, orbel, use_c_kepler_solver=False)',	1✔
145	setup=setup, number=numloops)
146	print("Ran %d model calculations in %5.3f seconds" % (numloops, tp))	1✔
147	if tc is not None and tc > 0:	1✔
148	print("The C version runs %5.2f times faster" % (tp/tc))	1✔
149
150	ecc = 0.7	1✔
151	numloops = 5000	1✔
152	print("\nECCENTRICITY = {}".format(ecc))	1✔
153
154	for size in [30]:	1✔
155	setup = """\	1✔
156	from radvel.kepler import rv_drive
157	import numpy as np
158	gc.enable()
159	ecc = %f
160	orbel = [32.468, 2456000, ecc, np.pi/2, 10.0]
161	t = np.linspace(2455000, 2457000, %d)
162	""" % (ecc, size)
163
164	if _check_cext():	1✔
165	print("\nProfiling pure C code for an RV time series with {} "	1✔
166	"observations".format(size))
167	tc = timeit.timeit('rv_drive(t, orbel, use_c_kepler_solver=True)',	1✔
168	setup=setup, number=numloops)
169	print("Ran %d model calculations in %5.3f seconds" % (numloops, tc))	1✔
170	else:
NEW 171	print("\nC extension not available; skipping C profiling for {} observations".format(size))	×
NEW 172	tc = None	×
173
174	print("Profiling Python code for an RV time series with {} "	1✔
175	"observations".format(size))
176	tp = timeit.timeit('rv_drive(t, orbel, use_c_kepler_solver=False)',	1✔
177	setup=setup, number=numloops)
178	print("Ran %d model calculations in %5.3f seconds" % (numloops, tp))	1✔
179	if tc is not None and tc > 0:	1✔
180	print("The C version runs %5.2f times faster" % (tp / tc))	1✔

California-Planet-Search / radvel / 17447393926

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