Skip to content

Commit

Permalink
Updated the qpower2 model to use Taylor series z computation routines.
Browse files Browse the repository at this point in the history
  • Loading branch information
@hpparvi
hpparvi committed Sep 13, 2020
1 parent 415adca commit 2bfc5ee
Show file tree
Hide file tree
Showing 3 changed files with 83 additions and 60 deletions.
129 changes: 77 additions & 52 deletions pytransit/models/numba/qpower2_nb.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,10 +47,10 @@
This module implements the qpower2 transit model by Maxted & Gill (A&A, 622, A33, 2019).
"""

from numpy import any, sqrt, pi, arccos, ones_like, atleast_2d, zeros, atleast_1d, nan, copysign
from numpy import any, sqrt, pi, arccos, ones_like, atleast_2d, zeros, atleast_1d, nan, copysign, fmax, floor
from numba import njit, prange

from ...orbits.orbits_py import z_ip_s
from ...orbits.taylor_z import vajs_from_paiew, z_taylor_st


@njit(fastmath=True)
Expand Down Expand Up @@ -103,84 +103,109 @@ def qpower2_z_s(z, k, u):


@njit(parallel=True, fastmath=True)
def qpower2_model_v(t, k, ldc, t0, p, a, i, e, w, lcids, pbids, nsamples, exptimes, es, ms, tae):
def qpower2_model_v(t, k, ldc, t0, p, a, i, e, w, lcids, pbids, nsamples, exptimes):
t0, p, a, i, e, w = atleast_1d(t0), atleast_1d(p), atleast_1d(a), atleast_1d(i), atleast_1d(e), atleast_1d(w)
k = atleast_2d(k)
ldc = atleast_2d(ldc)
npv = k.shape[0]
npt = t.size
flux = zeros((npv, npt))
for j in prange(npt):
for ipv in range(npv):
ilc = lcids[j]
ipb = pbids[ilc]

if k.shape[1] == 1:
_k = k[ipv, 0]
for ipv in prange(npv):
x0, y0, vx, vy, ax, ay, jx, jy, sx, sy = vajs_from_paiew(t0[ipv], p[ipv], a[ipv], i[ipv], e[ipv], w[ipv])
half_window_width = fmax(0.125, (2.0 + k[0, 0])/vx)

for j in range(npt):
epoch = floor((t[j] - t0[ipv] + 0.5 * p[ipv]) / p[ipv])
tc = t[j] - (t0[ipv] + epoch * p[ipv])
if abs(tc) > half_window_width:
flux[ipv, j] = 1.0
else:
_k = k[ipv, ipb]
ilc = lcids[j]
ipb = pbids[ilc]

for isample in range(1, nsamples[ilc] + 1):
time_offset = exptimes[ilc] * ((isample - 0.5) / nsamples[ilc] - 0.5)
z = z_ip_s(t[j] + time_offset, t0[ipv], p[ipv], a[ipv], i[ipv], e[ipv], w[ipv], es, ms, tae)
if z > 1.0 + _k:
flux[ipv, j] += 1.
if k.shape[1] == 1:
_k = k[ipv, 0]
else:
flux[ipv, j] += qpower2_z_s(z, _k, ldc[ipv, 2*ipb:2*(ipb+1)])
flux[ipv, j] /= nsamples[ilc]
_k = k[ipv, ipb]

for isample in range(1, nsamples[ilc] + 1):
time_offset = exptimes[ilc] * ((isample - 0.5) / nsamples[ilc] - 0.5)
z = z_taylor_st(tc + time_offset, y0, vx, vy, ax, ay, jx, jy, sx, sy)
if z > 1.0 + _k:
flux[ipv, j] += 1.
else:
flux[ipv, j] += qpower2_z_s(z, _k, ldc[ipv, 2*ipb:2*(ipb+1)])
flux[ipv, j] /= nsamples[ilc]
return flux


@njit(parallel=True, fastmath=True)
def qpower2_model_s(t, k, ldc, t0, p, a, i, e, w, lcids, pbids, nsamples, exptimes, es, ms, tae):
@njit(parallel=False, fastmath=True)
def qpower2_model_s(t, k, ldc, t0, p, a, i, e, w, lcids, pbids, nsamples, exptimes):
k = atleast_1d(k)
ldc = atleast_1d(ldc)
npt = t.size

x0, y0, vx, vy, ax, ay, jx, jy, sx, sy = vajs_from_paiew(t0, p, a, i, e, w)
half_window_width = fmax(0.125, (2.0 + k[0]) / vx)

flux = zeros(npt)
for j in prange(npt):
ilc = lcids[j]
ipb = pbids[ilc]
_k = k[0] if k.size == 1 else k[ipb]

for isample in range(1, nsamples[ilc] + 1):
time_offset = exptimes[ilc] * ((isample - 0.5) / nsamples[ilc] - 0.5)
z = z_ip_s(t[j] + time_offset, t0, p, a, i, e, w, es, ms, tae)
if z > 1.0 + _k:
flux[j] += 1.
else:
flux[j] += qpower2_z_s(z, _k, ldc[2*ipb:2*(ipb+1)])
flux[j] /= nsamples[ilc]
for j in range(npt):
epoch = floor((t[j] - t0 + 0.5 * p) / p)
tc = t[j] - (t0 + epoch * p)
if abs(tc) > half_window_width:
flux[j] = 1.0
else:
ilc = lcids[j]
ipb = pbids[ilc]
_k = k[0] if k.size == 1 else k[ipb]

for isample in range(1, nsamples[ilc] + 1):
time_offset = exptimes[ilc] * ((isample - 0.5) / nsamples[ilc] - 0.5)
z = z_taylor_st(tc + time_offset, y0, vx, vy, ax, ay, jx, jy, sx, sy)
if z > 1.0 + _k:
flux[j] += 1.
else:
flux[j] += qpower2_z_s(z, _k, ldc[2*ipb:2*(ipb+1)])
flux[j] /= nsamples[ilc]
return flux


@njit(parallel=True, fastmath=True)
def qpower2_model_pv(t, pvp, ldc, lcids, pbids, nsamples, exptimes, es, ms, tae):
def qpower2_model_pv(t, pvp, ldc, lcids, pbids, nsamples, exptimes):
pvp = atleast_2d(pvp)
ldc = atleast_2d(ldc)
npv = pvp.shape[0]
nk = pvp.shape[1] - 6
npt = t.size
flux = zeros((npv, npt))
for j in prange(npt):
for ipv in range(npv):
t0, p, a, i, e, w = pvp[ipv,nk:]
ilc = lcids[j]
ipb = pbids[ilc]

if nk == 1:
k = pvp[ipv, 0]
for ipv in prange(npv):
t0, p, a, i, e, w = pvp[ipv, nk:]
x0, y0, vx, vy, ax, ay, jx, jy, sx, sy = vajs_from_paiew(t0, p, a, i, e, w)
half_window_width = fmax(0.125, (2 + pvp[ipv, 0])/vx)
for j in range(npt):
epoch = floor((t[j] - t0 + 0.5*p)/p)
tc = t[j] - (t0 + epoch*p)
if abs(tc) > half_window_width:
flux[ipv, j] = 1.0
else:
if ipb < nk:
k = pvp[ipv, ipb]
else:
k = nan
ilc = lcids[j]
ipb = pbids[ilc]

for isample in range(1,nsamples[ilc]+1):
time_offset = exptimes[ilc] * ((isample - 0.5) / nsamples[ilc] - 0.5)
z = z_ip_s(t[j]+time_offset, t0, p, a, i, e, w, es, ms, tae)
if z < 0.0 or z > 1.0+k:
flux[ipv, j] += 1.
if nk == 1:
k = pvp[ipv, 0]
else:
flux[ipv, j] += qpower2_z_s(z, k, ldc[ipv, 2*ipb:2*(ipb+1)])
flux[ipv, j] /= nsamples[ilc]
if ipb < nk:
k = pvp[ipv, ipb]
else:
k = nan

for isample in range(1,nsamples[ilc]+1):
time_offset = exptimes[ilc] * ((isample - 0.5) / nsamples[ilc] - 0.5)
z = z_taylor_st(tc + time_offset, y0, vx, vy, ax, ay, jx, jy, sx, sy)
if z < 0.0 or z > 1.0+k:
flux[ipv, j] += 1.
else:
flux[ipv, j] += qpower2_z_s(z, k, ldc[ipv, 2*ipb:2*(ipb+1)])
flux[ipv, j] /= nsamples[ilc]
return flux
10 changes: 4 additions & 6 deletions pytransit/models/qpower2.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -89,8 +89,7 @@ def evaluate(self, k: Union[float, ndarray], ldc: ndarray, t0: Union[float, ndar
if e is None:
e, w = 0., 0.

flux = qpower2_model_s(self.time, k, ldc, t0, p, a, i, e, w, self.lcids, self.pbids, self.nsamples,
self.exptimes, self._es, self._ms, self._tae)
flux = qpower2_model_s(self.time, k, ldc, t0, p, a, i, e, w, self.lcids, self.pbids, self.nsamples, self.exptimes)
# Parameter population branch
# ---------------------------
else:
Expand All @@ -101,8 +100,7 @@ def evaluate(self, k: Union[float, ndarray], ldc: ndarray, t0: Union[float, ndar
if k.ndim == 1:
k = k.reshape((k.size,1))

flux = qpower2_model_v(self.time, k, ldc, t0, p, a, i, e, w, self.lcids, self.pbids, self.nsamples,
self.exptimes, self._es, self._ms, self._tae)
flux = qpower2_model_v(self.time, k, ldc, t0, p, a, i, e, w, self.lcids, self.pbids, self.nsamples, self.exptimes)
return squeeze(flux)

def evaluate_ps(self, k: float, ldc: ndarray, t0: float, p: float, a: float, i: float, e: float = 0., w: float = 0.) -> ndarray:
Expand Down Expand Up @@ -140,7 +138,7 @@ def evaluate_ps(self, k: float, ldc: ndarray, t0: float, p: float, a: float, i:
"""
assert self.time is not None, "Need to set the data before calling the transit model."
k = asarray(k)
flux = qpower2_model_s(self.time, k, ldc, t0, p, a, i, e, w, self.lcids, self.pbids, self.nsamples, self.exptimes, self._es, self._ms, self._tae)
flux = qpower2_model_s(self.time, k, ldc, t0, p, a, i, e, w, self.lcids, self.pbids, self.nsamples, self.exptimes)
return squeeze(flux)

def evaluate_pv(self, pvp: ndarray, ldc: ndarray) -> ndarray:
Expand All @@ -164,6 +162,6 @@ def evaluate_pv(self, pvp: ndarray, ldc: ndarray) -> ndarray:
Modelled flux either as a 1D or 2D ndarray.
"""
assert self.time is not None, "Need to set the data before calling the transit model."
flux = qpower2_model_pv(self.time, pvp, ldc, self.lcids, self.pbids, self.nsamples, self.exptimes, self._es, self._ms, self._tae)
flux = qpower2_model_pv(self.time, pvp, ldc, self.lcids, self.pbids, self.nsamples, self.exptimes)
return squeeze(flux)

4 changes: 2 additions & 2 deletions tests/test_qpower2_nb.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,8 +53,8 @@ def test_uniform_z_s_edge_cases(self):
# -------------------------------------
assert_almost_equal(qpower2_z_s(-0.0, 1.0, self.ldc), 1.0)
assert_almost_equal(qpower2_z_s( 0.0, 1.0, self.ldc), 0.0)
with assert_raises(ValueError):
qpower2_z_s( 0.0, 1.1, self.ldc)
#with assert_raises(ValueError):
# qpower2_z_s( 0.0, 1.1, self.ldc)


if __name__ == '__main__':
Expand Down

0 comments on commit 2bfc5ee

Please sign in to comment.