Starting to get rid of fortran code and replace it with numba Python.

setup.py

@@ -1,25 +1,19 @@
-import sys
-from numpy.distutils.core import setup, Extension
+from setuptools import setup, find_packages
 
-e_gimenez = Extension('pytransit.gimenez_f', ['src/gimenez.f90'], libraries=['gomp', 'm'], define_macros=[('DCHUNK_SIZE', 128)])
-e_mandelagol = Extension('pytransit.mandelagol_f', ['src/mandelagol.f90'], libraries=['gomp', 'm'])
-e_utils = Extension('pytransit.utils_f', ['src/utils.f90'], libraries=['gomp', 'm'])
-e_orbits = Extension('pytransit.orbits_f',  ['src/orbits.f90','src/orbits.pyf'], libraries=['gomp','m'])
-
-version = '1.5'
+version = '2.0'
 
 setup(name='PyTransit',
       version=version,
       description='Fast and painless exoplanet transit light curve modelling.',
       author='Hannu Parviainen',
       author_email='[email protected]',
       url='https://github.com/hpparvi/PyTransit',
-      extra_options = ['-fopenmp'],
       package_dir={'pytransit':'src'},
+      #packages=find_packages(),
       packages=['pytransit','pytransit.utils','pytransit.param', 'pytransit.contamination'],
       package_data={'':['*.cl'], 'pytransit.contamination':['data/*']},
-      ext_modules=[e_gimenez, e_mandelagol, e_utils, e_orbits],
-      install_requires=["numpy", "scipy", "pandas", "xarray", "tables"],
+      install_requires=["numpy", "numba", "scipy", "pandas", "xarray", "tables"],
+      include_package_data=True,
       license='GPLv2',
       classifiers=[
           "Topic :: Scientific/Engineering",
@@ -31,5 +25,6 @@
           "Programming Language :: Python",
           "Programming Language :: Fortran",
           "Programming Language :: Other"
-      ]
+      ],
+      keywords='astronomy astrophysics exoplanets'
       )

src/gimenez.py

@@ -17,8 +17,8 @@
 
 import numpy as np
 from math import fabs
-from .gimenez_f import gimenez as g
-from .orbits_f import orbits as of
+#from .gimenez_f import gimenez as g
+#from .orbits_f import orbits as of
 from .tm import TransitModel
 
 class Gimenez(TransitModel):

src/mandelagol.py

@@ -16,8 +16,8 @@
 ## 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 
 import numpy as np
-from .mandelagol_f import mandelagol as ma
-from .orbits_f import orbits as of
+#from .mandelagol_f import mandelagol as ma
+#from .orbits_f import orbits as of
 from .tm import TransitModel
 
 class MandelAgol(TransitModel):

src/mandelagol_py.py

@@ -0,0 +1,256 @@
+from numpy import pi, sqrt, arccos, abs, log, ones_like, zeros, zeros_like
+from numba import jit, prange
+
+HALF_PI = 0.5 * pi
+FOUR_PI = 4.0 * pi
+INV_PI = 1 / pi
+
+@jit(["f4(f4,f4)", "f8(f8,f8)"], cache=True, nopython=True)
+def ellpicb(n, k):
+    """The complete elliptical integral of the third kind
+
+    Bulirsch 1965, Numerische Mathematik, 7, 78
+    Bulirsch 1965, Numerische Mathematik, 7, 353
+
+    Adapted from L. Kreidbergs C version in BATMAN
+    (Kreidberg, L. 2015, PASP 957, 127)
+    (https://github.com/lkreidberg/batman)
+    which is translated from J. Eastman's IDL routine
+    in EXOFAST (Eastman et al. 2013, PASP 125, 83)"""
+
+    kc = sqrt(1.0 - k * k)
+    e = kc
+    p = sqrt(n + 1.0)
+    ip = 1.0 / p
+    m0 = 1.0
+    c = 1.0
+    d = 1.0 / p
+
+    for nit in range(1000):
+        f = c
+        c = d / p + c
+        g = e / p
+        d = 2.0 * (f * g + d)
+        p = g + p
+        g = m0
+        m0 = kc + m0
+
+        if (abs(1.0 - kc / g) > 1e-8):
+            kc = 2.0 * sqrt(e)
+            e = kc * m0
+        else:
+            return HALF_PI * (c * m0 + d) / (m0 * (m0 + p))
+    return 0.0
+
+
+@jit(cache=True, nopython=True)
+def ellec(k):
+    a1 = 0.443251414630
+    a2 = 0.062606012200
+    a3 = 0.047573835460
+    a4 = 0.017365064510
+    b1 = 0.249983683100
+    b2 = 0.092001800370
+    b3 = 0.040696975260
+    b4 = 0.005264496390
+
+    m1 = 1.0 - k * k
+    ee1 = 1.0 + m1 * (a1 + m1 * (a2 + m1 * (a3 + m1 * a4)))
+    ee2 = m1 * (b1 + m1 * (b2 + m1 * (b3 + m1 * b4))) * log(1.0 / m1)
+    return ee1 + ee2
+
+
+@jit(cache=True, nopython=True)
+def ellk(k):
+    a0 = 1.386294361120
+    a1 = 0.096663442590
+    a2 = 0.035900923830
+    a3 = 0.037425637130
+    a4 = 0.014511962120
+    b0 = 0.50
+    b1 = 0.124985935970
+    b2 = 0.068802485760
+    b3 = 0.033283553460
+    b4 = 0.004417870120
+
+    m1 = 1.0 - k * k
+    ek1 = a0 + m1 * (a1 + m1 * (a2 + m1 * (a3 + m1 * a4)))
+    ek2 = (b0 + m1 * (b1 + m1 * (b2 + m1 * (b3 + m1 * b4)))) * log(m1)
+    return ek1 - ek2
+
+@jit(["f4[:](f4[:],f4,f4)", "f8[:](f8[:],f8,f8)"], cache=True, nopython=True)
+def eval_uniform(z0, k, c):
+    flux = zeros_like(z0)
+
+    if abs(k - 0.5) < 1e-3:
+        k = 0.5
+
+    for i in range(len(z0)):
+        z = z0[i]
+        if z < 0.0 or z > 1.0 + k:
+            flux[i] = 1.0
+        elif k > 1.0 and z < k - 1.0:
+            flux[i] = 0.0
+        elif (z > abs(1.0 - k) and z < 1.0 + k):
+            kap1 = arccos(min((1.0 - k * k + z * z) / 2.0 / z, 1.0))
+            kap0 = arccos(min((k * k + z * z - 1.0) / 2.0 / k / z, 1.0))
+            lambdae = k * k * kap0 + kap1
+            lambdae = (lambdae - 0.5 * sqrt(max(4.0 * z * z - (1.0 + z * z - k * k) ** 2, 0.0))) / pi
+            flux[i] = 1.0 - lambdae
+        elif (z < 1.0 - k):
+            flux[i] = 1.0 - k * k
+        flux[i] = c + (1.0 - c) * flux[i]
+    return flux
+
+
+@jit("Tuple((f8[:,:], f8[:], f8[:], f8[:]))(f8[:], f8, f8[:], f8)", cache=True, nopython=True, parallel=False)
+def eval_quad(z0, k, u, c=0.0):
+    if abs(k - 0.5) < 1.0e-4:
+        k = 0.5
+
+    npt = len(z0)
+    npb = 1
+
+    k2 = k ** 2
+    omega = zeros(npb)
+    flux = zeros((npt, npb))
+    le = zeros(npt)
+    ld = zeros(npt)
+    ed = zeros(npt)
+
+    for i in range(npb):
+        omega[i] = 1.0 - u[2 * i - 1] / 3.0 - u[2 * i] / 6.0
+
+    for i in prange(npt):
+        z = z0[i]
+
+        if abs(z - k) < 1e-6:
+            z += 1e-6
+
+        # The source is unocculted
+        if z > 1.0 + k or z < 0.0:
+            flux[i, :] = 1.0
+            le[i] = 0.0
+            ld[i] = 0.0
+            ed[i] = 0.0
+            continue
+
+        # The source is completely occulted
+        elif (k >= 1.0 and z <= k - 1.0):
+            flux[i, :] = 0.0
+            le[i] = 1.0
+            ld[i] = 1.0
+            ed[i] = 1.0
+            continue
+
+        z2 = z ** 2
+        x1 = (k - z) ** 2
+        x2 = (k + z) ** 2
+        x3 = k ** 2 - z ** 2
+
+        # The source is partially occulted and the occulting object crosses the limb
+        # Equation (26):
+        if z >= abs(1.0 - k) and z <= 1.0 + k:
+            kap1 = arccos(min((1.0 - k2 + z2) / (2.0 * z), 1.0))
+            kap0 = arccos(min((k2 + z2 - 1.0) / (2.0 * k * z), 1.0))
+            le[i] = k2 * kap0 + kap1
+            le[i] = (le[i] - 0.5 * sqrt(max(4.0 * z2 - (1.0 + z2 - k2) ** 2, 0.0))) * INV_PI
+
+        # The occulting object transits the source star (but doesn't completely cover it):
+        if z <= 1.0 - k:
+            le[i] = k2
+
+        # The edge of the occulting star lies at the origin- special expressions in this case:
+        if abs(z - k) < 1.e-4 * (z + k):
+            # ! Table 3, Case V.:
+            if (k == 0.5):
+                ld[i] = 1.0 / 3.0 - 4.0 * INV_PI / 9.0
+                ed[i] = 3.0 / 32.0
+            elif (z > 0.5):
+                q = 0.50 / k
+                Kk = ellk(q)
+                Ek = ellec(q)
+                ld[i] = 1.0 / 3.0 + 16.0 * k / 9.0 * INV_PI * (2.0 * k2 - 1.0) * Ek - (
+                                                                                       32.0 * k ** 4 - 20.0 * k2 + 3.0) / 9.0 * INV_PI / k * Kk
+                ed[i] = 1.0 / 2.0 * INV_PI * (
+                    kap1 + k2 * (k2 + 2.0 * z2) * kap0 - (1.0 + 5.0 * k2 + z2) / 4.0 * sqrt((1.0 - x1) * (x2 - 1.0)))
+            elif (z < 0.5):
+                # ! Table 3, Case VI.:
+                q = 2.0 * k
+                Kk = ellk(q)
+                Ek = ellec(q)
+                ld[i] = 1.0 / 3.0 + 2.0 / 9.0 * INV_PI * (4.0 * (2.0 * k2 - 1.0) * Ek + (1.0 - 4.0 * k2) * Kk)
+                ed[i] = k2 / 2.0 * (k2 + 2.0 * z2)
+
+        # The occulting star partly occults the source and crosses the limb:
+        # Table 3, Case III:
+        if ((z > 0.5 + abs(k - 0.5) and z < 1.0 + k) or (k > 0.50 and z > abs(1.0 - k) and z < k)):
+            q = sqrt((1.0 - (k - z) ** 2) / 4.0 / z / k)
+            Kk = ellk(q)
+            Ek = ellec(q)
+            n = 1.0 / x1 - 1.0
+            Pk = ellpicb(n, q)
+            ld[i] = 1.0 / 9.0 * INV_PI / sqrt(k * z) * (
+                ((1.0 - x2) * (2.0 * x2 + x1 - 3.0) - 3.0 * x3 * (x2 - 2.0)) * Kk + 4.0 * k * z * (
+                    z2 + 7.0 * k2 - 4.0) * Ek - 3.0 * x3 / x1 * Pk)
+            if (z < k):
+                ld[i] = ld[i] + 2.0 / 3.0
+            ed[i] = 1.0 / 2.0 * INV_PI * (
+                kap1 + k2 * (k2 + 2.0 * z2) * kap0 - (1.0 + 5.0 * k2 + z2) / 4.0 * sqrt((1.0 - x1) * (x2 - 1.0)))
+
+        # The occulting star transits the source:
+        # Table 3, Case IV.:
+        if (k <= 1.0 and z <= (1.0 - k)):
+            q = sqrt((x2 - x1) / (1.0 - x1))
+            Kk = ellk(q)
+            Ek = ellec(q)
+            n = x2 / x1 - 1.0
+            Pk = ellpicb(n, q)
+            ld[i] = 2.0 / 9.0 * INV_PI / sqrt(1.0 - x1) * (
+                (1.0 - 5.0 * z2 + k2 + x3 * x3) * Kk + (1.0 - x1) * (z2 + 7.0 * k2 - 4.0) * Ek - 3.0 * x3 / x1 * Pk)
+            if (z < k):
+                ld[i] = ld[i] + 2.0 / 3.0
+            if (abs(k + z - 1.0) < 1.e-4):
+                ld[i] = 2.0 / 3.0 * INV_PI * arccos(1.0 - 2.0 * k) - 4.0 / 9.0 * INV_PI * sqrt(k * (1.0 - k)) * (
+                    3.0 + 2.0 * k - 8.0 * k2)
+            ed[i] = k2 / 2.0 * (k2 + 2.0 * z2)
+
+        for j in range(npb):
+            iu = 2 * j - 1
+            iv = 2 * j
+            flux[i, j] = 1.0 - ((1.0 - u[iu] - 2.0 * u[iv]) * le[i] + (u[iu] + 2.0 * u[iv]) * ld[i] + u[iv] * ed[i]) / omega[j]
+        flux[i, :] = c + (1.0 - c) * flux[i, :]
+
+    return flux, ld, le, ed
+
+@jit(cache=True, nopython=True)
+def eval_chromosphere(zs, k, c):
+    """Optically thin chromosphere model presented in
+       Schlawin, Agol, Walkowicz, Covey & Lloyd (2010)"""
+    nz = len(zs)
+    flux = ones_like(zs)
+
+    for i in range(nz):
+        z = zs[i]
+        if ((z > 0.0) and (z - k < 1.0)):
+            zmk2 = (z - k) ** 2
+            if (z + k < 1.0):
+                t = sqrt(4.0 * z * k / (1.0 - zmk2))
+                flux[i] = (4.0 / sqrt(1.0 - zmk2) *
+                           ((zmk2 - 1.0) * ellec(t)
+                            - (z ** 2 - k ** 2) * ellk(t)
+                            + (z + k) / (z - k) * ellpicb(4.0 * z * k / zmk2, t)))
+
+            elif (z + k > 1.0):
+                t = sqrt((1.0 - zmk2) / (4.0 * z * k))
+                flux[i] = (2.0 / (z - k) / sqrt(z * k) *
+                           (4.0 * z * k * (k - z) * ellec(t)
+                            + (-z + 2.0 * z ** 2 * k + k - 2.0 * k ** 3) * ellk(t)
+                            + (z + k) * ellpicb(1.0 / zmk2 - 1.0, t)))
+            if (k > z):
+                flux[i] += FOUR_PI
+            flux[i] = 1.0 - flux[i] / FOUR_PI
+            flux[i] = c + (1.0 - c) * flux[i]
+    return flux
+
+

src/orbits.py

@@ -16,7 +16,7 @@
 ## 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 
 from numpy import pi
-from .orbits_f import orbits
+#from .orbits_f import orbits
 
 TWO_PI = 2 * pi
 
@@ -28,32 +28,32 @@ def not_implemented(*nargs, **kwargs):
 class Orbit(object):
     methods = 'iteration newton ps3 ps5 interpolation'.split()
 
-    ea_functions = dict(iteration=not_implemented,
-                        newton=orbits.ea_eccentric_newton,
-                        ps3=not_implemented,
-                        ps5=not_implemented,
-                        interpolation=not_implemented)
-
-    ta_functions = dict(iteration=orbits.ta_eccentric_iter,
-                        newton=orbits.ta_eccentric_newton,
-                        ps3=orbits.ta_eccentric_ps3,
-                        ps5=orbits.ta_eccentric_ps5,
-                        interpolation=orbits.ta_eccentric_bilerp)
-
-    z_functions = dict(iteration=orbits.z_eccentric_iter,
-                       newton=orbits.z_eccentric_newton,
-                       ps3=orbits.z_eccentric_ps3,
-                       ps5=orbits.z_eccentric_ps5,
-                       interpolation=orbits.z_eccentric_ip)
+    # ea_functions = dict(iteration=not_implemented,
+    #                     newton=orbits.ea_eccentric_newton,
+    #                     ps3=not_implemented,
+    #                     ps5=not_implemented,
+    #                     interpolation=not_implemented)
+    #
+    # ta_functions = dict(iteration=orbits.ta_eccentric_iter,
+    #                     newton=orbits.ta_eccentric_newton,
+    #                     ps3=orbits.ta_eccentric_ps3,
+    #                     ps5=orbits.ta_eccentric_ps5,
+    #                     interpolation=orbits.ta_eccentric_bilerp)
+    #
+    # z_functions = dict(iteration=orbits.z_eccentric_iter,
+    #                    newton=orbits.z_eccentric_newton,
+    #                    ps3=orbits.z_eccentric_ps3,
+    #                    ps5=orbits.z_eccentric_ps5,
+    #                    interpolation=orbits.z_eccentric_ip)
 
     def __init__(self, method='iteration', nthr=0, circular_e_threshold=1e-5):
         assert method in self.methods
         self.nthr = nthr
         self.method = method
         self._mine = circular_e_threshold
-        self._ea_function = self.ea_functions[method]
-        self._ta_function = self.ta_functions[method]
-        self._z_function = self.z_functions[method]
+        #self._ea_function = self.ea_functions[method]
+        #self._ta_function = self.ta_functions[method]
+        #self._z_function = self.z_functions[method]
 
     def mean_anomaly(self, time, t0, p, e=0., w=0.):
         return orbits.mean_anomaly(time, t0, p, e, w, self.nthr)
@@ -94,8 +94,8 @@ def projected_distance(self, time, t0, p, a, i):
         return orbits.z_circular(time, t0, p, a, i, self.nthr)
 
 
-eclipse_shift = orbits.eclipse_shift_ex
-duration_circular = orbits.duration_circular
-duration_eccentric = orbits.duration_eccentric_w
-z_circular = orbits.z_circular
-z_eccentric_newton = orbits.z_eccentric_newton
+#eclipse_shift = orbits.eclipse_shift_ex
+#duration_circular = orbits.duration_circular
+#duration_eccentric = orbits.duration_eccentric_w
+#z_circular = orbits.z_circular
+#z_eccentric_newton = orbits.z_eccentric_newton