1 files changed, 158 insertions, 57 deletions

diff --git a/src/main.cpp b/src/main.cpp
index f169982..e1f6671 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -37,12 +37,18 @@ struct GState
 	float gain_re;
 	float gain_im;
 	bool autogain;
-	float time_start;
-	float time_len;
-	float bode_start;
-	float bode_len;
-	float rlocus_start;
-	float rlocus_len;
+	struct {
+		float time;
+		float bode;
+		float nyquist;
+		float rlocus;
+	} start;
+	struct {
+		float time;
+		float bode;
+		float nyquist;
+		float rlocus;
+	} len;
 	size_t npoles, nzeros;
 	ImPlotPoint poles[MAX_ZP], zeros[MAX_ZP];
 };
@@ -51,7 +57,9 @@ struct CState
 {
 	bool proper;
 	ct::TimeSeries time;
+	ct::LocusSeries nyquist;
 	ct::LocusSeries rlocus;
+	ct::TransferFn tf_cl;
 };
 
 
@@ -73,11 +81,38 @@ ImPlotPoint step_getter(int idx, void *data)
 	return ImPlotPoint(ts->time(idx), idx != 0);
 }
 
+// Global variable because std::function is trash
+int rlocus_curr_row = 0;
 ImPlotPoint rlocus_getter(int idx, void *data)
 {
-	// return ImPlotPoint(ls->out(row, idx).real(), ls->out(row, idx).imag());
+	ct::LocusSeries *ls = (ct::LocusSeries *) data;
+	return ImPlotPoint(ls->out(rlocus_curr_row, idx).real(),
+			ls->out(rlocus_curr_row, idx).imag());
 }
 
+ImPlotPoint nyquist_getter(int idx, void *data)
+{
+	ct::LocusSeries *ls = (ct::LocusSeries *) data;
+	return ImPlotPoint(ls->out(0, idx).real(), ls->out(0, idx).imag());
+}
+
+ImPlotPoint nyquist_getter_conj(int idx, void *data)
+{
+	ct::LocusSeries *ls = (ct::LocusSeries *) data;
+	return ImPlotPoint(ls->out(0, idx).real(), -ls->out(0, idx).imag());
+}
+
+ImPlotPoint bode_getter_ampl(int idx, void *data)
+{
+	ct::LocusSeries *ls = (ct::LocusSeries *) data;
+	return ImPlotPoint(ls->in(idx), std::abs(ls->out(0, idx)));
+}
+
+ImPlotPoint bode_getter_phase(int idx, void *data)
+{
+	ct::LocusSeries *ls = (ct::LocusSeries *) data;
+	return ImPlotPoint(ls->in(idx), std::arg(ls->out(0, idx)));
+}
 
 static void glfw_error_callback(int error, const char* description)
 {
@@ -152,12 +187,18 @@ int main(int, char**)
 		.gain_re = 1.,
 		.gain_im = 0.,
 		.autogain = false,
-		.time_start = 0,
-		.time_len = 10,
-		.bode_start = 0.01,
-		.bode_len = 10000,
-		.rlocus_start = 0.001,
-		.rlocus_len = 1000,
+		.start = {
+			.time = 0.0,
+			.bode = 0.001,
+			.nyquist = 0.001,
+			.rlocus = 0.001,
+		},
+		.len = {
+			.time = 10.,
+			.bode = 1000.,
+			.nyquist = 1000.,
+			.rlocus = 1000.,
+		},
 		.npoles = 0,
 		.nzeros = 0,
 	};
@@ -167,10 +208,12 @@ int main(int, char**)
 	// Computation state
 	CState cstate = {
 		.proper = false,
-		.time = ct::TimeSeries(gstate.time_start,
-				gstate.time_start + gstate.time_len, gstate.n_samples.time),
-		.rlocus = ct::LocusSeries(gstate.rlocus_start,
-				gstate.rlocus_start + gstate.rlocus_len, gstate.n_samples.rlocus),
+		.time = ct::TimeSeries(gstate.start.time,
+				gstate.start.time + gstate.len.time, gstate.n_samples.time),
+		.nyquist = ct::LocusSeries(gstate.start.nyquist,
+				gstate.start.nyquist + gstate.len.nyquist, gstate.n_samples.nyquist),
+		.rlocus = ct::LocusSeries(gstate.start.rlocus,
+				gstate.start.rlocus + gstate.len.rlocus, gstate.n_samples.rlocus),
 	};
 
 	// Main loop
@@ -204,38 +247,41 @@ int main(int, char**)
 				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);
+
+				ct::complex gain(gstate.gain_re, gstate.gain_im);
 				// Autogain
 				if (gstate.autogain)
 				{
-					ct::complex gain = 1. / tf_cl.dc_gain();
-					tf_cl = gain * tf_cl;
+					gain = 1. / std::abs(tf.dc_gain());
 					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;
-				}
+
+				// tf = ct::cancel_zp(tf);
+				ct::TransferFn tf_cl = ct::feedback(tf, -gain);
+				cstate.tf_cl = tf_cl;
 
 				if (tf_cl.is_proper())
 				{
-					// Convert to state space
-					ct::SSModel ss = ct::ctrb_form(ct::feedback(tf));
-
 					// New time series
-					cstate.time = ct::TimeSeries(gstate.time_start,
-							gstate.time_start + gstate.time_len, gstate.n_samples.time);
+					cstate.time = ct::TimeSeries(gstate.start.time,
+							gstate.start.time + gstate.len.time, gstate.n_samples.time);
 
-					// Update time domain simulation
-					ct::step(ss, cstate.time);
+					// New root nyquist series
+					cstate.nyquist = ct::LocusSeries(gstate.start.nyquist,
+							gstate.start.nyquist + gstate.len.nyquist, gstate.n_samples.nyquist);
 
 					// New root locus series
-					cstate.rlocus = ct::LocusSeries(gstate.rlocus_start,
-							gstate.rlocus_start + gstate.rlocus_len, gstate.n_samples.rlocus);
+					cstate.rlocus = ct::LocusSeries(gstate.start.rlocus,
+							gstate.start.rlocus + gstate.len.rlocus, gstate.n_samples.rlocus);
+
+					// Convert to state space
+					ct::SSModel ss = ct::ctrb_form(ct::feedback(tf, -gain));
 
+					// Update plots
+					ct::step(ss, cstate.time);
 					ct::rlocus(tf, cstate.rlocus);
+					ct::nyquist(tf, cstate.nyquist);
 				}
 			}
 
@@ -262,8 +308,8 @@ int main(int, char**)
 			ImGui::SeparatorText("Plant G(s)");
 
 			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::SliderFloat("Gain (real)", &gstate.gain_re, 0.0001, 10000, NULL, ImGuiSliderFlags_Logarithmic);
+			ImGui::SliderFloat("Gain (imag)", &gstate.gain_im, 0.0001, 10000, NULL, ImGuiSliderFlags_Logarithmic);
 
 			ImGui::PushItemWidth(140);
 			if (ImGui::BeginListBox("Poles"))
@@ -301,12 +347,18 @@ int main(int, char**)
 			ImGui::PopItemWidth();
 
 			ImGui::SeparatorText("Simulation Ranges");
-			ImGui::InputFloat("Time Start", &gstate.time_start, 0, 100);
-			ImGui::SliderFloat("Time Length", &gstate.time_len, 0, 100, NULL, ImGuiSliderFlags_Logarithmic);
-			ImGui::InputFloat("Bode Start", &gstate.bode_start, 0, 100);
-			ImGui::SliderFloat("Bode Length", &gstate.bode_len, 0, 100, NULL, ImGuiSliderFlags_Logarithmic);
-			ImGui::InputFloat("Root Locus Start", &gstate.rlocus_start, 0, 100);
-			ImGui::SliderFloat("Root Locus Length", &gstate.rlocus_len, 0.001, 10000, NULL, ImGuiSliderFlags_Logarithmic);
+
+			ImGui::InputFloat("Time Start", &gstate.start.time, 0, 100);
+			ImGui::SliderFloat("Time Length", &gstate.len.time, 0, 100, NULL, ImGuiSliderFlags_Logarithmic);
+
+			ImGui::InputFloat("Bode Start", &gstate.start.bode, 0, 100);
+			ImGui::SliderFloat("Bode Length", &gstate.len.bode, 0, 100, NULL, ImGuiSliderFlags_Logarithmic);
+
+			ImGui::InputFloat("Nyquist Start", &gstate.start.nyquist, 0, 100);
+			ImGui::SliderFloat("Nyquist Length", &gstate.len.nyquist, 0.001, 10000, NULL, ImGuiSliderFlags_Logarithmic);
+
+			ImGui::InputFloat("Root Locus Start", &gstate.start.rlocus, 0, 100);
+			ImGui::SliderFloat("Root Locus Length", &gstate.len.rlocus, 0.001, 10000, NULL, ImGuiSliderFlags_Logarithmic);
 
 			ImGui::SeparatorText("Number of Samples");
 			ImGui::SliderInt("Time", &gstate.n_samples.time, 50, 5000);
@@ -322,18 +374,23 @@ int main(int, char**)
 			{
 				ImPlot::SetupAxisScale(ImAxis_X1, ImPlotScale_Log10);
 				ImPlot::SetupAxisScale(ImAxis_Y1, ImPlotScale_Log10);
-				ImPlot::SetupAxesLimits(0.01, 10000, 0.001, 10);
-				ImPlot::SetupAxis(ImAxis_X1, "Amplitude |G(iw)| dB");
-				ImPlot::SetupAxis(ImAxis_Y1, "Frequency w / (rad / s)");
+				// ImPlot::SetupAxisLimits(ImAxis_X1, gstate.start.bode,
+				// 		gstate.start.bode + gstate.len.bode, ImPlotCond_Always);
+				ImPlot::SetupAxis(ImAxis_X1, "Amplitude |G(iw)| dB", ImPlotAxisFlags_AutoFit);
+				ImPlot::SetupAxis(ImAxis_Y1, "Frequency w / (rad / s)", ImPlotAxisFlags_AutoFit);
+				ImPlot::PlotLineG("|G(iw)|", bode_getter_ampl, &cstate.nyquist, cstate.nyquist.n_samples);
 				ImPlot::EndPlot();
 			}
 
 			if (ImPlot::BeginPlot("Bode Phase", ImVec2(-1, 0)))
 			{
-				ImPlot::SetupAxisScale(ImAxis_X1, ImPlotScale_Linear);
-				ImPlot::SetupAxesLimits(0.1, 100, -180, 180);
-				ImPlot::SetupAxis(ImAxis_X1, "Phase arg G(iw) / deg");
-				ImPlot::SetupAxis(ImAxis_Y1, "Frequency w / (rad / s)");
+				ImPlot::SetupAxisScale(ImAxis_X1, ImPlotScale_Log10);
+				ImPlot::SetupAxisScale(ImAxis_Y1, ImPlotScale_Linear);
+				// ImPlot::SetupAxisLimits(ImAxis_X1, gstate.start.bode,
+				// 		gstate.start.bode + gstate.len.bode, ImPlotCond_Always);
+				ImPlot::SetupAxis(ImAxis_X1, "Phase arg G(iw) / deg", ImPlotAxisFlags_AutoFit);
+				ImPlot::SetupAxis(ImAxis_Y1, "Frequency w / (rad / s)", ImPlotAxisFlags_AutoFit);
+				ImPlot::PlotLineG("arg G(iw)", bode_getter_phase, &cstate.nyquist, cstate.nyquist.n_samples);
 				ImPlot::EndPlot();
 			}
 		}
@@ -345,8 +402,8 @@ int main(int, char**)
 			{
 				ImPlot::SetupAxisScale(ImAxis_X1, ImPlotScale_Linear);
 				ImPlot::SetupAxis(ImAxis_X1, "Time t / s");
-				ImPlot::SetupAxisLimits(ImAxis_X1, gstate.time_start,
-						gstate.time_start + gstate.time_len, ImPlotCond_Always);
+				ImPlot::SetupAxisLimits(ImAxis_X1, gstate.start.time,
+						gstate.start.time + gstate.len.time, ImPlotCond_Always);
 				ImPlot::SetupAxis(ImAxis_Y1, "Response y(t)", ImPlotAxisFlags_AutoFit);
 
 				ImPlot::PlotLineG("Step", step_getter, &cstate.time, cstate.time.n_samples);
@@ -360,8 +417,9 @@ int main(int, char**)
 		if (ImGui::Begin("Stability"))
 		{
 			ImGui::Text("Use CTRL + R / L Click to add a pole / zero");
+			ImGui::Text("Hold CTRL / SHIFT after dragging to keep it real / imaginary");
 
-			int rlflags = ImPlotFlags_Equal | ImPlotFlags_NoBoxSelect;
+			int rlflags = ImPlotFlags_Equal | ImPlotFlags_NoBoxSelect | ImPlotFlags_NoLegend;
 			if (ImGui::GetIO().KeyCtrl)
 				rlflags |= ImPlotFlags_Crosshairs;
 
@@ -373,7 +431,8 @@ int main(int, char**)
 				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);
+					rlocus_curr_row = i; // global variable!
+					ImPlot::PlotLineG("Root Locus", rlocus_getter, &cstate.rlocus, cstate.rlocus.n_samples);
 					ImGui::PopID();
 				}
 
@@ -404,16 +463,56 @@ int main(int, char**)
 
 				// Plot zeros and poles
 				const ImVec4 poleCol(0,0.9f,0,1);
+				const ImVec4 currPoleCol(0,0.5f,0,1);
 				const ImVec4 zeroCol(1,0.5f,1,1);
+				const ImVec4 currZeroCol(1,0.2f,1,1);
 				const float size = 4;
 				const int flags = ImPlotDragToolFlags_Delayed;
 
+				auto zeros = cstate.tf_cl.num.roots();
+				for (int i = 0; i < zeros.n_elem; i++)
+				{
+					const double x = zeros(i).real();
+					const double y = zeros(i).imag();
+					ImGui::PushID(i);
+					ImPlot::SetNextMarkerStyle(ImPlotMarker_Diamond, size, currZeroCol, 0);
+					ImPlot::PlotScatter("Zeros", &x, &y, 1);
+					ImGui::PopID();
+				}
+
+				auto roots = cstate.tf_cl.den.roots();
+				for (int i = 0; i < roots.n_elem; i++)
+				{
+					const double x = roots(i).real();
+					const double y = roots(i).imag();
+					ImGui::PushID(i);
+					ImPlot::SetNextMarkerStyle(ImPlotMarker_Diamond, size, currPoleCol, 0);
+					ImPlot::PlotScatter("Roots", &x, &y, 1);
+					ImGui::PopID();
+				}
+
 				for (int i = 0; i < gstate.nzeros; i++)
-					ImPlot::DragPoint(i, &gstate.zeros[i].x, &gstate.zeros[i].y, zeroCol, size, flags);
+				{
+					bool dragging = ImPlot::DragPoint(i, &gstate.zeros[i].x, &gstate.zeros[i].y, zeroCol, size, flags);
+
+					if (dragging && ImGui::GetIO().KeyCtrl)
+						gstate.zeros[i].y = 0;
+
+					if (dragging && ImGui::GetIO().KeyShift)
+						gstate.zeros[i].x = 0;
+				}
 
 				for (int i = 0; i < gstate.npoles; i++)
-					ImPlot::DragPoint(gstate.nzeros + i, &gstate.poles[i].x, &gstate.poles[i].y,
-							poleCol, size, flags);
+				{
+					bool dragging = ImPlot::DragPoint(gstate.nzeros + i, &gstate.poles[i].x,
+							&gstate.poles[i].y, poleCol, size, flags);
+
+					if (dragging && ImGui::GetIO().KeyCtrl)
+						gstate.poles[i].y = 0;
+
+					if (dragging && ImGui::GetIO().KeyShift)
+						gstate.poles[i].x = 0;
+				}
 
 				ImPlot::EndPlot();
 			}
@@ -421,7 +520,9 @@ int main(int, char**)
 			if (ImPlot::BeginPlot("Nyquist Diagram", ImVec2(-1, 0), ImPlotFlags_Equal))
 			{
 				ImPlot::SetupAxesLimits(-1, 1, -1, 1);
-				ImPlot::SetNextMarkerStyle(ImPlotMarker_Cross);
+				// ImPlot::SetNextMarkerStyle(ImPlotMarker_Plus);
+				ImPlot::PlotLineG("Nyquist Curve", nyquist_getter, &cstate.nyquist, cstate.nyquist.n_samples);
+				ImPlot::PlotLineG("Nyquist Curve Conj", nyquist_getter_conj, &cstate.nyquist, cstate.nyquist.n_samples);
 				ImPlot::EndPlot();
 			}
 		}