1 files changed, 50 insertions, 0 deletions
diff --git a/modules/core/c++/equilibrium.hpp b/modules/core/c++/equilibrium.hpp
new file mode 100644
index 0000000..eb2b043
--- /dev/null
+++ b/modules/core/c++/equilibrium.hpp
@@ -0,0 +1,50 @@
+#pragma once
+
+#include "descriptor.hpp"
+
+namespace kel {
+namespace lbm {
+template<typename T, typename Descriptor>
+saw::data<sch::FixedArray<T, Descriptor::Q>> equilibrium(saw::data<sch::Scalar<T>> rho, saw::data<sch::Vector<T,Descriptor::D>> vel){
+	using dfi = df_info<T, Descriptor>;
+
+	saw::data<sch::FixedArray<T,Descriptor::Q>> eq;
+	// Brain broken, here's an owl
+	//   ^
+	//  0.0
+	//  / \
+	//  | |
+	//
+	//  Velocity * Velocity meaning || vel ||_2^2 or <vel,vel>_2
+	saw::data<T> vel_vel{0.0};
+	for(uint64_t j = 0u; j < Descriptor::D; ++j){
+		vel_vel = vel_vel + vel.at({{j}}) * vel.at({{j}});
+	}
+
+	/**
+	 * Calculate equilibrium
+	 */
+	for(uint64_t i = 0u; i < Descriptor::Q; ++i){
+		saw::data<T> vel_c{};
+		for(uint64_t j = 0u; j < Descriptor::D; ++j){
+			// <vel,c_i>_2
+			vel_c = vel_c + (vel.at({{j}}) * saw::data<T>{static_cast<saw::native_data_type<T>::type>(dfi::directions[i][j])});
+		}
+
+		auto vel_c_cs2 = vel_c * saw::data<T>{dfi::inv_cs2};
+
+		eq.at(i).set(
+			dfi::weights[i] * rho.at({}).get() * 
+			(
+			 1.0
+			 + vel_c_cs2.get()
+			 - dfi::inv_cs2 * 0.5 * vel_vel.get()
+			 + vel_c_cs2.get() * vel_c_cs2.get() * 0.5
+			)
+		);
+	}
+
+	return eq;
+}
+}
+}