1 files changed, 27 insertions, 16 deletions

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)
 				{