Sim.cpp

 #include "Sim.h"
#include <stdexcept>
#include <algorithm>
#include <cmath>
#include <limits>
#include <fstream>
#include <sstream>
#include <iomanip>

Sim::Sim() {
}

Sim::~Sim() {
}

void Sim::addVertex(Eigen::Vector3d pos) {
	Vert v;
	v.pos = pos;
	v.invMass = std::numeric_limits<float>::infinity();
	vels.push_back(Eigen::Vector3d::Zero());
	verts.push_back(v);
	vertexLastSegment.push_back(-1);
}

void Sim::addSegment(Eigen::Vector2i i, float kss, float density) {
	if (i(0) < 0 || i(0) >= verts.size() ||
		i(1) < 0 || i(1) >= verts.size())
		throw std::invalid_argument("Segment vertex index out of bounds");

	Seg s;
	s.i = i;

	Eigen::Vector3d dir = verts[i(1)].pos - verts[i(0)].pos;
	s.l = dir.norm();
	if (s.l <= 0) throw std::invalid_argument("Invalid segment length");

	s.q = Eigen::Quaterniond::FromTwoVectors(Eigen::Vector3d(1, 0, 0), dir.normalized());
	s.kss = kss * s.l;

	// Minimize twisting
	if (vertexLastSegment[i(0)] >= 0) {
		auto lastSeg = segs[vertexLastSegment[i(0)]];
		auto t = s.q.inverse() * lastSeg.q;
		if (t.x() != 0 || t.w() != 0) {
			// Create quaternion with only x and w components
			Eigen::Quaterniond tNorm(t.w(), t.x(), 0, 0);
			tNorm.normalize();
			s.q = s.q * tNorm;
		}
	}
	else if (vertexLastSegment[i(1)] >= 0) {
		auto lastSeg = segs[vertexLastSegment[i(1)]];
		auto t = s.q.inverse() * lastSeg.q;
		if (t.x() != 0 || t.w() != 0) {
			// Create quaternion with only x and w components
			Eigen::Quaterniond tNorm(t.w(), t.x(), 0, 0);
			tNorm.normalize();
			s.q = s.q * tNorm;
		}
	}

	vertexLastSegment[i(0)] = segs.size();
	vertexLastSegment[i(1)] = segs.size();

	segs.push_back(s);

	const float m = 0.5f * density * s.l;
	if (verts[i(0)].invMass > 0) verts[i(0)].invMass = 1 / (1 / verts[i(0)].invMass + m);
	if (verts[i(1)].invMass > 0) verts[i(1)].invMass = 1 / (1 / verts[i(1)].invMass + m);
}

void Sim::addRestAngle(Eigen::Vector2i i, float kbt) {
	if (i(0) < 0 || i(0) >= segs.size() ||
		i(1) < 0 || i(1) >= segs.size())
		throw std::invalid_argument("Angle segment index out of bounds");
	if (i(0) == i(1)) throw std::invalid_argument("Angle segment indices are the same");

	RestAngle ra;
	ra.i = i;
	ra.qRest = segs[i(0)].q.inverse() * segs[i(1)].q;
	ra.kbt = 4 * kbt / (0.5f * (segs[i(0)].l + segs[i(1)].l));
	restAngles.push_back(ra);
}

void Sim::pinVertex(int i) {
	if (i < 0 || i >= verts.size())
		throw std::invalid_argument("Vertex index out of bounds");
	verts[i].invMass = 0;
}

void Sim::pinSegment(int i) {
	if (i < 0 || i >= segs.size())
		throw std::invalid_argument("Segment index out of bounds");
	segs[i].kss = -abs(segs[i].kss);
}

void Sim::init() {
	// Initialize vertex-segment indices
	vector<int> counts(verts.size() + 1, 0);
	for (auto& seg : segs) {
		counts[seg.i(0)]++;
		counts[seg.i(1)]++;
	}

	int runningSum = 0;
	vertSegIndices.resize(verts.size() + 1);
	for (size_t i = 0; i < vertSegIndices.size(); i++) {
		vertSegIndices[i] = runningSum;
		runningSum += counts[i];
		counts[i] = 0;
	}

	vertSegIds.resize(runningSum);
	for (size_t i = 0; i < segs.size(); i++) {
		Eigen::Vector2i ind = segs[i].i;
		vertSegIds[vertSegIndices[ind(0)] + counts[ind(0)]++] = i;
		vertSegIds[vertSegIndices[ind(1)] + counts[ind(1)]++] = i;
	}

	// Initialize segment-rest angle indices
	counts.resize(segs.size() + 1);
	for (int& c : counts) c = 0;
	for (auto& ra : restAngles) {
		counts[ra.i(0)]++;
		counts[ra.i(1)]++;
	}

	runningSum = 0;
	segAngleIndices.resize(segs.size() + 1);
	for (size_t i = 0; i < segAngleIndices.size(); i++) {
		segAngleIndices[i] = runningSum;
		runningSum += counts[i];
		counts[i] = 0;
	}

	segAngleIds.resize(runningSum);
	for (size_t i = 0; i < restAngles.size(); i++) {
		Eigen::Vector2i ind = restAngles[i].i;
		segAngleIds[segAngleIndices[ind(0)] + counts[ind(0)]++] = i;
		segAngleIds[segAngleIndices[ind(1)] + counts[ind(1)]++] = i;
	}

	// Init memory
	vertexLastSegment.clear();
	lastVels.resize(verts.size(), Eigen::Vector3d::Zero());
	dxs.resize(verts.size(), Eigen::Vector3d::Zero());

	// Enable multithreading if there are enough segments
	if (segs.size() > COSSERAT_THREAD_BATCH * 2)
		multithreaded = true;
}

float Sim::advance(float h) {
	if (h <= 0) return 0;

	int steps = std::max(1, (int)std::ceil(h / maxH));
	meta.lastH = meta.h;
	meta.h = h / steps;

	for (int s = 0; s < steps; s++, stepCounter++) {
		startIterate();

		for (int itr = 0; itr < meta.numItr; itr++)
			iterate();
		endIterate();
		meta.time += meta.h;
	}

	return h;
}

void Sim::startIterate() {
	const float h = meta.h;
	const float invLastH = 1 / meta.lastH;
	const Eigen::Vector3d g = meta.gravity;

#pragma omp parallel for if(multithreaded) schedule(static, COSSERAT_THREAD_BATCH_STATIC)
	for (int tid = 0; tid < verts.size(); tid++) {
		auto vel = vels[tid];
		Eigen::Vector3d dx = vel * h;
		Eigen::Vector3d y = dx;

		float invMass = verts[tid].invMass;
		if (invMass > 0) {
			Eigen::Vector3d aExt = g;
			if (externalForce)
				aExt += externalForce(this, verts[tid], meta.time) * invMass;

			y += (h * h) * aExt;

			double ae2 = aExt.squaredNorm();
			if (ae2 > 0) {
				Eigen::Vector3d a = (vel - lastVels[tid]) * invLastH;
				float s = std::clamp((a.dot(aExt)) / ae2, 0.0, 1.0);
				dx += (h * h * s) * aExt;
			}
		}

		lastVels[tid] = vel;
		vels[tid] = y;
		dxs[tid] = dx;
	}
}

void Sim::endIterate() {
	const float h = meta.h;
	const float invH = 1 / h;

#pragma omp parallel for if(multithreaded) schedule(static, COSSERAT_THREAD_BATCH_STATIC)
	for (int tid = 0; tid < verts.size(); tid++) {
		Eigen::Vector3d dx = dxs[tid];
		if (verts[tid].invMass > 0)
			vels[tid] = dx * (invH * (1 - meta.drag * h));
		else
			vels[tid] = dx * invH;
		verts[tid].pos += dx;
	}
}

void Sim::iterate() {
	iterateVertVBD();
	iterateSegLambda();
}

void Sim::iterateVertVBD() {
	const float h = meta.h;
	const float damping = meta.damping / h;

#pragma omp parallel for if(multithreaded) schedule(dynamic, COSSERAT_THREAD_BATCH)
	for (int tid = 0; tid < verts.size(); tid++) {
		// Linear change
		auto v0 = verts[tid];

		if (v0.invMass > 0) {
			// Hessian H but we only track the diagonal
			float H = 1 / (v0.invMass * h * h);
			// vel has been overwritten to contain y - pos
			Eigen::Vector3d dx = dxs[tid];
			Eigen::Vector3d f = (1 / (h * h * v0.invMass)) * (vels[tid] - dx);

			// Cosserat stretching energy
			const int start = vertSegIndices[tid];
			const int end = vertSegIndices[tid + 1];
			for (int s = start; s < end; s++) {
				int segId = vertSegIds[s];
				auto seg = segs[segId];

				int oid = (seg.i(0) == tid) ? seg.i(1) : seg.i(0);
				auto p1 = verts[oid].pos;
				auto p1dx = dxs[oid];

				float invl = 1 / seg.l;
				float order = (seg.i(0) == tid) ? 1 : -1;
				Eigen::Vector3d c = ((p1 - v0.pos) + (p1dx - dx)) * (order * invl) - seg.q * Eigen::Vector3d(1, 0, 0);

				float k = abs(seg.kss) * invl;
				float d = k * invl;
				float damp = damping * d;
				f += order * k * c - damp * dx;
				H += d + damp;
			}

			// Local solve
			Eigen::Vector3d delta = f * (1 / H);
			dx += delta;

			dxs[tid] = dx;
		}
	}
}

void Sim::iterateSegLambda() {
#pragma omp parallel for if(multithreaded) schedule(dynamic, COSSERAT_THREAD_BATCH)
	for (int tid = 0; tid < segs.size(); tid++) {
		Seg seg = segs[tid];

		if (seg.kss < 0) continue;

		const Eigen::Vector3d v0 = verts[seg.i(0)].pos;
		const Eigen::Vector3d dx = dxs[seg.i(0)];
		const Eigen::Vector3d p1 = verts[seg.i(1)].pos;
		const Eigen::Vector3d p1dx = dxs[seg.i(1)];
		const Eigen::Vector3d v = (-2 * seg.kss / seg.l) * ((p1 - v0) + (p1dx - dx));

		Eigen::Vector4d b = Eigen::Vector4d::Zero();

		const int start = segAngleIndices[tid];
		const int end = segAngleIndices[tid + 1];
		for (int s = start; s < end; s++) {
			int angId = segAngleIds[s];
			RestAngle ang = restAngles[angId];

			const int oid = (ang.i(0) == tid) ? ang.i(1) : ang.i(0);
			const auto oq = segs[oid].q;
			auto qq = oq.inverse() * seg.q;
			if (ang.i(0) == tid) qq = qq.inverse();

			Eigen::Vector4d qq_vec(qq.x(), qq.y(), qq.z(), qq.w());
			Eigen::Vector4d qRest_vec(ang.qRest.x(), ang.qRest.y(), ang.qRest.z(), ang.qRest.w());
			float phi = (qq_vec.dot(qRest_vec) > 0) ? 1 : -1;

			if (ang.i(0) == tid) ang.qRest = ang.qRest.inverse();
			auto product = oq * ang.qRest;
			Eigen::Vector4d product_vec(product.x(), product.y(), product.z(), product.w());
			b += (ang.kbt * phi) * product_vec;
		}

		if (start == end) {
			Eigen::Vector3d from(1, 0, 0);
			Eigen::Vector3d to = (-v).normalized();
			seg.q = Eigen::Quaterniond::FromTwoVectors(from, to);
		} else {
			float lambda = v.norm() + b.norm();
			
			// Rotor(v): {v.x, v.y, v.z, 0}
			Eigen::Vector4d rotor_v(v.x(), v.y(), v.z(), 0);
			
			// Rotor(b): {b.x, b.y, b.z, b.w}  
			Eigen::Vector4d rotor_b = b;
			
			// Rotor(1): {1, 0, 0, 0}
			Eigen::Vector4d rotor_1(1, 0, 0, 0);
			
			// Manual quaternion multiplication for Rotor format {x,y,z,w}
			auto mult = [](const Eigen::Vector4d& a, const Eigen::Vector4d& b) -> Eigen::Vector4d {
				return Eigen::Vector4d(
					a.w() * b.x() + b.w() * a.x() + a.y() * b.z() - a.z() * b.y(),
					a.w() * b.y() + b.w() * a.y() + a.z() * b.x() - a.x() * b.z(),
					a.w() * b.z() + b.w() * a.z() + a.x() * b.y() - a.y() * b.x(),
					a.w() * b.w() - a.x() * b.x() - a.y() * b.y() - a.z() * b.z()
				);
			};
			
			auto combined = mult(mult(rotor_v, rotor_b), rotor_1);
			Eigen::Vector4d final_v = combined + lambda * b;
			
			final_v.normalize();
			seg.q = Eigen::Quaterniond(final_v.w(), final_v.x(), final_v.y(), final_v.z());
		}
		segs[tid].q = seg.q;
	}
}

void Sim::exportToVTK(const std::string& filename) const {
	std::ofstream file(filename);
	if (!file.is_open()) {
		throw std::runtime_error("Failed to open file for writing: " + filename);
	}

	// Write VTK header
	file << "# vtk DataFile Version 3.0\n";
	file << "Cosserat Rod Simulation\n";
	file << "ASCII\n";
	file << "DATASET POLYDATA\n";

	// Write vertices
	file << "POINTS " << verts.size() << " float\n";
	for (const auto& vert : verts) {
		file << std::fixed << std::setprecision(6) 
			 << vert.pos.x() << " " << vert.pos.y() << " " << vert.pos.z() << "\n";
	}

	// Write lines (segments)
	file << "LINES " << segs.size() << " " << (segs.size() * 3) << "\n";
	for (const auto& seg : segs) {
		file << "2 " << seg.i.x() << " " << seg.i.y() << "\n";
	}

	file.close();
}