Bezier Curve Example in Slang

examples/bezier2d/README.md

@@ -0,0 +1,40 @@
+## Bezier Curve Example 
+
+This example has primitives to compute Bezier Curve Coefficients given $N$ control points. Refer to `bezier.slang` file for the slang kernels. Along with computing Bezier Curve Coefficients, and interpolating along the curves, there's also primitives to compute SDF of the Bezier Curve, using implicitization.
+
+Refer to `notebook.ipynb` for an example of initializing a Bezier Curve with 4 control points, and computing SDF. Example Images from the notebook shown below. 
+
+![alt text](assets/image.png)
+
+![alt text](assets/image-1.png)
+
+There are 2 toy application examples built on top of the Bezier Curve Primitive
+
+### Nearest point to Bezier Curve from an initialization using SDF based tracing.
+
+A Bezier Curve is initialized with random control points. Then from a random 2D point as an initialization, the point closest to the Bezier curve is reached by minimizing the absolute value of the SDF. `sdf_descent.py` runs the code for this, and has the initialization setup inside. On running the code, an image with the name `sdf_descent_{N}pts.png` will be generated in the same folder. $N$ is the number of control points, specified inside the file. An example of the image is shown below - 
+
+![SDF Descent Example](assets/sdf_descent_6pts.png)
+
+You can run the file as 
+
+```
+python3 sdf_descent.py 
+```
+
+Parameters such as number of control points, learning rate etc can be changed inside the file.
+
+
+### Fitting Bezier Curves to Arbitrary shapes 
+
+We show that we can optimize the locations of the control points of Bezier curves to fit to arbitrary parameteric shapes. Refer to `bezier_curvefit.py`. In the file, we have 3 Shapes - `HEART`, `ELLIPSE` and `ASTRID`. The image below shows fitting a heart shape, starting from a randomly initialized Bezier Curve. In the image below, the predicted curve is intentionally scaled so that the ground truth and the prediction are clearly visible. The fit is actually perfect, and makes it hard to see both the curves overlayed separately on the plot. 
+
+![Initialization](assets/init.png)
+![Heart Fit](assets/control_pts_descent.png)
+
+You can run the file as 
+
+```
+python3 bezier_curvefit.py 
+```
+On running the file as is, the optimization loss curve, initial Bezier Curve, and the final fit will be saved in a sub-folder `heart_20` in the current directory. This can be modified from inside the code, along with the number of control points, learning rate etc.

examples/bezier2d/bezier.slang

@@ -0,0 +1,128 @@
+static const int N = NUM_CTRL_PTS;
+static const int c = DIM;
+static const int N1 = c * (N- 1);
+
+int nCi(int n, int i) {
+    if (i > n) return 0;
+    if (i == 0 || i == n) return 1;
+    if (i > n - i) i = n - i;
+
+    int result = 1;
+    for (int k = 1; k <= i; ++k) {
+        result *= n - k + 1;
+        result /= k;
+    }
+
+    return result;
+}
+
+[PreferRecompute]
+int fact(int n) {
+    int result = 1;
+    for (int i = 1; i <= n; ++i) {
+        result *= i;
+    }
+    return result;
+}
+
+/*
+ * We bottleneck the component calculation through a single function tagged as [PreferRecompute]
+ * to avoid interediate memory allocations for contents of loops.
+ */
+[PreferRecompute]
+[Differentiable]
+float calc_component(uint i, uint j, uint k, DiffTensorView control_pts) {
+    return pow(-1, i + j) * control_pts[i, k] / (fact(i) * fact(j - i));
+}
+
+
+// Function to assemble matrix to compute determinant of to compute SDF.
+[Differentiable]
+void asm_mat(uint index, DiffTensorView output, matrix<float, N, c> coeffs) {
+    /** Function to create the matrix whose determinant is to be evaluated to get the sdf
+    @param coeffs: Tensor (N,c)
+    **/
+
+    for (int i = 0; i < N - 1; i++)
+        for (int j = 0; j < N; j++)
+            for (int k = 0; k < c; k++)
+            {
+                output.storeOnce(uint3(index, (k * (N - 1) + i), j + i), coeffs[j][k]);
+            }
+}
+
+[AutoPyBindCUDA]
+[CUDAKernel]
+[Differentiable]
+void bezier2D(DiffTensorView t, DiffTensorView control_pts, DiffTensorView output) {
+    /** @param t (tensor Mx1) : indices between 0-1 to traverse across the Bezier curve
+    **  @param control_pts (Nx2): N - Degree of Bezier Curve 2D
+    */
+    uint3 tIdx = cudaThreadIdx() + cudaBlockIdx() * cudaBlockDim();
+
+    // If the thread index is beyond the input size, exit early.
+    if (tIdx.x > t.size(0))
+        return;
+    [ForceUnroll]
+    for (int i = 0; i <= N - 1; i++)
+    {
+        output[tIdx.x, 0] = output[tIdx.x, 0] + nCi(N - 1, i) * pow((1 - t[tIdx.x]), (N - 1 - i)) * pow(t[tIdx.x], i) * control_pts[i, 0];
+        output[tIdx.x, 1] = output[tIdx.x, 1] + nCi(N - 1, i) * pow((1 - t[tIdx.x]), (N - 1 - i)) * pow(t[tIdx.x], i) * control_pts[i, 1];
+    }
+}
+
+[AutoPyBindCUDA]
+[CUDAKernel]
+[Differentiable]
+void bezier2DSDF(DiffTensorView xy, DiffTensorView bcoeffs, DiffTensorView output) {
+    /** @param xy  - M,c
+    **  @param bcoeffs - N,c
+    **  @return output - M, N1, N1  matrix for each point at which SDF is to be evaluated
+    ** Each thread computes the SDF value for a given xy coordinate from the determinant function above. Maybe change it up to be just differentiable, and not AutoPyBindCUDA
+    */
+    uint3 tIdx = cudaThreadIdx() + cudaBlockIdx() * cudaBlockDim();
+    matrix<float, N, c> coeffs; 
+
+    // copying coefficients to a separate variable for each thread.
+    for (int i = 0; i < N; i++)
+        for (int j = 0; j < c; j++)
+            coeffs[i][j] = bcoeffs[i, j];
+
+    int M = xy.size(0); // xy - shaped M,2
+    if (tIdx.x > M) {
+        return;
+    }
+
+    float coord[c];
+    [ForceUnroll]
+    for (int i = 0; i < c; i++)
+        coord[i] = xy[tIdx.x, i];
+
+    [ForceUnroll]
+    for (int i = 0; i < c; i++)
+        coeffs[0][i] -= coord[i];
+
+    asm_mat(tIdx.x, output, coeffs);
+}
+
+[AutoPyBindCUDA]
+[CudaKernel]
+[Differentiable]
+void compute_coeffs(DiffTensorView control_pts, DiffTensorView output) {
+    // Compute the coefficients a_i for t^i, for bezier polynomial \sum a_i . t^i
+    for (int k = 0; k < c; k++)
+    {
+        for (int j = 0; j < N; j++)
+        {
+            int nCj = fact(N - 1) / fact(N - 1 - j); // degree of the polynomial is N-1
+            float sum = 0;
+
+            for (int i = 0; i < N; i++)
+            {
+                if (i <= j)
+                    sum += calc_component(i, j, k, control_pts);
+            }
+            output.storeOnce(uint2(j, k), nCj * sum);
+        }
+    }
+}

examples/bezier2d/bezier_curvefit.py

@@ -0,0 +1,126 @@
+## Fit Bezier Curve to Heart Shaped Equation
+import torch
+import slangpy
+import os 
+import matplotlib.pyplot as plt
+
+N = 20
+c = 2 
+m = slangpy.loadModule('bezier.slang', defines={"NUM_CTRL_PTS": N, "DIM":c})
+
+
+def heart(t):
+    t = t*2*torch.pi
+    x = 16*(torch.sin(t))**3
+    y = 13*torch.cos(t) - 5*torch.cos(2*t) -2*torch.cos(3*t) - torch.cos(4*t)
+    return torch.hstack([x.reshape(-1,1),y.reshape(-1,1)])
+
+def ellipse(t, a, b):
+    t = t*2*torch.pi
+    x = a * (torch.cos(t))
+    y = b * (torch.sin(t))
+    return torch.hstack([x.reshape(-1,1),y.reshape(-1,1)])
+
+def astrid(t, a):
+    t = t*2*torch.pi
+    x = a * (torch.cos(t))**3
+    y = a * (torch.sin(t))**3
+    return torch.hstack([x.reshape(-1,1),y.reshape(-1,1)])
+
+def curve_from_coeffs(t, coeffs):
+    """ To check if coefficients are correct """
+    output = torch.zeros(t.shape[0], coeffs.shape[1]).cuda()
+    for i in range(coeffs.shape[0]):
+        output = output + (t**i).view(-1,1) * coeffs[i].view(1,-1)
+    return output 
+
+class Bezier2D(torch.autograd.Function):
+	@staticmethod
+	def forward(ctx, t, control_pts):
+		"""
+		t: M,1 (torch.tensor) on GPU, parameter for bezier curves
+		control_pts: N,2 (torch.tensor) 
+		"""
+		outputs = torch.zeros(t.shape[0], control_pts.shape[1]).cuda()
+		kernel_with_args = m.bezier2D(t=t, control_pts=control_pts, output=outputs)
+		NUM_BLOCKS = 1 + t.shape[0] // 1024
+		kernel_with_args.launchRaw(
+			blockSize=(NUM_BLOCKS, 1, 1),
+			gridSize=(1024, 1, 1))
+		ctx.save_for_backward(t, control_pts, outputs)
+		return outputs
+
+	@staticmethod
+	def backward(ctx, grad_outputs):
+		(t, control_pts, outputs) = ctx.saved_tensors
+		grad_ctrl_pts = torch.zeros_like(control_pts).cuda()
+		grad_t  = torch.zeros_like(t).cuda()
+  		# Note: When using DiffTensorView, grad_output gets 'consumed' during the reverse-mode.
+		# If grad_output may be reused, consider calling grad_output = grad_output.clone()
+
+		kernel_with_args = m.bezier2D.bwd(t=(t, grad_t),
+                                                       control_pts=(control_pts, grad_ctrl_pts),
+                                                       output=(outputs, grad_outputs))
+		NUM_BLOCKS = 1 + t.shape[0] // 1024
+		kernel_with_args.launchRaw(
+			blockSize=(NUM_BLOCKS, 1, 1),
+			gridSize=(1024, 1, 1))
+
+		return grad_t, grad_ctrl_pts
+
+
+
+
+### Initializing Control points, and Target Curve
+num_pts = 100
+t = torch.linspace(0.0, 1, num_pts, dtype=torch.float).cuda()
+
+savedir =  "./heart_20"
+os.makedirs(savedir, exist_ok=True)
+# gt_pts = ellipse(t, 3.0, 4.0)
+# gt_pts = astrid(t, 3.0)
+gt_pts = heart(t)
+control_pts = torch.rand((N, 2), dtype=torch.float).cuda()
+control_pts.requires_grad_(True)
+
+
+### Experiment 1 - Learning control points to match heart
+# Define a custom parameter, for example, a single value parameter.
+opt_param = torch.nn.Parameter(control_pts)
+pts = Bezier2D.apply(t, opt_param)
+
+plt.figure()
+plt.plot(pts[:,0].detach().cpu().numpy()/0.9, pts[:,1].detach().cpu().numpy()/0.9, color='red',label='predicted')
+plt.title('Bezier Curve Initialization')
+plt.savefig(os.path.join(savedir, 'init.png'))
+
+# Use an optimizer, for example, SGD, and register the custom parameter with it.
+lr_init = 0.01
+optimizer = torch.optim.Adam([opt_param], lr=lr_init)
+
+loss_curve = []
+for epoch in range(10000):  # Assuming 10000 epochs
+    pts = Bezier2D.apply(t,  opt_param)
+    loss = ((torch.linalg.norm(pts  - gt_pts, dim=1))).mean()
+    optimizer.zero_grad()
+    for pg in optimizer.param_groups:
+        pg['lr'] = lr_init * 0.99
+    loss.backward()
+    optimizer.step()
+    loss_curve.append(loss.item())
+    print(f"Epoch {epoch+1}, Loss: {loss.item()}")
+
+plt.figure()
+pts = Bezier2D.apply(t,  opt_param)
+plt.plot(pts[:,0].detach().cpu().numpy()/0.95, pts[:,1].detach().cpu().numpy()/0.95, color='red',label='predicted')
+plt.plot(gt_pts[:,0].detach().cpu().numpy(), gt_pts[:,1].detach().cpu().numpy(), color='green',label='gt')
+plt.title('HEART')
+plt.legend(['Predicted', 'Ground Truth'])
+plt.savefig(os.path.join(savedir, 'control_pts_descent.png'))
+
+plt.figure()
+plt.plot(loss_curve)
+plt.title('Loss Curve')
+plt.xlabel('Iterations')
+plt.ylabel('Loss Value')
+plt.savefig(os.path.join(savedir, 'Loss_Curve.png'))