30_gen_dataset_ellipsoid.cpp

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

#include "contact_solver_factory.hpp"
#include "distribution_uniform.hpp"
#include "igl_wrapper.hpp"
#include "model_linear_spring.hpp"
#include "particle.hpp"
#include "shape_ellipsoid.hpp"
#include "shape_plane.hpp"
#include "solver_gjk_pp.hpp"
#include "spherical_voronoi.hpp"
#include "utils_math.hpp"
#include <fstream>
#include <iostream>
#include <random>
#include <sstream>
#include <string>
using namespace netdem;
using namespace std;
 
void SaveDatasetEllipsoidDetection(int num_samples, double (*ds_inputs)[7],
                                   bool *ds_cnt_flags, string filename) {
  stringbuf buf;
  ostream os(&buf);
  int os_width = 24;
  os.setf(ios::scientific);
  os.precision(15);
  for (int i = 0; i < num_samples; i++) {
    for (int i_inputs = 0; i_inputs < 7; i_inputs++) {
      os.width(os_width);
      os << ds_inputs[i][i_inputs] << ", ";
    }
    os.width(3);
    os << (ds_cnt_flags[i] ? 1 : 0) << endl;
  }
 
  ofstream outfile;
  outfile.open(filename);
  if (!outfile) {
    cout << "cannot open file: " << filename << endl;
  }
  outfile << buf.str();
  outfile.close();
  cout << "data saved to: " << filename << endl;
}
 
void SaveDatasetEllipsoidResolution(int num_samples, double (*ds_inputs)[7],
                                    double (*ds_cnt_feats)[7],
                                    string filename) {
  stringbuf buf;
  ostream os(&buf);
  int os_width = 24;
  os.setf(ios::scientific);
  os.precision(15);
  for (int i = 0; i < num_samples; i++) {
    // skip the non-contact data
    if (ds_cnt_feats[i][0] < 0)
      continue;
 
    for (int i_inputs = 0; i_inputs < 7; i_inputs++) {
      os.width(os_width);
      os << ds_inputs[i][i_inputs] << ", ";
    }
    for (int i_outputs = 0; i_outputs < 6; i_outputs++) {
      os.width(os_width);
      os << ds_cnt_feats[i][i_outputs] << ", ";
    }
    os.width(os_width);
    os << ds_cnt_feats[i][6] << endl;
  }
 
  ofstream outfile;
  outfile.open(filename);
  if (!outfile) {
    cout << "cannot open file: " << filename << endl;
  }
  outfile << buf.str();
  outfile.close();
  cout << "data saved to: " << filename << endl;
}
 
void GenDatasetEllipsoid(int num_samples = 100) {
  // load particle
  Ellipsoid ellipsoid = Ellipsoid(1, 1, 2);
  ellipsoid.SetSize(1.0);
 
  Particle obj_p1 = Particle(&ellipsoid);
  Particle obj_p2 = Particle(&ellipsoid);
  cout << "particle created ... " << endl;
 
  // allocate memory
  double(*ds_inputs)[7] = new double[num_samples][7];
  bool *ds_cnt_flags = new bool[num_samples];
  double(*ds_cnt_feats)[7] = new double[num_samples][7];
 
  // random generator
  UniformDistribution uniform_dist(0.0, 1.0);
 
  // use spherical centroidal voronoi to sample uniform unit VecXT
  VecXT<Vec3d> vertices = SphericalVoronoi::Solve(1000, 10000, 1.0e-4);
  VecXT<Vec3i> facets;
  IGLWrapper::ConvexHull(vertices, &vertices, &facets);
 
  // random cases
  for (int trial = 0, i = 0; trial < num_samples * 100; trial++) {
    // random normal
    int id_facet = floor(uniform_dist.Get() * facets.size());
    auto vert_0 = vertices[facets[id_facet][0]];
    auto vert_1 = vertices[facets[id_facet][1]];
    auto vert_2 = vertices[facets[id_facet][2]];
 
    double u_vert = uniform_dist.Get();
    double v_vert = uniform_dist.Get() * (1 - u_vert);
    double w_vert = 1 - u_vert - v_vert;
 
    Vec3d dir_n;
    dir_n[0] = u_vert * vert_0[0] + v_vert * vert_1[0] + w_vert * vert_2[0];
    dir_n[1] = u_vert * vert_0[1] + v_vert * vert_1[1] + w_vert * vert_2[1];
    dir_n[2] = u_vert * vert_0[2] + v_vert * vert_1[2] + w_vert * vert_2[2];
    Math::Normalize(&dir_n);
 
    // random overlap
    double cnt_prob = uniform_dist.Get();
    double len_min = cnt_prob > 0.6 ? 0 : -ellipsoid.GetBoundSphereRadius();
    double len_max = cnt_prob > 0.6 ? 0.1 : 0;
    double len_n = len_min + uniform_dist.Get() * (len_max - len_min);
 
    // random rotation
    Vec4d quat;
    quat[0] = uniform_dist.Get() * 2.0 - 1.0;
    quat[1] = uniform_dist.Get() * 2.0 - 1.0;
    quat[2] = uniform_dist.Get() * 2.0 - 1.0;
    quat[3] = uniform_dist.Get() * 2.0 - 1.0;
    Math::Quaternion::Normalize(&quat);
 
    // calculate p2 position
    Vec3d sup_pos_p1, sup_dir_p1{-dir_n[0], -dir_n[1], -dir_n[2]};
    sup_pos_p1 = ellipsoid.SupportPoint(sup_dir_p1);
 
    Vec3d sup_pos_p2, sup_pos_p2_ref, sup_pos_p2_o;
    sup_pos_p2 = sup_pos_p1 + len_n * dir_n;
 
    Vec4d quat_conj = Math::Quaternion::Conjugate(quat);
    Vec3d sup_dir_p2_ref = Math::Rotate(dir_n, quat_conj);
    sup_pos_p2_ref = ellipsoid.SupportPoint(sup_dir_p2_ref);
    sup_pos_p2_o = Math::Rotate(sup_pos_p2_ref, quat);
 
    Vec3d pos = sup_pos_p2 - sup_pos_p2_o;
 
    // apply to particle
    obj_p2.SetPosition(pos[0], pos[1], pos[2]);
    obj_p2.SetQuaternion(quat[0], quat[1], quat[2], quat[3]);
 
    if (len_n < 0) {
      // skip the cases pruuned away in the narrow phase contact detection
      if (Math::NormL2(pos) > ellipsoid.GetBoundSphereRadius() * 2)
        continue;
 
      // use distance from particle centroid to plane and plane normal as inputs
      ds_inputs[i][0] = obj_p2.pos[0];
      ds_inputs[i][1] = obj_p2.pos[1];
      ds_inputs[i][2] = obj_p2.pos[2];
      ds_inputs[i][3] = obj_p2.quaternion[0];
      ds_inputs[i][4] = obj_p2.quaternion[1];
      ds_inputs[i][5] = obj_p2.quaternion[2];
      ds_inputs[i][6] = obj_p2.quaternion[3];
 
      ds_cnt_flags[i] = 0;
 
      // volume and directional cross-section area of intersection as outputs
      ds_cnt_feats[i][0] = -1;
      ds_cnt_feats[i][1] = -1;
      ds_cnt_feats[i][2] = -1;
      ds_cnt_feats[i][3] = -1;
      ds_cnt_feats[i][4] = -1;
      ds_cnt_feats[i][5] = -1;
      ds_cnt_feats[i][6] = -1;
 
      i++;
    } else {
      // use distance from particle centroid to plane and plane normal as inputs
      ds_inputs[i][0] = obj_p2.pos[0];
      ds_inputs[i][1] = obj_p2.pos[1];
      ds_inputs[i][2] = obj_p2.pos[2];
      ds_inputs[i][3] = obj_p2.quaternion[0];
      ds_inputs[i][4] = obj_p2.quaternion[1];
      ds_inputs[i][5] = obj_p2.quaternion[2];
      ds_inputs[i][6] = obj_p2.quaternion[3];
 
      ds_cnt_flags[i] = 1;
 
      // volume and directional cross-section area of intersection as outputs
      ds_cnt_feats[i][0] = len_n;
      ds_cnt_feats[i][1] = dir_n[0];
      ds_cnt_feats[i][2] = dir_n[1];
      ds_cnt_feats[i][3] = dir_n[2];
      ds_cnt_feats[i][4] = sup_pos_p1[0] + 0.5 * len_n * dir_n[0];
      ds_cnt_feats[i][5] = sup_pos_p1[1] + 0.5 * len_n * dir_n[1];
      ds_cnt_feats[i][6] = sup_pos_p1[2] + 0.5 * len_n * dir_n[2];
 
      // SolverGJKPP cnt_solver;
      // LinearSpring cnt_model;
      // cnt_solver.Init(&obj_p1, &obj_p2);
      // auto cnt = ContactPP(&obj_p1, &obj_p2);
      // cnt.SetCollisionModel(&cnt_model);
      // cnt_solver.Detect();
      // cnt_solver.ResolveInit(&cnt, 1.0e-4);
      // auto &cnt_geoms = cnt.collision_entries[0].cnt_geoms;
 
      // cout << cnt_geoms.len_n << ", " << cnt_geoms.dir_n[0] << ", "
      //      << cnt_geoms.dir_n[1] << ", " << cnt_geoms.dir_n[2] << ", "
      //      << cnt_geoms.pos[0] << ", " << cnt_geoms.pos[1] << ", "
      //      << cnt_geoms.pos[2] << endl;
 
      // cout << ds_cnt_feats[i][0] << ", " << ds_cnt_feats[i][1] << ", "
      //      << ds_cnt_feats[i][2] << ", " << ds_cnt_feats[i][3] << ", "
      //      << ds_cnt_feats[i][4] << ds_cnt_feats[i][5] << ", "
      //      << ds_cnt_feats[i][6] << endl;
 
      i++;
    }
 
    if (((i + 1) % 100) == 0) {
      cout << "number of samples: " << i + 1 << " ..." << endl;
    }
    if (i >= num_samples) {
      break;
    }
  }
 
  string root_dir = "local/examples/netdem/ann_models/ellipsoid_ellipsoid/";
  SaveDatasetEllipsoidDetection(num_samples, ds_inputs, ds_cnt_flags,
                                root_dir + "dataset_detection.txt");
  SaveDatasetEllipsoidResolution(num_samples, ds_inputs, ds_cnt_feats,
                                 root_dir + "dataset_resolution.txt");
 
  delete[] ds_inputs;
  delete[] ds_cnt_flags;
  delete[] ds_cnt_feats;
}