z1_controller/include/common/utilities/PyPlot.h

261 lines
7.4 KiB
C++

#ifndef PYPLOT_H
#define PYPLOT_H
#include <map>
#include <thread>
#include <vector>
#include <string>
#include <Eigen/Dense>
#include <iostream>
#include "common/utilities/matplotlibcpp.h"
#include "unitree_arm_sdk/timer.h"
namespace plt = matplotlibcpp;
// enum class PlotType{
// STATIC,
// ANIME
// };
struct Curve{
std::vector<double> x;
std::vector<double> y;
void printXY(double xRough, int pointNum){
for(int i(0); i<x.size(); ++i){
if(xRough < x[i]){
for(int j(0); j < pointNum; ++j){
std::cout << " X: " << x[i+j] << ", Y: " << y[i+j] << std::endl;
}
break;
}
}
}
};
struct Plot{
std::vector<Curve*> curves;
std::vector<std::string> labels;
std::string plotName;
int curveCount;
std::map<std::string, int> curveName2ID;
Plot(std::string name, int count, std::vector<std::string> labelVec)
:plotName(name), curveCount(count), labels(labelVec){
for(int i(0); i < count; ++i){
curveName2ID.insert(std::pair<std::string, int>(labels[i], i));
curves.push_back(new Curve());
}
}
~Plot(){
for(int i(0); i < curveCount; ++i){
delete curves[i];
}
}
double getX(long long startT){
return (double)(getSystemTime() - startT) * 1e-6;
}
void printXY(std::string curveName, double xRough, int pointNum){
std::cout << "[DEBUG] Plot: " << plotName << ", Curve: " << curveName << std::endl;
curves[curveName2ID[curveName]]->printXY(xRough, pointNum);
}
};
class PyPlot{
public:
PyPlot();
~PyPlot();
void addPlot(std::string plotName, int curveCount, std::vector<std::string> labelVec);
void addPlot(std::string plotName, int curveCount);
void showPlot(std::string plotName);
void showPlot(std::vector<std::string> plotNameVec);
void showPlotAll();
void printXY(std::string plotName, std::string curveName, double xRough, int pointNum = 1);
void addFrame(std::string plotName, double value);
void addFrame(std::string plotName, double x, double value);
template <typename T>
void addFrame(std::string plotName, T* valueArray);
template <typename T>
void addFrame(std::string plotName, double x, T* valueArray);
template <typename T>
void addFrame(std::string plotName, const Eigen::MatrixBase<T> &vec);
template <typename T>
void addFrame(std::string plotName, double x, const Eigen::MatrixBase<T> &vec);
template <typename T>
void addFrame(std::string plotName, const std::vector<T> &vec);
template <typename T>
void addFrame(std::string plotName, double x, const std::vector<T> &vec);
private:
void _checkStart();
int _plotCount = 0;
std::map<std::string, int> _plotName2ID;
std::vector< Plot* > _plots;
// std::thread _plotThread;
long long _pointNum;
Plot* _getPlotPtr(std::string plotName);
bool start;
long long startT;
};
inline PyPlot::PyPlot(){
start = false;
// _plotThread = std::thread(&PyPlot::_plotThreadFunc, this);
}
inline PyPlot::~PyPlot(){
for(int i(0); i < _plotCount; ++i){
delete _plots[i];
}
}
inline void PyPlot::_checkStart(){
if(!start){
start = true;
startT = getSystemTime();
}
}
inline void PyPlot::printXY(std::string plotName, std::string curveName, double xRough, int pointNum){
_plots[_plotName2ID[plotName]]->printXY(curveName, xRough, pointNum);
}
inline void PyPlot::addPlot(std::string plotName, int curveCount, std::vector<std::string> labelVec){
if(_plotName2ID.count(plotName) == 0){
_plotName2ID.insert(std::pair<std::string, int>(plotName, _plotCount));
++_plotCount;
_plots.push_back( new Plot(plotName, curveCount, labelVec) );
}else{
std::cout << "[ERROR] Already has same Plot: " << plotName << std::endl;
exit(-1);
}
}
inline void PyPlot::addPlot(std::string plotName, int curveCount){
std::vector<std::string> label;
for(int i(0); i < curveCount; ++i){
label.push_back(std::to_string(i+1));
}
addPlot(plotName, curveCount, label);
}
inline void PyPlot::showPlot(std::string plotName){
Plot* plot = _getPlotPtr(plotName);
plt::figure();
plt::title(plot->plotName);
for(int i(0); i < plot->curveCount; ++i){
plt::named_plot(plot->labels[i], plot->curves[i]->x, plot->curves[i]->y);
}
plt::legend();
plt::show();
}
inline void PyPlot::showPlot(std::vector<std::string> plotNameVec){
for(std::vector<std::string>::iterator itName = plotNameVec.begin(); itName != plotNameVec.end(); ++itName){
plt::figure();
Plot* plot = _plots[_plotName2ID[*itName]];
plt::title(plot->plotName);
for(int i(0); i < plot->curveCount; ++i){
plt::named_plot(plot->labels[i], plot->curves[i]->x, plot->curves[i]->y);
}
plt::legend();
}
plt::show();
}
inline void PyPlot::showPlotAll(){
for(int i(0); i < _plotCount; ++i){
plt::figure();
Plot* plot = _plots[i];
plt::title(plot->plotName);
for(int j(0); j < plot->curveCount; ++j){
plt::named_plot(plot->labels[j], plot->curves[j]->x, plot->curves[j]->y);
}
plt::legend();
}
plt::show();
}
inline Plot* PyPlot::_getPlotPtr(std::string plotName){
if(_plotName2ID.count(plotName) == 0){
std::cout << "[ERROR] Plot " << plotName << " does not exist" << std::endl;
exit(-1);
}else{
return _plots[_plotName2ID[plotName]];
}
}
inline void PyPlot::addFrame(std::string plotName, double value){
_checkStart();
Plot* plot = _getPlotPtr(plotName);
addFrame(plotName, plot->getX(startT), value);
}
inline void PyPlot::addFrame(std::string plotName, double x, double value){
Plot* plot = _getPlotPtr(plotName);
plot->curves[0]->x.push_back(x);
plot->curves[0]->y.push_back(value);
}
template <typename T>
inline void PyPlot::addFrame(std::string plotName, T* valueArray){
_checkStart();
Plot* plot = _getPlotPtr(plotName);
addFrame(plotName, plot->getX(startT), valueArray);
}
template <typename T>
inline void PyPlot::addFrame(std::string plotName, double x, T* valueArray){
Plot* plot = _getPlotPtr(plotName);
for(int i(0); i < plot->curveCount; ++i){
plot->curves[i]->x.push_back(x);
plot->curves[i]->y.push_back(valueArray[i]);
}
}
template <typename T>
inline void PyPlot::addFrame(std::string plotName, const Eigen::MatrixBase<T> &vec){
_checkStart();
Plot* plot = _getPlotPtr(plotName);
addFrame(plotName, plot->getX(startT), vec);
}
template <typename T>
inline void PyPlot::addFrame(std::string plotName, double x, const Eigen::MatrixBase<T> &vec){
Plot* plot = _getPlotPtr(plotName);
for(int i(0); i < plot->curveCount; ++i){
plot->curves[i]->x.push_back(x);
plot->curves[i]->y.push_back(vec(i));
}
}
template <typename T>
inline void PyPlot::addFrame(std::string plotName, const std::vector<T> &vec){
_checkStart();
Plot* plot = _getPlotPtr(plotName);
addFrame(plotName, plot->getX(startT), vec);
}
template <typename T>
inline void PyPlot::addFrame(std::string plotName, double x, const std::vector<T> &vec){
Plot* plot = _getPlotPtr(plotName);
for(int i(0); i < plot->curveCount; ++i){
plot->curves[i]->x.push_back(x);
plot->curves[i]->y.push_back(vec[i]);
}
}
#endif // PYPLOT_H