21_test_dataset_ellipsoid_plane.cpp

This is an example of how to use the netdem library.

#include "contact_solver_factory.hpp"
#include "model_linear_spring.hpp"
#include "particle.hpp"
#include "shape_ellipsoid.hpp"
#include "shape_plane.hpp"
#include "solver_gjk_pw.hpp"
#include "utils_math.hpp"
#include <fstream>
#include <iostream>
#include <random>
#include <sstream>
#include <string>
 
using namespace netdem;
using namespace std;
 
VecXT<VecXT<double>> LoadDatasetEllipsoidPlane(string filename) {
  ifstream file_handle;
  file_handle.open(filename, ios::binary);
  if (!file_handle) {
    cout << "Fail to read file: " << filename << endl;
    return VecXT<VecXT<double>>();
  }
 
  file_handle.seekg(0, ios::end);
  long context_size = file_handle.tellg();
  char *buffer = new char[context_size];
  file_handle.seekg(0);
  file_handle.read(buffer, context_size);
  file_handle.close();
 
  stringstream ss;
  ss.str("");
  ss.clear();
  ss.str(buffer);
  VecXT<VecXT<double>> dataset;
  string str_line, str_num;
  VecXT<double> data_line;
 
  getline(ss, str_line);
  istringstream readstr(str_line);
  while (getline(readstr, str_num, ',')) {
    data_line.push_back(atof(str_num.c_str()));
  }
  getline(readstr, str_num);
  data_line.push_back(atof(str_num.c_str()));
  dataset.push_back(data_line);
  int num_data_per_line = data_line.size();
 
  while (getline(ss, str_line)) {
    if (str_line.find(",") == string::npos)
      continue;
 
    data_line.resize(num_data_per_line);
    istringstream readstr(str_line);
    for (int i = 0; i < num_data_per_line - 1; i++) {
      getline(readstr, str_num, ',');
      data_line[i] = atof(str_num.c_str());
    }
    getline(readstr, str_num);
    data_line[num_data_per_line - 1] = atof(str_num.c_str());
    dataset.push_back(data_line);
  }
  dataset.shrink_to_fit();
 
  delete[] buffer;
  return dataset;
}
 
void TestDatasetEllipsoidPlaneDetection() {
  // load particle
  Ellipsoid ellipsoid = Ellipsoid(1, 1, 2);
  ellipsoid.SetSize(1.0);
  Particle obj_p = Particle(&ellipsoid);
  cout << "particle created ... " << endl;
 
  // load wall
  Plane plane(0, 0, 0, 0, 0, 1);
  Wall obj_w = Wall(&plane);
  cout << "wall created ... " << endl;
 
  // solver
  SolverGJKPW cnt_solver;
  LinearSpring cnt_model;
 
  string root_dir = "local/examples/netdem/ann_models/ellipsoid_plane/";
  VecXT<VecXT<double>> dataset =
      LoadDatasetEllipsoidPlane(root_dir + "dataset_detection.txt");
 
  int num_samples = dataset.size();
  double(*ds_inputs)[4] = new double[num_samples][4];
  bool *ds_cnt_flags = new bool[num_samples];
 
  for (int i = 0, skip = 1; i < num_samples; i += skip) {
    for (int i_inputs = 0; i_inputs < 4; i_inputs++) {
      ds_inputs[i][i_inputs] = dataset[i][i_inputs];
    }
    ds_cnt_flags[i] = (bool)dataset[i][4];
  }
 
  for (int i = 0; i < num_samples; i++) {
    double dist_pc_to_plane;
    Vec3d dir_n;
 
    // unpack the inputs
    dist_pc_to_plane = ds_inputs[i][0];
    dir_n[0] = ds_inputs[i][1];
    dir_n[1] = ds_inputs[i][2];
    dir_n[2] = ds_inputs[i][3];
 
    // transform the inputs to the wall position and rotation
    // this is only for verification. No need to do this in contact solver
    Vec3d dir_n_ref{0, 0, 1};
    auto rot_axis = Math::Cross(dir_n_ref, dir_n);
 
    Vec4d quat;
    quat[0] = 1 + Math::Dot(dir_n, dir_n_ref);
    quat[1] = rot_axis[0];
    quat[2] = rot_axis[1];
    quat[3] = rot_axis[2];
    Math::Quaternion::Normalize(&quat);
 
    obj_w.SetPosition(-dist_pc_to_plane * dir_n[0],
                      -dist_pc_to_plane * dir_n[1],
                      -dist_pc_to_plane * dir_n[2]);
    obj_w.SetQuaternion(quat[0], quat[1], quat[2], quat[3]);
 
    cnt_solver.Init(&obj_p, &obj_w);
 
    bool cnt_flag = cnt_solver.Detect();
    cout << "true vs data: " << cnt_flag << ", " << ds_cnt_flags[i] << endl;
  }
 
  delete[] ds_inputs;
  delete[] ds_cnt_flags;
}
 
void TestDatasetEllipsoidPlaneResolution() {
  // load particle
  Ellipsoid ellipsoid = Ellipsoid(1, 1, 2);
  ellipsoid.SetSize(1.0);
  Particle obj_p = Particle(&ellipsoid);
  cout << "particle created ... " << endl;
 
  // load wall
  Plane plane(0, 0, 0, 0, 0, 1);
  Wall obj_w = Wall(&plane);
  cout << "wall created ... " << endl;
 
  // solver
  SolverGJKPW cnt_solver;
  LinearSpring cnt_model;
 
  string root_dir = "local/examples/netdem/ann_models/ellipsoid_plane/";
  VecXT<VecXT<double>> dataset =
      LoadDatasetEllipsoidPlane(root_dir + "dataset_resolution.txt");
 
  int num_samples = dataset.size();
  double(*ds_inputs)[4] = new double[num_samples][4];
  bool *ds_cnt_flags = new bool[num_samples];
  double(*ds_cnt_feats)[4] = new double[num_samples][4];
 
  for (int i = 0, skip = 1; i < num_samples; i += skip) {
    for (int i_inputs = 0; i_inputs < 4; i_inputs++) {
      ds_inputs[i][i_inputs] = dataset[i][i_inputs];
    }
    for (int i_outputs = 0; i_outputs < 4; i_outputs++) {
      ds_cnt_feats[i][i_outputs] = dataset[i][4 + i_outputs];
    }
  }
 
  for (int i = 0; i < num_samples; i++) {
    double dist_pc_to_plane;
    Vec3d dir_n;
 
    // unpack the inputs
    dist_pc_to_plane = ds_inputs[i][0];
    dir_n[0] = ds_inputs[i][1];
    dir_n[1] = ds_inputs[i][2];
    dir_n[2] = ds_inputs[i][3];
 
    // transform the inputs to the wall position and rotation
    // this is only for verification. No need to do this in contact solver
    Vec3d dir_n_ref{0, 0, 1};
    auto rot_axis = Math::Cross(dir_n_ref, dir_n);
    Math::Normalize(&rot_axis);
    auto rot_angle = acos(Math::Dot(dir_n, dir_n_ref));
    auto quat = Math::Quaternion::FromRodrigues(rot_angle, rot_axis);
 
    obj_w.SetPosition(-dist_pc_to_plane * dir_n[0],
                      -dist_pc_to_plane * dir_n[1],
                      -dist_pc_to_plane * dir_n[2]);
    obj_w.SetQuaternion(quat[0], quat[1], quat[2], quat[3]);
 
    cnt_solver.Init(&obj_p, &obj_w);
 
    // contact detection and resolution
    if (cnt_solver.Detect()) {
      ContactPW cnt = ContactPW(&obj_p, &obj_w);
      cnt.SetCollisionModel(&cnt_model);
      cnt_solver.ResolveInit(&cnt, 1.0e-4);
 
      auto &cnt_geoms = cnt.collision_entries[0].cnt_geoms;
      cout << ">> true: " << cnt_geoms.len_n << ", " << cnt_geoms.pos[0] << ", "
           << cnt_geoms.pos[1] << ", " << cnt_geoms.pos[2] << endl;
 
      cout << ">> data: " << ds_cnt_feats[i][0] << ", " << ds_cnt_feats[i][1]
           << ", " << ds_cnt_feats[i][2] << ", " << ds_cnt_feats[i][3] << endl;
    }
  }
 
  delete[] ds_inputs;
  delete[] ds_cnt_feats;
}
 
void TestDatasetEllipsoidPlane() {
  TestDatasetEllipsoidPlaneDetection();
  TestDatasetEllipsoidPlaneResolution();
}