Merge pull request #27 from lnls-fac/add-matrix-pm

Add matrix pass_method

Author: fernandohds564

include/trackcpp/auxiliary.h

@@ -34,7 +34,8 @@ class PassMethodsClass {
   static const int pm_thinquad_pass                  = 6;
   static const int pm_thinsext_pass                  = 7;
   static const int pm_kicktable_pass                 = 8;
-  static const int pm_nr_pms                         = 9;
+  static const int pm_matrix_pass                    = 9;
+  static const int pm_nr_pms                         = 10;
   PassMethodsClass() {
     passmethods.push_back("identity_pass");
     passmethods.push_back("drift_pass");
@@ -45,6 +46,7 @@ class PassMethodsClass {
     passmethods.push_back("thinquad_pass");
     passmethods.push_back("thinsext_pass");
     passmethods.push_back("kicktable_pass");
+    passmethods.push_back("matrix_pass");
   }
   int size() const { return passmethods.size(); }
   std::string operator[](const int i) const { return passmethods[i]; }
@@ -64,7 +66,8 @@ struct PassMethod {
         pm_thinquad_pass                  = 6,
         pm_thinsext_pass                  = 7,
         pm_kicktable_pass                 = 8,
-        pm_nr_pms                         = 9
+        pm_matrix_pass                    = 9,
+        pm_nr_pms                         = 10,
     };
 };
 
@@ -79,7 +82,8 @@ const std::vector<std::string> pm_dict = {
         "cavity_pass",
         "thinquad_pass",
         "thinsext_pass",
-        "kicktable_pass"
+        "kicktable_pass",
+        "matrix_pass"
 };
 
 struct Status {

include/trackcpp/elements.h

@@ -19,6 +19,7 @@
 
 #include "kicktable.h"
 #include "auxiliary.h"
+#include "linalg.h"
 #include <vector>
 #include <string>
 #include <fstream>
@@ -61,8 +62,9 @@ class Element {
 
   std::vector<double> polynom_a = default_polynom;
   std::vector<double> polynom_b = default_polynom;
+  Matrix              matrix66 = Matrix(6);
 
-  const Kicktable*    kicktable   = nullptr;
+  const Kicktable*    kicktable = nullptr;
 
   double              t_in[6],  t_out[6];
   double              r_in[36], r_out[36];
@@ -84,6 +86,7 @@ class Element {
   static Element vcorrector (const std::string& fam_name_, const double& length_, const double& vkick_);
   static Element corrector  (const std::string& fam_name_, const double& length_, const double& hkick_, const double& vkick_);
   static Element drift      (const std::string& fam_name_, const double& length_);
+  static Element matrix     (const std::string& fam_name_, const double& length_);
   static Element rbend      (const std::string& fam_name_, const double& length_,
                              const double& angle_, const double& angle_in_ = 0, const double& angle_out_ = 0,
                              const double& gap_ = 0, const double& fint_in_ = 0, const double& fint_out_ = 0,
@@ -103,6 +106,7 @@ class Element {
 void initialize_marker(Element& element);
 void initialize_corrector(Element& element, const double& hkick, const double& vkick);
 void initialize_drift(Element& element);
+void initialize_matrix(Element &element);
 void initialize_rbend(Element& element, const double& angle, const double& angle_in, const double& angle_out,
             const double& gap, const double& fint_in, const double& fint_out,
             const std::vector<double>& polynom_a, const std::vector<double>& polynom_b,

include/trackcpp/linalg.h

@@ -68,4 +68,6 @@ void matrix6_set_identity_posvec(std::vector<Pos<T> >& m, const T& a = 1) {
   m[0].rx = m[1].px = m[2].ry = m[3].py = m[4].de = m[5].dl = a;
 }
 
+void multiply_transf_matrix66(Matrix &m, const double k1);
+
 #endif

include/trackcpp/passmethods.h

@@ -32,6 +32,7 @@ template <typename T> Status::type pm_cavity_pass                (Pos<T> &pos, c
 template <typename T> Status::type pm_thinquad_pass              (Pos<T> &pos, const Element &elem, const Accelerator& accelerator);
 template <typename T> Status::type pm_thinsext_pass              (Pos<T> &pos, const Element &elem, const Accelerator& accelerator);
 template <typename T> Status::type pm_kicktable_pass             (Pos<T> &pos, const Element &elem, const Accelerator& accelerator);
+template <typename T> Status::type pm_matrix_pass                (Pos<T> &pos, const Element &elem, const Accelerator& accelerator);
 
 #include "passmethods.hpp"
 

include/trackcpp/passmethods.hpp

@@ -29,6 +29,7 @@
 #include "pos.h"
 #include "auxiliary.h"
 #include "tpsa.h"
+#include "linalg.h"
 #include <cmath>
 
 // constants for 4th-order symplectic integrator
@@ -121,6 +122,21 @@ Status::type kicktablethinkick(Pos<T>& pos, const Kicktable* kicktable,
   return status;
 }
 
+template <typename T>
+void matthinkick(Pos<T> &pos, const Matrix &m) {
+
+  T rx = pos.rx, px = pos.px;
+  T ry = pos.ry, py = pos.py;
+  T de = pos.de, dl = pos.dl;
+
+  pos.rx = m[0][0]*rx+m[0][1]*px+m[0][2]*ry+m[0][3]*py+m[0][4]*de+m[0][5]*dl;
+  pos.px = m[1][0]*rx+m[1][1]*px+m[1][2]*ry+m[1][3]*py+m[1][4]*de+m[1][5]*dl;
+  pos.ry = m[2][0]*rx+m[2][1]*px+m[2][2]*ry+m[2][3]*py+m[2][4]*de+m[2][5]*dl;
+  pos.py = m[3][0]*rx+m[3][1]*px+m[3][2]*ry+m[3][3]*py+m[3][4]*de+m[3][5]*dl;
+  pos.de = m[4][0]*rx+m[4][1]*px+m[4][2]*ry+m[4][3]*py+m[4][4]*de+m[4][5]*dl;
+  pos.dl = m[5][0]*rx+m[5][1]*px+m[5][2]*ry+m[5][3]*py+m[5][4]*de+m[5][5]*dl;
+}
+
 template <typename T>
 void strthinkick(Pos<T>& pos, const double& length,
                  const std::vector<double>& polynom_a,
@@ -230,7 +246,6 @@ void local_2_global(Pos<T> &pos, const Element &elem) {
   translate_pos(pos, elem.t_out);
 }
 
-
 template <typename T>
 Status::type pm_identity_pass(Pos<T> &pos, const Element &elem,
                               const Accelerator& accelerator) {
@@ -259,9 +274,9 @@ Status::type pm_str_mpole_symplectic4_pass(Pos<T> &pos, const Element &elem,
   const std::vector<double> &polynom_a = elem.polynom_a;
   const std::vector<double> &polynom_b = elem.polynom_b;
   for(unsigned int i=0; i<elem.nr_steps; ++i) {
-    drift(pos, l1);
+    drift<T>(pos, l1);
     strthinkick<T>(pos, k1, polynom_a, polynom_b, accelerator);
-    drift(pos, l2);
+    drift<T>(pos, l2);
     strthinkick<T>(pos, k2, polynom_a, polynom_b, accelerator);
     drift<T>(pos, l2);
     strthinkick<T>(pos, k1, polynom_a, polynom_b, accelerator);
@@ -301,7 +316,6 @@ Status::type pm_bnd_mpole_symplectic4_pass(Pos<T> &pos, const Element &elem,
   return Status::success;
 }
 
-
 template <typename T>
 Status::type pm_corrector_pass(Pos<T> &pos, const Element &elem,
                                const Accelerator& accelerator) {
@@ -333,7 +347,6 @@ Status::type pm_corrector_pass(Pos<T> &pos, const Element &elem,
   return Status::success;
 }
 
-
 template <typename T>
 Status::type pm_cavity_pass(Pos<T> &pos, const Element &elem,
                             const Accelerator& accelerator) {
@@ -410,6 +423,37 @@ Status::type pm_kicktable_pass(Pos<T> &pos, const Element &elem,
   return status;
 }
 
+template <typename T>
+Status::type pm_matrix_pass(Pos<T> &pos, const Element &elem,
+                            const Accelerator& accelerator) {
+
+  global_2_local(pos, elem);
+  if (elem.length > 0) {
+    double sl = elem.length / float(elem.nr_steps);
+    double l1 = sl * DRIFT1;
+    double l2 = sl * DRIFT2;
+    double k1 = KICK1 / float(elem.nr_steps);
+    double k2 = KICK2 / float(elem.nr_steps);
+    Matrix mat1 = elem.matrix66;
+    Matrix mat2 = elem.matrix66;
+    multiply_transf_matrix66(mat1, k1);
+    multiply_transf_matrix66(mat2, k2);
+    for(unsigned int i=0; i<elem.nr_steps; ++i) {
+      drift<T>(pos, l1);
+      matthinkick<T>(pos, mat1);
+      drift<T>(pos, l2);
+      matthinkick<T>(pos, mat2);
+      drift<T>(pos, l2);
+      matthinkick<T>(pos, mat1);
+      drift<T>(pos, l1);
+    }
+  } else {
+    matthinkick<T>(pos, elem.matrix66);
+  }
+  local_2_global(pos, elem);
+  return Status::success;
+}
+
 #undef DRIFT1
 #undef DRIFT2
 #undef KICK1

include/trackcpp/tracking.h

@@ -68,6 +68,9 @@ Status::type track_elementpass (
 	case PassMethod::pm_kicktable_pass:
 		if ((status = pm_kicktable_pass<T>(orig_pos, el, accelerator)) != Status::success) return status;
 		break;
+	case PassMethod::pm_matrix_pass:
+		if ((status = pm_matrix_pass<T>(orig_pos, el, accelerator)) != Status::success) return status;
+		break;
 	default:
 		return Status::passmethod_not_defined;
 	}

python_package/interface.cpp

@@ -81,6 +81,10 @@ Element drift_wrapper(const std::string &fam_name_, const double &length_) {
     return Element::drift(fam_name_, length_);
 }
 
+Element matrix_wrapper(const std::string &fam_name_, const double &length_) {
+    return Element::matrix(fam_name_, length_);
+}
+
 Element rbend_wrapper(const std::string& fam_name_, const double& length_,
 					  const double& angle_, const double& angle_in_, const double& angle_out_,
 					  const double& gap_, const double& fint_in_, const double& fint_out_,

python_package/interface.h

@@ -68,6 +68,7 @@ Element hcorrector_wrapper(const std::string& fam_name_, const double& length_,
 Element vcorrector_wrapper(const std::string& fam_name_, const double& length_, const double& vkick_);
 Element corrector_wrapper(const std::string& fam_name_, const double& length_, const double& hkick_, const double& vkick_);
 Element drift_wrapper(const std::string& fam_name_, const double& length_);
+Element matrix_wrapper(const std::string& fam_name_, const double& length_);
 Element quadrupole_wrapper(const std::string& fam_name_, const double& length_, const double& K_, const int nr_steps_ = 1);
 Element sextupole_wrapper(const std::string& fam_name_, const double& length_, const double& S_, const int nr_steps_ = 1);
 Element rfcavity_wrapper(const std::string& fam_name_, const double& length_, const double& frequency_, const double& voltage_, const double& phase_lag);

src/elements.cpp

@@ -34,7 +34,7 @@ Element::Element(const std::string& fam_name_, const double& length_) :
       }
     }
   }
-
+  matrix66.eye();
 }
 
 const std::string& Element::get_pass_method() {
@@ -69,6 +69,12 @@ Element Element::drift (const std::string& fam_name_, const double& length_) {
   return e;
 }
 
+Element Element::matrix(const std::string& fam_name_, const double& length_) {
+  Element e = Element(fam_name_, length_);
+    initialize_matrix(e);
+  return e;
+}
+
 Element Element::hcorrector(const std::string& fam_name_, const double& length_, const double& hkick_) {
   Element e = Element(fam_name_, length_);
     initialize_corrector(e, hkick_, 0.0);
@@ -160,6 +166,12 @@ bool Element::operator==(const Element& o) const {
     if (this->phase_lag != o.phase_lag) return false;
     if (this->polynom_a != o.polynom_a) return false;
     if (this->polynom_b != o.polynom_b) return false;
+    const Matrix& m = this->matrix66;
+    const Matrix& mo = o.matrix66;
+    for(unsigned int i=0; i<m.size(); ++i){
+      for(unsigned int j=0; j<m[i].size(); ++j)
+        if (m[i][j] != mo[i][j]) return false;
+    }
     for(unsigned int i=0; i<6; ++i) {
       if (this->t_in[i] != o.t_in[i]) return false;
       if (this->t_out[i] != o.t_out[i]) return false;
@@ -214,6 +226,10 @@ void initialize_drift(Element &element) {
     element.pass_method = PassMethod::pm_drift_pass;
 }
 
+void initialize_matrix(Element &element) {
+    element.pass_method = PassMethod::pm_matrix_pass;
+}
+
 void initialize_rbend(Element& element, const double& angle, const double& angle_in, const double& angle_out, const double& gap, const double& fint_in, const double& fint_out, const std::vector<double>& polynom_a, const std::vector<double>& polynom_b, const double& K, const double& S, const int nr_steps) {
     element.pass_method = PassMethod::pm_bnd_mpole_symplectic4_pass;
     element.angle = angle;

src/flat_file.cpp

@@ -17,6 +17,7 @@
 #include <trackcpp/flat_file.h>
 #include <trackcpp/elements.h>
 #include <trackcpp/auxiliary.h>
+#include <trackcpp/linalg.h>
 #include <fstream>
 #include <string>
 #include <sstream>
@@ -28,8 +29,10 @@ static bool read_boolean_string(std::istringstream& ss);
 static std::string get_boolean_string(bool value);
 static bool has_t_vector(const double* t);
 static bool has_r_matrix(const double* r);
+static bool has_matrix66(const Matrix& r);
 static bool has_polynom(const std::vector<double>& p);
 static void write_6d_vector(std::ostream& fp, const std::string& label, const double* t);
+static void write_6d_vector(std::ostream& fp, const std::string& label, const std::vector<double>& t);
 static void write_polynom(std::ostream& fp, const std::string& label, const std::vector<double>& p);
 static void synchronize_polynomials(Element& e);
 static void read_polynomials(std::ifstream& fp, Element& e);
@@ -135,6 +138,14 @@ void write_flat_file_trackcpp(std::ostream& fp, const Accelerator& accelerator)
       write_6d_vector(fp, "de|r_out", &e.r_out[6*4]);
       write_6d_vector(fp, "dl|r_out", &e.r_out[6*5]);
     }
+    if (has_matrix66(e.matrix66)) {
+      write_6d_vector(fp, "rx|matrix66", e.matrix66[0]);
+      write_6d_vector(fp, "px|matrix66", e.matrix66[1]);
+      write_6d_vector(fp, "ry|matrix66", e.matrix66[2]);
+      write_6d_vector(fp, "py|matrix66", e.matrix66[3]);
+      write_6d_vector(fp, "de|matrix66", e.matrix66[4]);
+      write_6d_vector(fp, "dl|matrix66", e.matrix66[5]);
+    }
 
     fp << '\n';
   }
@@ -218,6 +229,12 @@ Status::type read_flat_file_trackcpp(std::istream& fp, Accelerator& accelerator)
     if (cmd.compare("py|r_out")  == 0) { for(auto i=0; i<6; ++i) ss >> e.r_out[3*6+i]; continue; }
     if (cmd.compare("de|r_out")  == 0) { for(auto i=0; i<6; ++i) ss >> e.r_out[4*6+i]; continue; }
     if (cmd.compare("dl|r_out")  == 0) { for(auto i=0; i<6; ++i) ss >> e.r_out[5*6+i]; continue; }
+    if (cmd.compare("rx|matrix66")  == 0) { for(auto i=0; i<6; ++i) ss >> e.matrix66[0][i]; continue; }
+    if (cmd.compare("px|matrix66")  == 0) { for(auto i=0; i<6; ++i) ss >> e.matrix66[1][i]; continue; }
+    if (cmd.compare("ry|matrix66")  == 0) { for(auto i=0; i<6; ++i) ss >> e.matrix66[2][i]; continue; }
+    if (cmd.compare("py|matrix66")  == 0) { for(auto i=0; i<6; ++i) ss >> e.matrix66[3][i]; continue; }
+    if (cmd.compare("de|matrix66")  == 0) { for(auto i=0; i<6; ++i) ss >> e.matrix66[4][i]; continue; }
+    if (cmd.compare("dl|matrix66")  == 0) { for(auto i=0; i<6; ++i) ss >> e.matrix66[5][i]; continue; }
     if (cmd.compare("pass_method") == 0) {
       std::string pass_method; ss >> pass_method;
       bool found_pm = false;
@@ -454,6 +471,16 @@ static bool has_r_matrix(const double* r) {
   return false;
 }
 
+static bool has_matrix66(const Matrix& m) {
+  for (int i=0; i<6; ++i)
+    for (int j=0; j<6; ++j)
+      if ((i != j) & (m[i][j] != 0.0))
+        return true;
+      else if ((i == j) & (m[i][j] != 1.0))
+        return true;
+  return false;
+}
+
 static bool has_polynom(const std::vector<double>& p) {
   for (int i=0; i<p.size(); ++i)
     if (p[i] != 0)
@@ -469,6 +496,13 @@ static void write_6d_vector(std::ostream& fp, const std::string& label, const do
   fp << '\n';
 }
 
+static void write_6d_vector(std::ostream& fp, const std::string& label, const std::vector<double>& t) {
+  fp << std::setw(pw) << label;
+  for (int i=0; i<6; ++i)
+    fp << t[i] << "  ";
+  fp << '\n';
+}
+
 static void write_polynom(std::ostream& fp, const std::string& label, const std::vector<double>& p) {
   fp << std::setw(pw) << label;
   for (int i=0; i<p.size(); ++i)