added flexibility

src/motions/base.h

@@ -27,6 +27,7 @@ public:
     [[nodiscard]] double getEta() const { return m_eta; }
     void setEta(const double eta) { m_eta = eta; }
     [[nodiscard]] std::string getName() const { return m_name; }
+    [[nodiscard]] double getInitOmega() const { return m_initial_omega; };
 
     static Motion* createFromInput(const std::string& input, std::mt19937_64& rng);
 
@@ -36,6 +37,7 @@ protected:
     double m_eta{0.};
     std::mt19937_64& m_rng;
     std::uniform_real_distribution<> m_uni_dist;
+    double m_initial_omega{0.};
 };
 
 std::ostream& operator<<(std::ostream& os, const Motion& m);

src/motions/bimodalangle.cpp

@@ -11,6 +11,7 @@ BimodalAngle::BimodalAngle(std::mt19937_64 &rng) : Motion(std::string("BimodalAn
 
 void BimodalAngle::initialize() {
     m_prev_pos = draw_position();
+    m_initial_omega = omega_q(m_prev_pos);
 };
 
 double BimodalAngle::jump() {

src/motions/isosmallangle.cpp

@@ -11,6 +11,7 @@ SmallAngle::SmallAngle(std::mt19937_64 &rng) : Motion(std::string("IsotropicAngl
 
 void SmallAngle::initialize() {
     m_prev_pos = draw_position();
+    m_initial_omega = omega_q(m_prev_pos);
 };
 
 double SmallAngle::jump() {

src/motions/random.cpp

@@ -6,7 +6,9 @@ RandomJump::RandomJump(const double delta, const double eta, std::mt19937_64 &rn
 
 RandomJump::RandomJump(std::mt19937_64 &rng) : Motion(std::string("RandomJump"), rng) {}
 
-void RandomJump::initialize() {}
+void RandomJump::initialize() {
+    m_initial_omega = RandomJump::jump();
+}
 
 double RandomJump::jump() {
     return omega_q(draw_position());

src/motions/tetrahedral.cpp

@@ -19,6 +19,7 @@ void TetrahedralJump::initialize() {
         auto corner_pos = rotate(pos, m_beta, alpha + (i-1) * 2*M_PI/3.);
         m_corners[i] = omega_q(corner_pos);
     }
+    m_initial_omega = TetrahedralJump::jump();
 }
 
 double TetrahedralJump::jump() {

src/simulation/sims.cpp

@@ -49,8 +49,9 @@ void run_spectrum(
     for (int mol_i = 0; mol_i < num_walker; mol_i++){
         std::vector<double> traj_time{};
         std::vector<double> traj_phase{};
+        std::vector<double> traj_omega{};
 
-        make_trajectory(motion, dist, tmax, traj_time, traj_phase);
+        make_trajectory(motion, dist, tmax, traj_time, traj_phase, traj_omega);
 
         for (auto& [t_echo_j, fid_j] : fid_dict) {
             // get phase at echo pulse
@@ -96,8 +97,10 @@ void run_ste(
     for (auto t_evo_i: evolution_times) {
         cc_dict[t_evo_i] = std::vector<double>(num_mix_times);
         ss_dict[t_evo_i] = std::vector<double>(num_mix_times);
+        std::fill(cc_dict[t_evo_i].begin(), cc_dict[t_evo_i].end(), 0.);
         std::fill(ss_dict[t_evo_i].begin(), ss_dict[t_evo_i].end(), 0.);
     }
+    std::vector<double> f2(num_mix_times);
 
     // each trajectory must have a duration of at least tmax
     const double tmax = *std::max_element(evolution_times.begin(), evolution_times.end()) * 2 + *std::max_element(mixing_times.begin(), mixing_times.end());
@@ -113,8 +116,16 @@ void run_ste(
     for (int mol_i = 0; mol_i < num_walker; mol_i++){
         std::vector<double> traj_time{};
         std::vector<double> traj_phase{};
+        std::vector<double> traj_omega{};
 
-        make_trajectory(motion, dist, tmax, traj_time, traj_phase);
+        make_trajectory(motion, dist, tmax, traj_time, traj_phase, traj_omega);
+
+        int f2_pos = 0;
+        for (int f2_idx=0; f2_idx < num_mix_times; f2_idx++) {
+            const double t_mix_f2 = mixing_times[f2_idx];
+            f2_pos = nearest_index(traj_time, t_mix_f2, f2_pos);
+            f2[f2_idx] += traj_omega[f2_pos] * motion.getInitOmega() / num_walker;
+        }
 
         for (auto& [t_evo_j, _] : cc_dict) {
             auto& cc_j = cc_dict[t_evo_j];
@@ -149,6 +160,7 @@ void run_ste(
     save_parameter_to_file("ste", motion.getName(), dist.getName(), parameter, optional);
     save_data_to_file("coscos", motion.getName(), dist.getName(), mixing_times, cc_dict, optional);
     save_data_to_file("sinsin", motion.getName(), dist.getName(), mixing_times, ss_dict, optional);
+    save_data_to_file("f2", motion.getName(), dist.getName(), mixing_times, f2, optional);
 
     printEnd(start);
 }
@@ -159,7 +171,8 @@ void make_trajectory(
     Distribution& dist,
     const double t_max,
     std::vector<double>& out_time,
-    std::vector<double>& out_phase
+    std::vector<double>& out_phase,
+    std::vector<double>& out_omega
     ) {
     // Starting position
     double t_passed = 0;
@@ -168,16 +181,21 @@ void make_trajectory(
     motion.initialize();
     dist.initialize();
 
+    double omega = motion.getInitOmega();
+
     out_time.emplace_back(t_passed);
-    out_phase.emplace_back(0);
+    out_phase.emplace_back(phase);
+    out_omega.emplace_back(omega);
 
     while (t_passed < t_max) {
         const double t = dist.tau_wait();
         t_passed += t;
-        phase += motion.jump() * t;
+        omega = motion.jump();
+        phase += omega * t;
 
         out_time.emplace_back(t_passed);
         out_phase.emplace_back(phase);
+        out_omega.emplace_back(omega);
     }
 }
 

src/simulation/sims.h

@@ -36,8 +36,9 @@ void run_ste(std::unordered_map<std::string, double>& parameter, std::unordered_
  * @param t_max Double that defines maximum time of trajectory
  * @param out_time Vector of waiting times
  * @param out_phase Vector of phase between waiting times
+ * @param out_omega Vector of omega at jump time
  */
-void make_trajectory(Motion& motion, Distribution& dist, double t_max, std::vector<double>& out_time, std::vector<double>& out_phase);
+void make_trajectory(Motion& motion, Distribution& dist, double t_max, std::vector<double>& out_time, std::vector<double>& out_phase, std::vector<double>& out_omega);
 
 std::chrono::system_clock::time_point printStart(std::unordered_map<std::string, double> &optional);
 void printEnd(std::chrono::system_clock::time_point start);

src/utils/io.cpp

@@ -193,3 +193,33 @@ void save_data_to_file(
         }
     }
 }
+
+void save_data_to_file(
+    const std::string& resulttype,
+    const std::string& motiontype,
+    const std::string& disttype,
+    const std::vector<double>& x,
+    const std::vector<double>& y,
+    std::unordered_map<std::string, double>& optional
+    ) {
+    // make file name
+    std::ostringstream datafile_name;
+    datafile_name << resulttype << "_" << motiontype << "_" << disttype;
+    datafile_name << std::setprecision(6) << std::scientific;
+    for (const auto& [key, value]: optional) {
+        datafile_name << "_" << key << "=" << value;
+    }
+    datafile_name << ".dat";
+
+    {
+        // write data to file, columns are secondary axis (echo delay, evolution times)
+        std::string datafile = datafile_name.str();
+        std::ofstream filestream(datafile, std::ios::out);
+
+        // write values to file
+        auto size = x.size();
+        for (unsigned int i = 0; i < size; i++) {
+            filestream << x[i] << "\t" << y[i] << "\n";
+        }
+    }
+}

src/utils/io.h

@@ -21,7 +21,9 @@ std::pair<std::string, double> get_optional_parameter(std::vector<std::string>::
 
 std::unordered_map<std::string, double> read_parameter(const std::filesystem::path&);
 
+
 void save_parameter_to_file(const std::string&, const std::string&, const std::string&, std::unordered_map<std::string, double>&, std::unordered_map<std::string, double>&);
 void save_data_to_file(const std::string&, const std::string&, const std::string&, const std::vector<double>&, const std::map<double, std::vector<double>>&, std::unordered_map<std::string, double>&);
+void save_data_to_file(const std::string&, const std::string&, const std::string&, const std::vector<double>&, const std::vector<double>&, std::unordered_map<std::string, double>&);
 
 #endif