Notation tweak, add gradient binning support

demos/gradients.py

@@ -14,7 +14,7 @@
 ctx = pocky.Context.default()
 ctx1 = greenlantern.Context(ctx)
 
-nt = 1000
+nt = 100
 time = np.linspace(-0.3, 0.3, nt)
 time = pocky.BufferPair(ctx, time.astype(np.float32))
 time.copy_to_device()
@@ -29,7 +29,10 @@
 dflux = np.empty((12, nt), dtype=np.float32)
 dflux = pocky.BufferPair(ctx, dflux)
 
-ctx1.ellipsoid_transit_flux_dual(time, params, flux=flux, dflux=dflux)
+ctx1.ellipsoid_transit_flux_dual(time, params, flux=flux, dflux=dflux, binsize=0.01*Porb)
 plt.plot(time.host, dflux.host.T, lw=2)
 
+ctx1.ellipsoid_transit_flux_dual(time, params, flux=flux, dflux=dflux)
+plt.plot(time.host, dflux.host.T, lw=1, ls='dashed')
+
 plt.show()

demos/spherical.py

@@ -60,15 +60,15 @@
 
     params = np.array([[rr, rr, rr, semimajor, t0, beta, zeta, eta, xi, q1, q2, Porb, ecc, omega]], dtype=np.float32)
     params = pocky.BufferPair(ctx, params)
-    ctx1.ellipsoid_transit_flux(time_dev, params, flux=flux, eccentric=True)
+    ctx1.ellipsoid_transit_flux(time_dev, params, flux=flux, eccentric=True, locked=False)
 
     axs[0].plot(time, model_flux, lw=2, ls='solid', alpha=0.6, c=f'C{i}')
     axs[0].plot(time, flux.host[0], lw=2, ls='dashed', alpha=1, c=f'C{i}')
 
     axs[1].scatter(time, 1e6 * (flux.host[0] - model_flux),
         alpha=0.05, s=8, c=f'C{i}', rasterized=True)
 
-    label = rf'$R_p / R_\star = {rr:.2f}, d_\star = {semimajor:.0f}, i = {90.-beta_deg:.0f}^\circ, u_1 = {u1:.2f}, u_2 = {u2:.2f}$'
+    label = rf'$R_p / R_\star = {rr:.2f}, d / R_\star = {semimajor:.0f}, i = {90.-beta_deg:.0f}^\circ, u_1 = {u1:.2f}, u_2 = {u2:.2f}$'
     handles.append(Line2D([0], [0], c=f'C{i}', lw=2, label=label))
 
 axs[1].set_xlabel(r'Time $t$ (code units)')

pyproject.toml

@@ -2,7 +2,7 @@
 build-backend = "setuptools.build_meta"
 requires = [
     "setuptools",
-    "numpy<2",
+    "numpy",
     "pocky @ git+https://github.com/emprice/pocky.git@main",
 ]
 

setup.py

@@ -16,11 +16,13 @@ def gen_kernel_paths():
         yield kernel_name, kernel_input
 
 kernel_paths = list(gen_kernel_paths())
-kernel_frag_template = textwrap.dedent('''\
-const char kernel_{kernel_name}[] = R"===(
-{kernel_content}
-)===";\n\n''')
-kernel_agg_template = textwrap.dedent('''\
+
+def make_kernel_frag(kernel_name, kernel_content):
+    return f'const char kernel_{kernel_name}[] = \n"' + \
+        '\\n"\n"'.join(kernel_content.replace('"', '\\"').split('\n')) + '\\n";\n\n'
+
+def make_kernel_agg(num_kernels, kernel_vars):
+    return textwrap.dedent(f'''\
 const cl_uint num_kernel_frags = {num_kernels};
 const char *kernel_frags[] = {{ {kernel_vars} }};
 ''')
@@ -31,12 +33,10 @@ def run(self):
             for kernel_name, kernel_input in kernel_paths:
                 with open(kernel_input, 'r') as kernel_in:
                     kernel_content = kernel_in.read()
-                    kernel_frag = kernel_frag_template.format(
-                        kernel_content=kernel_content, kernel_name=kernel_name)
+                    kernel_frag = make_kernel_frag(kernel_name, kernel_content)
                     kernel_def.write(kernel_frag)
 
-            kernel_agg = kernel_agg_template.format(
-                num_kernels=len(kernel_paths),
+            kernel_agg = make_kernel_agg(len(kernel_paths),
                 kernel_vars=', '.join([f'kernel_{name}' for name, _ in kernel_paths]))
             kernel_def.write(kernel_agg)
 

src/greenlantern/__init__.py

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

src/greenlantern/ext/ellipsoid_dual.c

@@ -22,8 +22,8 @@ PyObject *ellipsoid_transit_flux_dual(greenlantern_context_object *context,
     cl_event event;
     cl_ushort locked;
 
-    long worksz[1], dfluxsz[2];
-    size_t kernsz[1], locsz[2];
+    long worksz[2], dfluxsz[2];
+    size_t kernsz[2], locsz[2];
 
     if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O!O!|$O!O!O!O!O!", keys,
         pocky_api->bufpair_type, &time, pocky_api->bufpair_type, &params,
@@ -56,6 +56,7 @@ PyObject *ellipsoid_transit_flux_dual(greenlantern_context_object *context,
 
     /* Construct the global dimensions for the kernel */
     worksz[0] = PyArray_DIM((PyArrayObject *) time->host, 0);
+    worksz[1] = 1;
 
     /* Create a flux buffer if needed */
     if ((flux == NULL) &&
@@ -92,7 +93,7 @@ PyObject *ellipsoid_transit_flux_dual(greenlantern_context_object *context,
 
     /* Choose the appropriate kernel */
     if (!binsize) kernel = context->kernels.ellipsoid_transit_flux_dual;
-    else kernel = NULL; // FIXME context->kernels.ellipsoid_transit_flux_binned_dual;
+    else kernel = context->kernels.ellipsoid_transit_flux_binned_dual;
 
     /* Copy data to the device */
     if (time->dirty == Py_True)
@@ -130,27 +131,30 @@ PyObject *ellipsoid_transit_flux_dual(greenlantern_context_object *context,
     locked = (locked_flag == Py_True) ? 1 : 0;
     clSetKernelArg(kernel, 2, sizeof(cl_ushort), &locked);
 
-    locsz[0] = 32;
-
     if (!binsize)
     {
         clSetKernelArg(kernel, 3, sizeof(cl_mem), &(flux->device));
         clSetKernelArg(kernel, 4, sizeof(cl_mem), &(dflux->device));
-        locsz[1] = 1;
+
+        locsz[0] = 32;  /* number of threads working on summation */
+        locsz[1] = 1;   /* number of points to bin together (unused) */
     }
     else
     {
         float binsize_val = (float) PyFloat_AsDouble(binsize);
         clSetKernelArg(kernel, 3, sizeof(cl_float), &binsize_val);
         clSetKernelArg(kernel, 4, sizeof(cl_mem), &(flux->device));
         clSetKernelArg(kernel, 5, sizeof(cl_mem), &(dflux->device));
-        locsz[1] = 17;   /* number of points to bin together */
+
+        locsz[0] = 16;  /* number of threads working on summation */
+        locsz[1] = 17;  /* number of points to bin together */
     }
 
     kernsz[0] = worksz[0] * locsz[0];
+    kernsz[1] = worksz[1] * locsz[1];
 
     err = clEnqueueNDRangeKernel(queue, kernel,
-        1, NULL, kernsz, locsz, 0, NULL, &event);
+        2, NULL, kernsz, locsz, 0, NULL, &event);
     if (err != CL_SUCCESS)
     {
         snprintf(buf, BUFSIZ, greenlantern_ocl_fmt_internal,

src/greenlantern/ext/greenlantern_context.c

@@ -64,6 +64,11 @@ int greenlantern_context_init(greenlantern_context_object *context,
     if (err != CL_SUCCESS) return -1;
     clRetainKernel(context->kernels.ellipsoid_transit_flux_dual);
 
+    err = pocky_api->opencl_kernel_lookup_by_name(num_kernels, kernels,
+        "ellipsoid_transit_flux_binned_dual", &(context->kernels.ellipsoid_transit_flux_binned_dual));
+    if (err != CL_SUCCESS) return -1;
+    clRetainKernel(context->kernels.ellipsoid_transit_flux_binned_dual);
+
     /* Release any remaining kernel handles */
     for (idx = 0; idx < num_kernels; ++idx) clReleaseKernel(kernels[idx]);
     free(kernels);
@@ -82,7 +87,9 @@ void greenlantern_context_dealloc(greenlantern_context_object *self)
     clReleaseKernel(self->kernels.ellipsoid_transit_flux_binned_vector);
     clReleaseKernel(self->kernels.ellipsoid_eccentric_transit_flux_vector);
     clReleaseKernel(self->kernels.ellipsoid_eccentric_transit_flux_binned_vector);
+
     clReleaseKernel(self->kernels.ellipsoid_transit_flux_dual);
+    clReleaseKernel(self->kernels.ellipsoid_transit_flux_binned_dual);
 
     /* Release other handles */
     clReleaseProgram(self->program);

src/greenlantern/ext/include/greenlantern_context.h

@@ -18,6 +18,7 @@ typedef struct
         cl_kernel ellipsoid_eccentric_transit_flux_binned_vector;
 
         cl_kernel ellipsoid_transit_flux_dual;
+        cl_kernel ellipsoid_transit_flux_binned_dual;
     }
     kernels;    /**< Named kernel handles */
 }

src/greenlantern/ext/kernels/ellipsoid.cl

@@ -66,24 +66,24 @@ float ellipsoid_transit_flux_locked_workgroup_body(int gid, int gsz,
     float2 alpha, float2 beta, ang_params_t p, float3 ax, float ds,
     float u1, float u2)
 {
-    float afac1 = ax.x * (beta.y * (-alpha.x * p.r3 + alpha.y * p.r2) + beta.x * p.r1);
-    float bfac1 = ax.y * (beta.y * ( alpha.x * p.r6 - alpha.y * p.r5) - beta.x * p.r4);
-    float cfac1 = ax.z * (beta.y * (-alpha.x * p.r9 + alpha.y * p.r8) + beta.x * p.r7);
+    float afac1 = ax.x * (beta.y * (alpha.x * p.r3 + alpha.y * p.r2) - beta.x * p.r1);
+    float bfac1 = ax.y * (beta.y * (alpha.x * p.r6 + alpha.y * p.r5) - beta.x * p.r4);
+    float cfac1 = ax.z * (beta.y * (alpha.x * p.r9 + alpha.y * p.r8) - beta.x * p.r7);
 
-    float afac2 = ax.x * ( alpha.x * p.r2 + alpha.y * p.r3);
-    float bfac2 = ax.y * (-alpha.x * p.r5 - alpha.y * p.r6);
-    float cfac2 = ax.z * ( alpha.x * p.r8 + alpha.y * p.r9);
+    float afac2 = ax.x * (alpha.y * p.r3 - alpha.x * p.r2);
+    float bfac2 = ax.y * (alpha.y * p.r6 - alpha.x * p.r5);
+    float cfac2 = ax.z * (alpha.y * p.r9 - alpha.x * p.r8);
 
-    float abfac = ax.x * ax.y * (beta.y * p.r7 + beta.x * (alpha.x * p.r9 - alpha.y * p.r8));
-    float acfac = ax.x * ax.z * (beta.y * p.r4 + beta.x * (alpha.x * p.r6 - alpha.y * p.r5));
-    float bcfac = ax.y * ax.z * (beta.y * p.r1 + beta.x * (alpha.x * p.r3 - alpha.y * p.r2));
+    float abfac = ax.x * ax.y * (beta.y * p.r7 + beta.x * (alpha.x * p.r9 + alpha.y * p.r8));
+    float acfac = ax.x * ax.z * (beta.y * p.r4 + beta.x * (alpha.x * p.r6 + alpha.y * p.r5));
+    float bcfac = ax.y * ax.z * (beta.y * p.r1 + beta.x * (alpha.x * p.r3 + alpha.y * p.r2));
 
     float dx = sqrt(afac1 * afac1 + bfac1 * bfac1 + cfac1 * cfac1);
 
     float dy1 = (afac1 * afac2 + bfac1 * bfac2 + cfac1 * cfac2) / dx;
     float dy2 = sqrt(abfac * abfac + acfac * acfac + bcfac * bcfac) / dx;
 
-    float x0 = -ds * alpha.x * beta.y;
+    float x0 = ds * alpha.x * beta.y;
     float y0 = ds * alpha.y;
 
     float3 Ival = 0;
@@ -133,7 +133,7 @@ float ellipsoid_transit_flux_unlocked_workgroup_body(int gid, int gsz,
     float dy2 = sqrt(abfac * abfac + acfac * acfac + bcfac * bcfac) / dx;
 
     float x0 = ds * alpha.x * beta.y;
-    float y0 = -ds * alpha.y;
+    float y0 = ds * alpha.y;
 
     float3 Ival = 0;