1 files changed, 112 insertions, 0 deletions
diff --git a/source/blender/blenlib/intern/math_solvers.c b/source/blender/blenlib/intern/math_solvers.c
index 641e50c9bde..5a300e556a3 100644
--- a/source/blender/blenlib/intern/math_solvers.c
+++ b/source/blender/blenlib/intern/math_solvers.c
@@ -72,3 +72,115 @@ 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] * x[i-1] + b[i] * x[i] + c[i] * x[i+1] = d[i]
+ *
+ * Ignores a[0] and c[count-1]. Uses Thomas algorithm, e.g. see wiki.
+ *
+ * \param 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 *x, const int count)
+{
+	if (count < 0)
+		return false;
+
+	size_t bytes = sizeof(float)*(unsigned)count;
+	float *c1 = (float *)MEM_mallocN(bytes*2, "tridiagonal_c1d1");
+	float *d1 = c1 + count;
+
+	if (!c1)
+		return false;
+
+	int i;
+	double c_prev, d_prev, x_prev;
+
+	/* forward pass */
+
+	c_prev = ((double)c[0]) / b[0];
+	d_prev = ((double)d[0]) / b[0];
+
+	c1[0] = ((float)c_prev);
+	d1[0] = ((float)d_prev);
+
+	for (i = 1; i < count; i++) {
+		double denum = b[i] - a[i]*c_prev;
+
+		c_prev = c[i] / denum;
+		d_prev = (d[i] - a[i]*d_prev) / denum;
+
+		c1[i] = ((float)c_prev);
+		d1[i] = ((float)d_prev);
+	}
+
+	/* back pass */
+
+	x_prev = d_prev;
+	x[--i] = ((float)x_prev);
+
+	while (--i >= 0) {
+		x_prev = d1[i] - c1[i] * x_prev;
+		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 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 *x, const int count)
+{
+	if (count < 1)
+		return false;
+
+	float a0 = a[0], cN = c[count-1];
+
+	if (a0 == 0.0f && cN == 0.0f) {
+		return BLI_tridiagonal_solve(a, b, c, d, 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, x, count);
+
+	/* apply adjustment */
+	if (success) {
+		float coeff = (x[0] + x[count-1]) / (1.0f + tmp[0] + tmp[count-1]);
+
+		for (int i = 0; i < count; i++)
+			x[i] -= coeff * tmp[i];
+	}
+
+	MEM_freeN(tmp);
+
+	return success;
+}
+