Operators for TF

3 files changed, 167 insertions, 40 deletions

diff --git a/src/control.cpp b/src/control.cpp
index f57baec..4abc899 100644
--- a/src/control.cpp
+++ b/src/control.cpp
@@ -1,16 +1,30 @@
 #include "control.h"
 #include <armadillo>
+#include <cassert>
+#include <cmath>
 
 
 namespace ct
 {
-	TransferFn::TransferFn(complex gain)
-		: gain(gain)
-		, num(math::Poly<complex>{1})
-		, den(math::Poly<complex>{1})
+	TransferFn::TransferFn()
+		: num(1)
+		, den(1)
 	{
-		num(0) = 1.;
-		den(0) = 1.;
+		num(0) = 1;
+		den(0) = 1;
+	}
+
+	TransferFn::TransferFn(math::PolyCx num, math::PolyCx den)
+		: num(num)
+		, den(den)
+	{}
+
+	complex TransferFn::dc_gain() const
+	{
+		arma::cx_vec z(1, arma::fill::zeros);
+		arma::cx_vec n = arma::polyval(num.coeffs, z);
+		arma::cx_vec d = arma::polyval(den.coeffs, z);
+		return n(0) / d(0);
 	}
 
 	bool TransferFn::is_proper() const
@@ -23,9 +37,70 @@ namespace ct
 		return den.degree() >= num.degree();
 	}
 
-	TransferFn feedback(const TransferFn& tf, double k)
+	TransferFn operator + (const TransferFn& g, const TransferFn& h)
+	{
+		TransferFn tf(g.num * h.den + h.num * g.den, g.den * g.den);
+		tf.canonicalize();
+		return tf;
+	}
+
+	TransferFn operator - (const TransferFn& g, const TransferFn& h)
 	{
-		// TransferFn tf_cl;
+		TransferFn tf = g + (-1. * h);
+		tf.canonicalize();
+		return tf;
+	}
+
+	TransferFn operator * (const TransferFn& g, const TransferFn& h)
+	{
+		TransferFn tf(g.num * h.num, g.den * h.den);
+		tf.canonicalize();
+		return tf;
+	}
+
+	TransferFn operator / (const TransferFn& g, const TransferFn& h)
+	{
+		TransferFn hinv = TransferFn(h.den, h.num);
+		hinv.canonicalize();
+		TransferFn tf = g * hinv;
+		tf.canonicalize();
+		return tf;
+	}
+
+	TransferFn operator * (const complex k, const TransferFn& h)
+	{
+		TransferFn tf(k * h.num, h.den);
+		tf.canonicalize();
+		return tf;
+	}
+
+	TransferFn operator / (const TransferFn& h, const complex k)
+	{
+		TransferFn tf((1. / k) * h.num, h.den);
+		tf.canonicalize();
+		return tf;
+	}
+
+
+	TransferFn feedback(const TransferFn& tf, complex k)
+	{
+		const TransferFn unit(math::PolyCx({1.}), math::PolyCx({1.}));
+		const TransferFn ol = k * tf;
+		const TransferFn bo = unit - ol;
+		const TransferFn fb = unit / bo;
+		const TransferFn tf_cl = ol * fb;
+		// const TransferFn tf_cl = (k * tf) / (unit - k * tf);
+		return tf_cl;
+	}
+
+	void TransferFn::canonicalize()
+	{
+		complex an = den(0);
+		if (std::abs(an) > 0)
+		{
+			num.coeffs /= an;
+			den.coeffs /= an;
+		}
 	}
 
 	LocusSeries::LocusSeries(double start, double end, size_t nsamples)
@@ -77,7 +152,7 @@ namespace ct
 		{
 			ss.A(ord - 1, i) = tf.den(ord - i);
 			if (i < tf.num.coeffs.n_elem)
-				ss.C(i) = tf.num(i) * tf.gain;
+				ss.C(i) = tf.num(i);
 		}
 	
 		return ss;
diff --git a/src/control.h b/src/control.h
index 8ae9c5d..587b80e 100644
--- a/src/control.h
+++ b/src/control.h
@@ -1,9 +1,12 @@
 #pragma once
 
-#include <armadillo>
+#include "armadillo/include/armadillo"
+
+#include <utility>
 #include <complex>
 #include <cassert>
 #include <cmath>
+#include <cstddef>
 
 #ifndef CONTROL_H
 #define CONTROL_H
@@ -14,7 +17,8 @@
 
 namespace ct
 {
-	using complex = std::complex<double>;
+	typedef std::complex<double> complex;
+	typedef std::size_t size_t;
 
 	namespace math
 	{
@@ -23,11 +27,19 @@ namespace ct
 		{
 			arma::Col<T> coeffs;
 
-			Poly(size_t n_elem) : coeffs(n_elem) {};
+			Poly(size_t n_elem) : coeffs(n_elem) {}
 			Poly(arma::Col<T> c) : coeffs(c) {}
 
+			template<typename U>
+			Poly(std::initializer_list<U> l) : coeffs(l.size())
+			{
+				int i = 0;
+				for (const U& u : l)
+					coeffs(i++) = T(u);
+			}
+
 			int degree() const { return coeffs.n_elem - 1; }
-			void add_root(T r) { (*this) = (*this) * Poly({1, -r}); }
+			void add_root(T r) { (*this) = (*this) * Poly({1., -r}); }
 			arma::Col<T> roots() const { return arma::roots(coeffs); }
 
 			template<typename IndexT>
@@ -41,6 +53,7 @@ namespace ct
 		Poly<T> operator * (const Poly<T>& p, const Poly<T>& q)
 		{
 			arma::Col<T> coeffs = arma::conv(p.coeffs, q.coeffs);
+			CT_ASSERT(!coeffs.has_nan());
 			return Poly<T>(coeffs);
 		}
 
@@ -48,6 +61,15 @@ namespace ct
 		Poly<T> operator * (R scalar, const Poly<T>& p)
 		{
 			Poly<T> q(p.coeffs * scalar);
+			CT_ASSERT(!q.coeffs.has_nan());
+			return q;
+		}
+
+		template<typename T, typename R = T>
+		Poly<T> operator / (R scalar, const Poly<T>& p)
+		{
+			Poly<T> q(p.coeffs / scalar);
+			CT_ASSERT(!q.coeffs.has_nan());
 			return q;
 		}
 
@@ -57,34 +79,53 @@ namespace ct
 			const Poly<T>& big = (p.degree() > q.degree()) ? p : q;
 			const Poly<T>& small = (p.degree() > q.degree()) ? q : p;
 
-			Poly<T> s(big.degree() + 1);
-
-			for (int i = 0; i <= small.degree(); i++)
-				s(i) = big(i) + small(i);
+			Poly<T> s(small);
+			int rel = big.degree() - small.degree();
+			s.coeffs.insert_rows(0, rel);
 
-			for (int i = small.degree() + 1; i <= big.degree(); i++)
-				s(i) = big(i);
+			CT_ASSERT(!s.coeffs.has_nan());
+			return Poly<T>(big.coeffs + s.coeffs);
+		}
 
-			return s;
+		template<typename T>
+		Poly<T> operator - (const Poly<T>& p, const Poly<T>& q)
+		{
+			return p + (-1. * q);
 		}
+
+		typedef Poly<complex> PolyCx;
 	}
 
 	struct TransferFn
 	{
 		complex gain;
-		math::Poly<complex> num;
-		math::Poly<complex> den;
+		math::PolyCx num;
+		math::PolyCx den;
 
-		TransferFn(complex gain);
+		TransferFn(void);
+		TransferFn(math::PolyCx num, math::PolyCx den);
 
 		inline void add_pole(complex p) { den.add_root(p); }
 		inline void add_zero(complex z) { num.add_root(z); };
 
+		complex dc_gain() const;
+
 		bool is_proper() const;
 		bool is_strictly_proper() const;
+
+		void canonicalize();
 	};
 
-	TransferFn feedback(const TransferFn& tf, double k = -1);
+	TransferFn operator + (const TransferFn& g, const TransferFn& h);
+	TransferFn operator - (const TransferFn& g, const TransferFn& h);
+	TransferFn operator * (const TransferFn& g, const TransferFn& h);
+	TransferFn operator / (const TransferFn& g, const TransferFn& h);
+	TransferFn operator * (const complex k, const TransferFn& h);
+	TransferFn operator / (const TransferFn& h, const complex k);
+
+	// inline TransferFn operator / (const TransferFn&
+
+	TransferFn feedback(const TransferFn& tf, complex k = -1);
 
 	struct LocusSeries
 	{
diff --git a/src/main.cpp b/src/main.cpp
index 2d442b7..f169982 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -34,7 +34,8 @@ struct GState
 		int nyquist;
 		int rlocus;
 	} n_samples;
-	float gain;
+	float gain_re;
+	float gain_im;
 	bool autogain;
 	float time_start;
 	float time_len;
@@ -74,11 +75,7 @@ ImPlotPoint step_getter(int idx, void *data)
 
 ImPlotPoint rlocus_getter(int idx, void *data)
 {
-	ct::LocusSeries *ls = (ct::LocusSeries *) data;
-	// FIXME: this is a very ugly trick
-	idx = idx % ls->n_samples;
-	int row = (int) (idx / ls->n_samples);
-	return ImPlotPoint(ls->out(row, idx).real(), ls->out(row, idx).imag());
+	// return ImPlotPoint(ls->out(row, idx).real(), ls->out(row, idx).imag());
 }
 
 
@@ -152,7 +149,8 @@ int main(int, char**)
 			.nyquist = 500,
 			.rlocus = 2000,
 		},
-		.gain = 1.,
+		.gain_re = 1.,
+		.gain_im = 0.,
 		.autogain = false,
 		.time_start = 0,
 		.time_len = 10,
@@ -198,23 +196,33 @@ int main(int, char**)
 			if (gstate.npoles || gstate.nzeros)
 			{
 				// Create new transfer function
-				ct::TransferFn tf(gstate.gain);
+				ct::TransferFn tf;
+
 				for (int k = 0; k < gstate.npoles; k++)
 					tf.add_pole(ct::complex(gstate.poles[k].x, gstate.poles[k].y));
 
 				for (int k = 0; k < gstate.nzeros; k++)
 					tf.add_zero(ct::complex(gstate.zeros[k].x, gstate.zeros[k].y));
 
+				ct::TransferFn tf_cl = ct::feedback(tf);
 				// Autogain
 				if (gstate.autogain)
 				{
-					gstate.gain = abs(tf.den.coeffs.back() / tf.num.coeffs.back());
+					ct::complex gain = 1. / tf_cl.dc_gain();
+					tf_cl = gain * tf_cl;
+					gstate.gain_re = gain.real();
+					gstate.gain_im = gain.imag();
+				}
+				else
+				{
+					ct::complex gain(gstate.gain_re, gstate.gain_im);
+					tf_cl = gain * tf_cl;
 				}
 
-				if (tf.is_proper())
+				if (tf_cl.is_proper())
 				{
 					// Convert to state space
-					ct::SSModel ss = ct::ctrb_form(tf);
+					ct::SSModel ss = ct::ctrb_form(ct::feedback(tf));
 
 					// New time series
 					cstate.time = ct::TimeSeries(gstate.time_start,
@@ -253,9 +261,9 @@ int main(int, char**)
 
 			ImGui::SeparatorText("Plant G(s)");
 
-			ImGui::SliderFloat("Gain", &gstate.gain, -1000, 1000, NULL, ImGuiSliderFlags_Logarithmic);
-			ImGui::SameLine();
 			ImGui::Checkbox("Automatic", &gstate.autogain);
+			ImGui::SliderFloat("Gain (real)", &gstate.gain_re, -1000, 1000, NULL, ImGuiSliderFlags_Logarithmic);
+			ImGui::SliderFloat("Gain (imag)", &gstate.gain_im, -1000, 1000, NULL, ImGuiSliderFlags_Logarithmic);
 
 			ImGui::PushItemWidth(140);
 			if (ImGui::BeginListBox("Poles"))
@@ -362,9 +370,12 @@ int main(int, char**)
 				ImPlot::SetupAxesLimits(-10, 10, -10, 10);
 				
 				// Plot curren root locus
-				// FIXME: this is a very ugly trick
-				ImPlot::PlotLineG("Root Locus", rlocus_getter, &cstate.rlocus,
-						cstate.rlocus.n_samples * cstate.rlocus.out.n_rows);
+				for (int i = 0; i < cstate.rlocus.out.n_rows; i++)
+				{
+					ImGui::PushID(i);
+					// ImPlot::PlotLineG("Root Locus", rlocus_getter, &cstate.rlocus.out.row(i), cstate.rlocus.n_samples);
+					ImGui::PopID();
+				}
 
 				if (ImPlot::IsPlotHovered() && ImGui::GetIO().KeyCtrl)
 				{