Correct handling of tidally locked vs. unlocked orbits

README.md

@@ -27,6 +27,9 @@ at the University of Chicago in Chicago, IL, USA.
 
  - [x] Take time and orbital period as inputs instead of angle $\alpha$
  - [x] Optionally handle elliptical orbits
+ - [ ] Support `binsize` argument to dual transit flux
+ - [ ] Support `locked` argument to dual transit flux
+ - [ ] Support `eccentric` argument to dual transit flux
 
 ## :sparkles: Credits
 

demos/gradients.py

@@ -0,0 +1,35 @@
+import time
+import pocky
+import numpy as np
+import greenlantern
+import matplotlib.pyplot as plt
+from nordplotlib.png import install; install()
+
+a, b, c, ds, tc, beta, zeta, eta, xi, u1, u2, Porb = \
+    0.1, 0.2, 0.3, 5., 0., 0.01, 0.3, 0.2, 0.1, 0.1, -0.02, 2 * np.pi
+
+q1 = (u1 + u2)**2
+q2 = u1 / (2 * (u1 + u2)) if u1 > 0 else 0.
+
+ctx = pocky.Context.default()
+ctx1 = greenlantern.Context(ctx)
+
+nt = 1000
+time = np.linspace(-0.3, 0.3, nt)
+time = pocky.BufferPair(ctx, time.astype(np.float32))
+time.copy_to_device()
+time.dirty = False
+
+params = np.array([[a, b, c, ds, tc, beta, zeta, eta, xi, q1, q2, Porb]], dtype=np.float32)
+params = pocky.BufferPair(ctx, params)
+
+flux = np.empty((nt,), dtype=np.float32)
+flux = pocky.BufferPair(ctx, flux)
+
+dflux = np.empty((12, nt), dtype=np.float32)
+dflux = pocky.BufferPair(ctx, dflux)
+
+ctx1.ellipsoid_transit_flux_dual(time, params, flux=flux, dflux=dflux)
+plt.plot(time.host, dflux.host.T, lw=2)
+
+plt.show()

demos/versus_batman.py

@@ -7,8 +7,14 @@
 import matplotlib.pyplot as plt
 from nordplotlib.png import install; install()
 
-rr, semimajor, tc, beta, zeta, eta, xi, u1, u2, Porb, ecc, omega = \
-    0.1, 5., 0.01, -0.02, 0., 0., 0., 0.1, -0.03, 2 * np.pi, 0.3, 0.2
+rr, semimajor, tc, beta, zeta, eta, xi, u1, u2, Porb = \
+    0.1, 5., 0.01, -0.02, 0., 0., 0., 0.1, -0.03, 2 * np.pi
+test_eccentric = False
+
+if test_eccentric:
+    ecc, omega = 0.3, 0.2
+else:
+    ecc, omega = 0., -0.5 * np.pi
 
 q1 = (u1 + u2)**2
 q2 = u1 / (2 * (u1 + u2)) if u1 > 0 else 0.
@@ -48,16 +54,23 @@ def generate_timings():
         time_dev.copy_to_device()
         time_dev.dirty = False
 
-        params = np.array([[rr, rr, rr, semimajor, tc, beta, zeta, eta, xi, q1, q2, Porb, ecc, omega]], dtype=np.float32)
-        params = pocky.BufferPair(ctx, params)
-
-        flux = np.empty((params.host.shape[0], nt), dtype=np.float32)
+        flux = np.empty((1, nt), dtype=np.float32)
         flux = pocky.BufferPair(ctx, flux)
 
+        if test_eccentric:
+            params = np.array([[rr, rr, rr, semimajor, tc, beta,
+                zeta, eta, xi, q1, q2, Porb, ecc, omega]], dtype=np.float32)
+            params = pocky.BufferPair(ctx, params)
+        else:
+            params = np.array([[rr, rr, rr, semimajor, tc, beta,
+                zeta, eta, xi, q1, q2, Porb]], dtype=np.float32)
+            params = pocky.BufferPair(ctx, params)
+
         dt_greenlantern = 0
         for _ in range(ntries):
             t0 = time.time()
-            ctx1.ellipsoid_transit_flux(time_dev, params, flux=flux, eccentric=True)
+            ctx1.ellipsoid_transit_flux(time_dev, params,
+                flux=flux, eccentric=test_eccentric)
             t1 = time.time()
             dt_greenlantern += (t1 - t0) / flux.host.size
         dt_greenlantern /= ntries

setup.py

@@ -42,7 +42,7 @@ def run(self):
 
         return super().run()
 
-source_files = ['greenlantern.c', 'greenlantern_context.c', 'ellipsoid.c']
+source_files = ['greenlantern.c', 'greenlantern_context.c', 'ellipsoid.c', 'ellipsoid_dual.c']
 source_files = [os.path.join(source_dir, fname) for fname in source_files]
 
 header_files = ['greenlantern.h', 'greenlantern_context.h']

src/greenlantern/__init__.py

@@ -1,3 +1,3 @@
 from .ext import *
 
-__version__ = '1.4'
+__version__ = '1.5'

src/greenlantern/ext/ellipsoid.c

@@ -11,23 +11,25 @@ PyObject *ellipsoid_transit_flux(greenlantern_context_object *context,
     PyObject *args, PyObject *kwargs)
 {
     char buf[BUFSIZ];
-    char *keys[] = { "time", "params", "binsize", "eccentric", "flux", "queue", NULL };
+    char *keys[] = { "time", "params", "binsize", "locked", "eccentric", "flux", "queue", NULL };
 
-    PyObject *binsize = NULL, *queue_idx = NULL, *eccentric_flag = Py_False;
+    PyObject *binsize = NULL, *queue_idx = NULL, *locked_flag = Py_False, *eccentric_flag = Py_False;
     pocky_bufpair_object *input, *params, *flux = NULL;
 
     cl_int err;
     cl_command_queue queue;
     cl_kernel kernel = NULL;
     cl_event event;
+    cl_ushort locked;
 
     long worksz[2];
     size_t kernsz[2], locsz[2];
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O!O!|$O!O!O!O!", keys,
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O!O!|$O!O!O!O!O!", keys,
         pocky_api->bufpair_type, &input, pocky_api->bufpair_type, &params,
-        &PyFloat_Type, &binsize, &PyBool_Type, &eccentric_flag,
-        pocky_api->bufpair_type, &flux, &PyLong_Type, &queue_idx)) return NULL;
+        &PyFloat_Type, &binsize, &PyBool_Type, &locked_flag,
+        &PyBool_Type, &eccentric_flag, pocky_api->bufpair_type, &flux,
+        &PyLong_Type, &queue_idx)) return NULL;
 
     if ((input->host == NULL) ||
         (!PyArray_CheckExact((PyArrayObject *) input->host)) ||
@@ -123,18 +125,21 @@ PyObject *ellipsoid_transit_flux(greenlantern_context_object *context,
     clSetKernelArg(kernel, 0, sizeof(cl_mem), &(input->device));
     clSetKernelArg(kernel, 1, sizeof(cl_mem), &(params->device));
 
+    locked = (locked_flag == Py_True) ? 1 : 0;
+    clSetKernelArg(kernel, 2, sizeof(cl_ushort), &locked);
+
     locsz[0] = 32;
 
     if (!binsize)
     {
-        clSetKernelArg(kernel, 2, sizeof(cl_mem), &(flux->device));
+        clSetKernelArg(kernel, 3, sizeof(cl_mem), &(flux->device));
         locsz[1] = 1;
     }
     else
     {
         float binsize_val = (float) PyFloat_AsDouble(binsize);
-        clSetKernelArg(kernel, 2, sizeof(cl_float), &binsize_val);
-        clSetKernelArg(kernel, 3, sizeof(cl_mem), &(flux->device));
+        clSetKernelArg(kernel, 3, sizeof(cl_float), &binsize_val);
+        clSetKernelArg(kernel, 4, sizeof(cl_mem), &(flux->device));
         locsz[1] = 17;   /* number of points to bin together */
     }
 
@@ -178,188 +183,4 @@ PyObject *ellipsoid_transit_flux(greenlantern_context_object *context,
     return (PyObject *) flux;
 }
 
-PyObject *ellipsoid_transit_flux_dual(greenlantern_context_object *context,
-    PyObject *args, PyObject *kwargs)
-{
-    char buf[BUFSIZ];
-    char *keys[] = { "time", "params", "binsize", "flux", "dflux", "queue", NULL };
-
-    PyObject *binsize = NULL, *queue_idx = NULL;
-    pocky_bufpair_object *time, *params, *flux = NULL, *dflux = NULL;
-
-    cl_int err;
-    cl_command_queue queue;
-    cl_kernel kernel = NULL;
-    cl_event event;
-
-    long worksz[1], dfluxsz[2];
-    size_t kernsz[1], locsz[2];
-
-    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O!O!|$O!O!O!O!", keys,
-        pocky_api->bufpair_type, &time, pocky_api->bufpair_type, &params,
-        &PyFloat_Type, &binsize, pocky_api->bufpair_type, &flux,
-        pocky_api->bufpair_type, &dflux, &PyLong_Type, &queue_idx)) return NULL;
-
-    if ((time->host == NULL) ||
-        (!PyArray_CheckExact((PyArrayObject *) time->host)) ||
-        (PyArray_NDIM((PyArrayObject *) time->host) != 1) ||
-        (PyArray_TYPE((PyArrayObject *) time->host) != NPY_FLOAT32))
-    {
-        /* Invalid time array */
-        PyErr_SetString(PyExc_ValueError,
-            "Host time array should be a one-dimensional array of float32");
-        return NULL;
-    }
-
-    if ((params->host == NULL) ||
-        (!PyArray_CheckExact((PyArrayObject *) params->host)) ||
-        (PyArray_NDIM((PyArrayObject *) params->host) != 2) ||
-        (PyArray_DIM((PyArrayObject *) params->host, 0) != 1) ||
-        (PyArray_TYPE((PyArrayObject *) params->host) != NPY_FLOAT32))
-    {
-        /* Invalid params array */
-        PyErr_SetString(PyExc_ValueError,
-            "Host parameters should be a two-dimensional array of float32");
-        return NULL;
-    }
-
-    /* Construct the global dimensions for the kernel */
-    worksz[0] = PyArray_DIM((PyArrayObject *) time->host, 0);
-
-    /* Create a flux buffer if needed */
-    if ((flux == NULL) &&
-        (pocky_api->bufpair_empty_from_shape(context->pocky, 1, worksz, &flux)))
-    {
-        PyErr_SetString(PyExc_RuntimeError, "Could not create a flux array on-the-fly");
-        return NULL;
-    }
-
-    dfluxsz[0] = 12;
-    dfluxsz[1] = worksz[0];
-
-    /* Create a dflux buffer if needed */
-    if ((dflux == NULL) &&
-        (pocky_api->bufpair_empty_from_shape(context->pocky, 2, dfluxsz, &dflux)))
-    {
-        PyErr_SetString(PyExc_RuntimeError, "Could not create a dflux array on-the-fly");
-        return NULL;
-    }
-
-    if (worksz[0] == 0)
-    {
-        /* Early exit -- no work to do */
-        return PyTuple_Pack(2, flux, dflux);
-    }
-
-    /* Extract the queue handle (first by default) and kernel */
-    if (!queue_idx) queue = context->pocky->queues[0];
-    else
-    {
-        long idx = PyLong_AsLong(queue_idx);
-        queue = context->pocky->queues[idx];
-    }
-
-    /* Choose the appropriate kernel */
-    if (!binsize) kernel = context->kernels.ellipsoid_transit_flux_dual;
-    else kernel = NULL; // FIXME context->kernels.ellipsoid_transit_flux_binned_dual;
-
-    /* Copy data to the device */
-    if (time->dirty == Py_True)
-    {
-        err = clEnqueueWriteBuffer(queue, time->device,
-            CL_TRUE, 0, time->host_size * sizeof(cl_float),
-            PyArray_DATA((PyArrayObject *) time->host), 0, NULL, NULL);
-        if (err != CL_SUCCESS)
-        {
-            snprintf(buf, BUFSIZ, greenlantern_ocl_fmt_internal,
-                pocky_api->opencl_error_to_string(err), err);
-            PyErr_SetString(greenlantern_ocl_error, buf);
-            return NULL;
-        }
-    }
-
-    if (params->dirty == Py_True)
-    {
-        err = clEnqueueWriteBuffer(queue, params->device,
-            CL_TRUE, 0, params->host_size * sizeof(cl_float),
-            PyArray_DATA((PyArrayObject *) params->host), 0, NULL, NULL);
-        if (err != CL_SUCCESS)
-        {
-            snprintf(buf, BUFSIZ, greenlantern_ocl_fmt_internal,
-                pocky_api->opencl_error_to_string(err), err);
-            PyErr_SetString(greenlantern_ocl_error, buf);
-            return NULL;
-        }
-    }
-
-    /* Ready to run kernel */
-    clSetKernelArg(kernel, 0, sizeof(cl_mem), &(time->device));
-    clSetKernelArg(kernel, 1, sizeof(cl_mem), &(params->device));
-
-    locsz[0] = 32;
-
-    if (!binsize)
-    {
-        clSetKernelArg(kernel, 2, sizeof(cl_mem), &(flux->device));
-        clSetKernelArg(kernel, 3, sizeof(cl_mem), &(dflux->device));
-        locsz[1] = 1;
-    }
-    else
-    {
-        float binsize_val = (float) PyFloat_AsDouble(binsize);
-        clSetKernelArg(kernel, 2, sizeof(cl_float), &binsize_val);
-        clSetKernelArg(kernel, 3, sizeof(cl_mem), &(flux->device));
-        clSetKernelArg(kernel, 4, sizeof(cl_mem), &(dflux->device));
-        locsz[1] = 17;   /* number of points to bin together */
-    }
-
-    kernsz[0] = worksz[0] * locsz[0];
-
-    err = clEnqueueNDRangeKernel(queue, kernel,
-        1, NULL, kernsz, locsz, 0, NULL, &event);
-    if (err != CL_SUCCESS)
-    {
-        snprintf(buf, BUFSIZ, greenlantern_ocl_fmt_internal,
-            pocky_api->opencl_error_to_string(err), err);
-        PyErr_SetString(greenlantern_ocl_error, buf);
-        return NULL;
-    }
-
-    /* Block until the kernel has run */
-    err = clWaitForEvents(1, &event);
-    if (err != CL_SUCCESS)
-    {
-        snprintf(buf, BUFSIZ, greenlantern_ocl_fmt_internal,
-            pocky_api->opencl_error_to_string(err), err);
-        PyErr_SetString(greenlantern_ocl_error, buf);
-        return NULL;
-    }
-    clReleaseEvent(event);
-
-    /* Copy data back and block until we have all the results */
-    err = clEnqueueReadBuffer(queue, flux->device,
-        CL_TRUE, 0, flux->host_size * sizeof(cl_float),
-        PyArray_DATA((PyArrayObject *) flux->host), 0, NULL, NULL);
-    if (err != CL_SUCCESS)
-    {
-        snprintf(buf, BUFSIZ, greenlantern_ocl_fmt_internal,
-            pocky_api->opencl_error_to_string(err), err);
-        PyErr_SetString(greenlantern_ocl_error, buf);
-        return NULL;
-    }
-
-    err = clEnqueueReadBuffer(queue, dflux->device,
-        CL_TRUE, 0, dflux->host_size * sizeof(cl_float),
-        PyArray_DATA((PyArrayObject *) dflux->host), 0, NULL, NULL);
-    if (err != CL_SUCCESS)
-    {
-        snprintf(buf, BUFSIZ, greenlantern_ocl_fmt_internal,
-            pocky_api->opencl_error_to_string(err), err);
-        PyErr_SetString(greenlantern_ocl_error, buf);
-        return NULL;
-    }
-
-    return PyTuple_Pack(2, flux, dflux);
-}
-
 /* vim: set ft=c.doxygen: */