1 files changed, 108 insertions, 0 deletions
diff --git a/source/blender/blenlib/intern/math_solvers.c b/source/blender/blenlib/intern/math_solvers.c
index 641e50c9bde..b8a22900ba1 100644
--- a/source/blender/blenlib/intern/math_solvers.c
+++ b/source/blender/blenlib/intern/math_solvers.c
@@ -72,3 +72,111 @@ void BLI_svd_m3(const float m3[3][3], float r_U[3][3], float r_S[3], float r_V[3
 {
 	EIG_svd_square_matrix(3, (const float *)m3, (float *)r_U, (float *)r_S, (float *)r_V);
 }
+
+/***************************** Simple Solvers ************************************/
+
+/**
+ * \brief Solve a tridiagonal system of equations:
+ *
+ * a[i] * r_x[i-1] + b[i] * r_x[i] + c[i] * r_x[i+1] = d[i]
+ *
+ * Ignores a[0] and c[count-1]. Uses the Thomas algorithm, e.g. see wiki.
+ *
+ * \param r_x output vector, may be shared with any of the input ones
+ * \return true if success
+ */
+bool BLI_tridiagonal_solve(const float *a, const float *b, const float *c, const float *d, float *r_x, const int count)
+{
+	if (count < 1)
+		return false;
+
+	size_t bytes = sizeof(double) * (unsigned)count;
+	double *c1 = (double *)MEM_mallocN(bytes * 2, "tridiagonal_c1d1");
+	double *d1 = c1 + count;
+
+	if (!c1)
+		return false;
+
+	int i;
+	double c_prev, d_prev, x_prev;
+
+	/* forward pass */
+
+	c1[0] = c_prev = ((double)c[0]) / b[0];
+	d1[0] = d_prev = ((double)d[0]) / b[0];
+
+	for (i = 1; i < count; i++) {
+		double denum = b[i] - a[i] * c_prev;
+
+		c1[i] = c_prev = c[i] / denum;
+		d1[i] = d_prev = (d[i] - a[i] * d_prev) / denum;
+	}
+
+	/* back pass */
+
+	x_prev = d_prev;
+	r_x[--i] = ((float)x_prev);
+
+	while (--i >= 0) {
+		x_prev = d1[i] - c1[i] * x_prev;
+		r_x[i] = ((float)x_prev);
+	}
+
+	MEM_freeN(c1);
+
+	return isfinite(x_prev);
+}
+
+/**
+ * \brief Solve a possibly cyclic tridiagonal system using the Sherman-Morrison formula.
+ *
+ * \param r_x output vector, may be shared with any of the input ones
+ * \return true if success
+ */
+bool BLI_tridiagonal_solve_cyclic(const float *a, const float *b, const float *c, const float *d, float *r_x, const int count)
+{
+	if (count < 1)
+		return false;
+
+	float a0 = a[0], cN = c[count - 1];
+
+	/* if not really cyclic, fall back to the simple solver */
+	if (a0 == 0.0f && cN == 0.0f) {
+		return BLI_tridiagonal_solve(a, b, c, d, r_x, count);
+	}
+
+	size_t bytes = sizeof(float) * (unsigned)count;
+	float *tmp = (float *)MEM_mallocN(bytes * 2, "tridiagonal_ex");
+	float *b2 = tmp + count;
+
+	if (!tmp)
+		return false;
+
+	/* prepare the noncyclic system; relies on tridiagonal_solve ignoring values */
+	memcpy(b2, b, bytes);
+	b2[0] -= a0;
+	b2[count - 1] -= cN;
+
+	memset(tmp, 0, bytes);
+	tmp[0] = a0;
+	tmp[count - 1] = cN;
+
+	/* solve for partial solution and adjustment vector */
+	bool success =
+		BLI_tridiagonal_solve(a, b2, c, tmp, tmp, count) &&
+		BLI_tridiagonal_solve(a, b2, c, d, r_x, count);
+
+	/* apply adjustment */
+	if (success) {
+		float coeff = (r_x[0] + r_x[count - 1]) / (1.0f + tmp[0] + tmp[count - 1]);
+
+		for (int i = 0; i < count; i++) {
+			r_x[i] -= coeff * tmp[i];
+		}
+	}
+
+	MEM_freeN(tmp);
+
+	return success;
+}
+