ENH: Export PSO postitions to csv

BrownBiomechanics · Apr 5, 2023 · 4c20f62 · 4c20f62
1 parent f40447f
commit 4c20f62
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 4 deletions.
diff --git a/libautoscoper/src/PSO.cpp b/libautoscoper/src/PSO.cpp
@@ -1,6 +1,6 @@
 #include "PSO.hpp"
 #include <iostream>
-#include <string>
+#include <fstream>
 
 
 // New Particle Swarm Optimization
@@ -17,6 +17,23 @@ float host_fitness_function(float x[])
   return (float)total;
 }
 
+void write_positions(const std::string& output_directory, float* positions, unsigned int epoch_num,double xyzypr[]) {
+    // Write out the positions to a csv file 
+    std::string filename = output_directory + "pso_epoch_" + std::to_string(epoch_num) + ".csv";
+    std::ofstream file(filename);
+    if (!file.is_open()) {
+		std::cerr << "Could not open file: " << filename << std::endl;
+		exit(1);
+	}
+    for (int i = 0; i < NUM_OF_PARTICLES; i++) {
+        for (int j = 0; j < NUM_OF_DIMENSIONS; j++) {
+			file << xyzypr[j] + positions[i*NUM_OF_DIMENSIONS + j] << ",";
+		}
+		file << std::endl;
+	}
+    file.close();
+}
+
 void intializeRandom()
 {
   if (char* randomSeed = std::getenv("Autoscoper_RANDOM_SEED")) {
@@ -46,7 +63,7 @@ float getRandomClamped()
   return (float)rand() / (float)RAND_MAX;
 }
 
-void pso(float *positions, float *velocities, float *pBests, float *gBest, unsigned int MAX_EPOCHS, unsigned int MAX_STALL)
+void pso(float *positions, float *velocities, float *pBests, float *gBest, unsigned int MAX_EPOCHS, unsigned int MAX_STALL, double xyzypr[], const std::string& output_directory)
 {
   int stall_iter = 0;
   float tempParticle1[NUM_OF_DIMENSIONS];
@@ -78,6 +95,9 @@ void pso(float *positions, float *velocities, float *pBests, float *gBest, unsig
       positions[i] += velocities[i];
     }
 
+    if (directory != "")
+        write_positions(directory, positions, counter, xyzypr);
+
     OMEGA = OMEGA * 0.9f;
 
     for (int i = 0; i < NUM_OF_PARTICLES*NUM_OF_DIMENSIONS; i += NUM_OF_DIMENSIONS)

diff --git a/libautoscoper/src/PSO.hpp b/libautoscoper/src/PSO.hpp
@@ -12,6 +12,8 @@
 #include <stdlib.h>
 #include <time.h>
 #include <math.h>
+#include <string>
+
 const int NUM_OF_PARTICLES = 100;
 const int NUM_OF_DIMENSIONS = 6;
 const float c1 = 1.5f;
@@ -22,8 +24,10 @@ double PSO_FUNC(double *P);
 
 void intializeRandom();
 
+void write_positions(const std::string& output_directory, float* positions, unsigned int epoch_num, double xyzypr[]);
+
 float getRandom(float low, float high);
 float getRandomClamped();
 float host_fitness_function(float x[]);
 
-void pso(float *positions, float *velocities, float *pBests, float *gBest, unsigned int MAX_EPOCHS, unsigned int MAX_STALL);
+void pso(float* positions, float* velocities, float* pBests, float* gBest, unsigned int MAX_EPOCHS, unsigned int MAX_STALL, double xyzypr[] = {}, const std::string& output_directory = "");
diff --git a/libautoscoper/src/Tracker.cpp b/libautoscoper/src/Tracker.cpp
@@ -49,6 +49,7 @@
 #include <sstream>
 #include <vector>
 #include <cmath>
+#include <filesystem>
 #ifndef M_PI
 #define M_PI 3.14159265358979323846
 #endif
@@ -74,6 +75,11 @@
 #include "CoordFrame.hpp"
 #include "PSO.hpp"
 
+#ifdef WIN32
+#define OS_SEP "\\"
+#else
+#define OS_SEP "/"
+#endif
 
 using namespace std;
 
@@ -442,7 +448,25 @@ void Tracker::optimize(int frame, int dFrame, int repeats, int opt_method, unsig
 
       clock_t cpu_begin = clock();
 
-      pso(positions, velocities, pBests, gBest, MAX_EPOCHS, MAX_STALL);
+      // check for the environment variable Autoscoper_PSO_MAIN_OUTPUT_DIR
+      if (char* dirname = getenv("Autoscoper_PSO_MAIN_OUTPUT_DIR")) {
+        string output_dir(dirname);
+        cerr << "Autoscoper_PSO_MAIN_OUTPUT_DIR is set to " << output_dir << endl;
+        output_dir += (OS_SEP + string("Volume_") + to_string(trial_.current_volume));
+        output_dir += (OS_SEP + string("frame_") + to_string(frame));
+        filesystem::path p(output_dir);
+        error_code ec;
+        if (!filesystem::is_directory(p) || !filesystem::exists(p)) {
+            filesystem::create_directories(p,ec);
+            if (ec) {
+                throw exception("Error creating directory " + output_dir + ": " + ec.message());
+			}
+        }
+        output_dir += OS_SEP;
+        pso(positions, velocities, pBests, gBest, MAX_EPOCHS, MAX_STALL,xyzypr,output_dir);
+      }
+      else
+          pso(positions, velocities, pBests, gBest, MAX_EPOCHS, MAX_STALL);
       //cuda_pso(positions, velocities, pBests, gBest, MAX_EPOCHS);
 
       clock_t cpu_end = clock();