BFieldComponentBeamline.cc

/**************************************************************************
 * basf2 (Belle II Analysis Software Framework)                           *
 * Author: The Belle II Collaboration                                     *
 *                                                                        *
 * See git log for contributors and copyright holders.                    *
 * This file is licensed under LGPL-3.0, see LICENSE.md.                  *
 **************************************************************************/
 
#include <geometry/bfieldmap/BFieldComponentBeamline.h>
 
#include <framework/utilities/FileSystem.h>
#include <framework/logging/Logger.h>
#include <framework/utilities/MathHelpers.h>
 
#include <boost/iostreams/filtering_stream.hpp>
#include <boost/iostreams/device/file.hpp>
#include <boost/iostreams/filter/gzip.hpp>
 
#include <cmath>
#include <limits>
 
using namespace std;
namespace io = boost::iostreams;
namespace Belle2 {

  struct triangle_t {
    short int j0{0};
    short int j1{0};
    short int j2{0};
    short int n0{0};
    short int n1{0};
    short int n2{0};
  };

  struct xy_t {
    double x{0}; 
    double y{0}; 
  };

  class TriangularInterpolation {
  public:
    [[nodiscard]] const vector<xy_t>& getPoints() const { return m_points;}
    [[nodiscard]] const vector<triangle_t>& getTriangles() const { return m_triangles;}

    TriangularInterpolation() = default;

    TriangularInterpolation(vector<xy_t>& pc, vector<triangle_t>& ts, double d)
    {
      init(pc, ts, d);
    }

    ~TriangularInterpolation() = default;

    void init(vector<xy_t>& points, vector<triangle_t>& triangles, double d)
    {
      std::swap(points, m_points);
      std::swap(triangles, m_triangles);
      const double inf = numeric_limits<double>::infinity();
      double xmin = inf, xmax = -inf;
      double ymin = inf, ymax = -inf;
      auto limit = [&xmin, &xmax, &ymin, &ymax](const xy_t & p) -> void {
        xmin = std::min(xmin, p.x); xmax = std::max(xmax, p.x);
        ymin = std::min(ymin, p.y); ymax = std::max(ymax, p.y);
      };
      for (const auto& t : m_triangles) {
        limit(m_points[t.j0]);
        limit(m_points[t.j1]);
        limit(m_points[t.j2]);
      }
      double xw = xmax - xmin;
      double yw = ymax - ymin;
 
      xmin -= xw * (1. / 256), xmax += xw * (1. / 256);
      ymin -= yw * (1. / 256), ymax += yw * (1. / 256);
      int nx = lrint(xw * (1 + 1. / 128) / d), ny = lrint(yw * (1 + 1. / 128) / d);
      nx = std::max(nx, 2);
      ny = std::max(ny, 2);
      makeIndex(nx, xmin, xmax, ny, ymin, ymax);
    }

    [[nodiscard]] xy_t triangleCenter(const triangle_t& p) const
    {
      const xy_t& p0 = m_points[p.j0], &p1 = m_points[p.j1], &p2 = m_points[p.j2];
      return {(p0.x + p1.x + p2.x)* (1. / 3), (p0.y + p1.y + p2.y)* (1. / 3)};
    }

    void makeIndex(int nx, double xmin, double xmax, int ny, double ymin, double ymax)
    {
      m_nx = nx, m_ny = ny, m_xmin = xmin, m_xmax = xmax, m_ymin = ymin, m_ymax = ymax;
      m_ixnorm = (nx - 1) / (xmax - xmin), m_iynorm = (ny - 1) / (ymax - ymin);
      double dx = (xmax - xmin) / (nx - 1), dy = (ymax - ymin) / (ny - 1);
      //    cout<<nx<<" "<<xmin<<" "<<xmax<<" "<<ny<<" "<<ymin<<" "<<ymax<<endl;
      m_spatialIndex.resize(nx * ny);
      m_triangleCenters.resize(m_triangles.size());
      for (unsigned int i = 0; i < m_triangles.size(); i++) m_triangleCenters[i] = triangleCenter(m_triangles[i]);
 
      m_triangleAreas.resize(m_triangles.size());
      auto getTriangleArea = [this](const triangle_t& p) ->double{
        const xy_t& p0 = m_points[p.j0], &p1 = m_points[p.j1], &p2 = m_points[p.j2];
        double d21x = p2.x - p1.x, d21y = p2.y - p1.y;
        double d01x = p0.x - p1.x, d01y = p0.y - p1.y;
        return 1 / (d21x * d01y - d21y * d01x);
      };
      for (unsigned int i = 0; i < m_triangles.size(); i++) m_triangleAreas[i] = getTriangleArea(m_triangles[i]);
 
 
      for (int ix = 0; ix < nx; ix++) {
        double x = xmin + ix * dx;
        for (int iy = 0; iy < ny; iy++) {
          double y = ymin + iy * dy;
          int imin = -1; double dmin = 1e100;
          for (unsigned int i = 0; i < m_triangleCenters.size(); i++) {
            const xy_t& p = m_triangleCenters[i];
            double d = square(p.x - x) + square(p.y - y);
            if (d < dmin) {imin = i; dmin = d;}
          }
          int k = iy + ny * ix;
          m_spatialIndex[k] = imin;
        }
      }
    }

    [[nodiscard]] short int sideCross(short int prev, short int curr, const xy_t& r, const xy_t& v0) const
    {
      const double vx = r.x - v0.x, vy = r.y - v0.y;
      auto isCrossed = [&vx, &vy, &v0](const xy_t & p1, const xy_t & p0) -> bool {
        double u0x = p0.x, u0y = p0.y;
        double ux = p1.x - u0x, uy = p1.y - u0y;
        double dx = u0x - v0.x, dy = u0y - v0.y;
        double D = uy * vx - ux * vy;
        double t = dx * vy - dy * vx;
        double s = dx * uy - dy * ux;
        return ((t < D) != (t < 0)) && ((s < D) != (s < 0));
      };
 
      const triangle_t& p = m_triangles[curr];
      const xy_t& p0 = m_points[p.j0], &p1 = m_points[p.j1], &p2 = m_points[p.j2];
      if (p.n0 != prev && isCrossed(p1, p2)) return p.n0;
      if (p.n1 != prev && isCrossed(p2, p0)) return p.n1;
      if (p.n2 != prev && isCrossed(p0, p1)) return p.n2;
      return m_triangles.size();
    }

    void weights(short int i, const xy_t& r, double& w0, double& w1, double& w2) const
    {
      const triangle_t& p = m_triangles[i];
      const xy_t& p0 = m_points[p.j0], &p1 = m_points[p.j1], &p2 = m_points[p.j2];
      double dx2  = p2.x -  r.x, dy2  = p2.y -  r.y;
      double d21x = p2.x - p1.x, d21y = p2.y - p1.y;
      double d02x = p0.x - p2.x, d02y = p0.y - p2.y;
      w0 = (dx2 * d21y - dy2 * d21x) * m_triangleAreas[i];
      w1 = (dx2 * d02y - dy2 * d02x) * m_triangleAreas[i];
      w2 = 1 - (w0 + w1);
    }

    [[nodiscard]] short int findTriangle(const xy_t& r0) const
    {
      unsigned int ix = lrint((r0.x - m_xmin) * m_ixnorm), iy = lrint((r0.y - m_ymin) * m_iynorm);
      short int curr = (ix < m_nx && iy < m_ny) ? m_spatialIndex[iy + m_ny * ix] : 0;
      xy_t r = m_triangleCenters[curr];
      const short int end = m_triangles.size();
      short int prev = end;
      do {
        short int next = sideCross(prev, curr, r, r0);
        if (next == end) break;
        prev = curr;
        curr = next;
      } while (true);
      return curr;
    }
  protected:
    vector<triangle_t> m_triangles;
    vector<xy_t> m_points;
    vector<xy_t> m_triangleCenters;
    vector<double> m_triangleAreas;
    vector<short int> m_spatialIndex;
    double m_xmin;
    double m_xmax;
    double m_ymin;
    double m_ymax;
    unsigned int m_nx;
    unsigned int m_ny;
    double m_ixnorm{1};
    double m_iynorm{1};
  };

  class BeamlineFieldMapInterpolation {
  public:
    [[nodiscard]] const TriangularInterpolation& getTriangularInterpolation() const {return m_triInterpol;}
    BeamlineFieldMapInterpolation() = default;
    ~BeamlineFieldMapInterpolation() = default;

    void init(const string& fieldmapname, const string& interpolname, double validRadius)
    {
      if (fieldmapname.empty()) {
        B2ERROR("The filename for the beamline magnetic field component is empty !");
        return;
      }
      std::string l_fieldmapname = FileSystem::findFile("/data/" + fieldmapname);
 
      if (interpolname.empty()) {
        B2ERROR("The filename for the triangulation of the beamline magnetic field component is empty !");
        return;
      }
      std::string l_interpolname = FileSystem::findFile("/data/" + interpolname);
 
      m_rmax = validRadius;
 
      B2DEBUG(50, "Delaunay triangulation of the beamline field map: " << l_interpolname);
      B2DEBUG(50, "Beamline field map: " << l_fieldmapname);
 
      // load interpolation triangles
      ifstream INd(l_interpolname);
      int nts; INd >> nts;
      B2DEBUG(50, "Total number of triangles: " << nts);
      vector<triangle_t> ts;
      ts.reserve(nts);
      // load triangle definitions from file
      {
        triangle_t p;
        while (INd >> nts >> p.j0 >> p.j1 >> p.j2 >> p.n0 >> p.n1 >> p.n2) ts.push_back(p);
      }
 
      //Load the map file
      io::filtering_istream IN;
      IN.push(io::gzip_decompressor());
      IN.push(io::file_source(l_fieldmapname));
 
      int nrphi;
      IN >> nrphi >> m_rj >> m_nr >> m_nphi;
      IN >> m_nz >> m_zj >> m_dz0 >> m_dz1;
      m_idphi = m_nphi / M_PI;
      m_rj2 = m_rj * m_rj;
      m_idz0 = 1 / m_dz0;
      m_idz1 = 1 / m_dz1;
      int nz0 = 2 * int(m_zj / m_dz0);
      m_zmax = (m_nz - (nz0 + 1)) / 2 * m_dz1 + m_zj;
      m_nz1 = (m_nz - (nz0 + 1)) / 2;
      m_nz2 = m_nz1 + nz0;
      //    cout<<_zmax<<" "<<nz0<<" "<<m_nz1<<" "<<m_nz2<<endl;
 
      struct cs_t {double c, s;};
      vector<cs_t> cs(nrphi);
      vector<xy_t> pc;
      vector<ROOT::Math::XYZVector> tbc;
      pc.reserve(nrphi);
      char cbuf[256]; IN.getline(cbuf, 256);
      double rmax = 0;
      for (int j = 0; j < nrphi; j++) {
        IN.getline(cbuf, 256);
        char* next = cbuf;
        double r    = strtod(next, &next);
        double phi  = strtod(next, &next);
        strtod(next, &next);
        double Br   = strtod(next, &next);
        double Bphi = strtod(next, &next);
        double Bz   = strtod(next, nullptr);
        r *= 100;
        rmax = std::max(r, rmax);
        if (phi == 0) {
          cs[j] = { 1, 0};
        } else if (phi == 180) {
          cs[j] = { -1, 0};
        } else {
          phi *= M_PI / 180;
          cs[j] = {cos(phi), sin(phi)};
        }
        double x = r * cs[j].c, y = r * cs[j].s;
        pc.push_back({x, y});
        if (cs[j].s == 0) Bphi = 0;
        double Bx = Br * cs[j].c - Bphi * cs[j].s;
        double By = Br * cs[j].s + Bphi * cs[j].c;
        tbc.emplace_back(Bx, By, Bz);
      }
      //    cout<<"Field map has data points up to R="<<rmax<<" cm."<<endl;
 
      m_idr = m_nr / (rmax - m_rj);
      m_rmax = std::min(m_rmax, rmax);
      bool reduce = m_rj > m_rmax;
 
      // if valid radius is within the triangular mesh try to reduce
      // field map keeping only points which participate to the
      // interpolation
      vector<bool> ip;
      if (reduce) {
        ip.resize(nrphi, false);
        vector<bool> it(ts.size(), false);
        auto inside = [this](const xy_t & p)->bool{
          return p.x * p.x + p.y * p.y < m_rmax * m_rmax;
        };
 
        for (int i = 0, imax = ts.size(); i < imax; i++) {
          const triangle_t& p = ts[i];
          const xy_t& p0 = pc[p.j0], &p1 = pc[p.j1], &p2 = pc[p.j2];
          if (inside(p0) || inside(p1) || inside(p2)) {
            it[i] = true;
            ip[p.j0] = true;
            ip[p.j1] = true;
            ip[p.j2] = true;
          }
        }
 
        vector<short int> pindx(nrphi, -1);
        int rnp = 0;
        for (int i = 0, imax = ip.size(); i < imax; i++) {
          if (ip[i]) pindx[i] = rnp++;
        }
        vector<xy_t> rpc;
        rpc.reserve(rnp);
        for (int i = 0, imax = pc.size(); i < imax; i++) {
          if (ip[i]) rpc.push_back(pc[i]);
        }
 
        vector<short int> tindx(ts.size(), -1);
        short int rnt = 0;
        for (int i = 0, imax = it.size(); i < imax; i++) {
          if (it[i]) tindx[i] = rnt++;
        }
        vector<triangle_t> rts;
        rts.reserve(rnt);
        short int nt = ts.size();
        auto newind = [&nt, &tindx, &rnt](short int n) -> short int {return (n < nt) ? tindx[n] : rnt;};
        for (int i = 0, imax = ts.size(); i < imax; i++) {
          if (it[i]) {
            const triangle_t& t = ts[i];
            rts.push_back({pindx[t.j0], pindx[t.j1], pindx[t.j2], newind(t.n0), newind(t.n1), newind(t.n2)});
          }
        }
 
        B2DEBUG(50, "Reduce map size to cover only region R<" << m_rmax << " cm: Ntriangles=" << rnt << " Nxypoints = " << rnp <<
                " Nzslices=" << m_nz << " NBpoints = " << rnp * m_nz);
        std::swap(rpc, pc);
        std::swap(rts, ts);
      } else {
        ip.resize(nrphi, true);
      }
      m_nxy = pc.size();
 
      m_triInterpol.init(pc, ts, 0.1);
 
      vector<ROOT::Math::XYZVector> bc(m_nxy * m_nz);
      unsigned int count = 0;
      for (int i = 0; i < nrphi; i++) {
        if (ip[i]) bc[count++] = ROOT::Math::XYZVector(tbc[i]);
      }
 
      for (int i = 1; i < m_nz; ++i) {
        for (int j = 0; j < nrphi; j++) {
          IN.getline(cbuf, 256);
          if (!ip[j]) continue;
          char* next = cbuf;
          next = strchr(next, ' ');
          next = strchr(next + 1, ' ');
          next = strchr(next + 1, ' ');
          double Br   = strtod(next, &next);
          double Bphi = strtod(next, &next);
          double Bz   = strtod(next, nullptr);
          if (cs[j].s == 0) Bphi = 0;
          double Bx = Br * cs[j].c - Bphi * cs[j].s;
          double By = Br * cs[j].s + Bphi * cs[j].c;
          bc[count++].SetXYZ(Bx, By, Bz);
        }
      }
      assert(count == bc.size());
      swap(bc, m_B);
    }

    int zIndexAndWeight(double z, double& w1) const
    {
      if (std::abs(z) > m_zmax) return -1;
      double fz;
      int iz = 0;
      if (z < - m_zj) {
        fz = (z + m_zmax) * m_idz1;
      } else if (z < m_zj) {
        fz = (z + m_zj) * m_idz0;
        iz = m_nz1;
      } else {
        fz = (z - m_zj) * m_idz1;
        iz = m_nz2;
      }
      int jz = static_cast<int>(fz);
      w1 = fz - jz;
      iz += jz;
      if (iz == m_nz) {
        --iz;
        w1 = 1;
      }
      return iz;
    }

    [[nodiscard]] bool inRange(const ROOT::Math::XYZVector& v) const
    {
      if (std::abs(v.Z()) > m_zmax) return false;
      double R2 = v.X() * v.X() + v.Y() * v.Y();
      if (R2 > m_rmax * m_rmax) return false;
      return true;
    }

    [[nodiscard]] ROOT::Math::XYZVector interpolateField(const ROOT::Math::XYZVector& v) const
    {
      ROOT::Math::XYZVector res = {0, 0, 0};
      double R2 = v.X() * v.X() + v.Y() * v.Y();
      if (R2 > m_rmax * m_rmax) return res;
      double wz1;
      int iz = zIndexAndWeight(v.Z(), wz1);
      if (iz < 0) return res;
      double wz0 = 1 - wz1;
 
      if (R2 < m_rj2) { // triangular interpolation
        xy_t xy = {v.X(), std::abs(v.Y())};
        double w0, w1, w2;
        short int it = m_triInterpol.findTriangle(xy);
        m_triInterpol.weights(it, xy, w0, w1, w2);
        auto t = m_triInterpol.getTriangles().begin() + it;
        int j0 = t->j0, j1 = t->j1, j2 = t->j2;
        const ROOT::Math::XYZVector* B = m_B.data() + m_nxy * iz;
        ROOT::Math::XYZVector b = (B[j0] * w0 + B[j1] * w1 + B[j2] * w2) * wz0;
        B += m_nxy; // next z-slice
        b += (B[j0] * w0 + B[j1] * w1 + B[j2] * w2) * wz1;
        res = b;
      } else {// r-phi grid
        double r = sqrt(R2), phi = atan2(std::abs(v.Y()), v.X());
        double fr = (r - m_rj) * m_idr;
        double fphi = phi * m_idphi;
 
        int ir = static_cast<int>(fr);
        int iphi = static_cast<int>(fphi);
 
        ir -= (ir == m_nr);
        iphi -= (iphi == m_nphi);
 
        double wr1 = fr - ir, wr0 = 1 - wr1;
        double wphi1 = fphi - iphi, wphi0 = 1 - wphi1;
        const int nr1 = m_nr + 1, nphi1 = m_nphi + 1;
        int j00 = m_nxy - nr1 * (nphi1 - iphi) + ir;
        int j01 = j00 + 1;
        int j10 = j00 + nr1;
        int j11 = j01 + nr1;
 
        double w00 = wr0 * wphi0, w01 = wphi0 * wr1, w10 = wphi1 * wr0, w11 = wphi1 * wr1;
        const ROOT::Math::XYZVector* B = m_B.data() + m_nxy * iz;
        ROOT::Math::XYZVector b = (B[j00] * w00 + B[j01] * w01 + B[j10] * w10 + B[j11] * w11) * wz0;
        B += m_nxy; // next z-slice
        b += (B[j00] * w00 + B[j01] * w01 + B[j10] * w10 + B[j11] * w11) * wz1;
        res = b;
      }
      if (v.Y() < 0) res.SetY(-res.Y());
      return res;
    }
  protected:
    vector<ROOT::Math::XYZVector> m_B;
    TriangularInterpolation m_triInterpol;
    int m_nxy{0};
    int m_nz{0};
    int m_nz1{0};
    int m_nz2{0};
    int m_nr{0};
    int m_nphi{0};
    double m_rj{0};
    double m_rj2{0};
    double m_zj{0};
    double m_dz0{0};
    double m_dz1{0};
    double m_idz0{0};
    double m_idz1{0};
    double m_idphi{0};
    double m_idr{0};
    double m_rmax{0};
    double m_zmax{0};
  };
 
  void BFieldComponentBeamline::initialize()
  {
    if (!m_ler) m_ler = new BeamlineFieldMapInterpolation;
    if (!m_her) m_her = new BeamlineFieldMapInterpolation;
    m_ler->init(m_mapFilename_ler, m_interFilename_ler, m_mapRegionR[1]);
    m_her->init(m_mapFilename_her, m_interFilename_her, m_mapRegionR[1]);
  }
 
  BFieldComponentBeamline::~BFieldComponentBeamline()
  {
    if (m_ler) delete m_ler;
    if (m_her) delete m_her;
  }
 
  bool BFieldComponentBeamline::isInRange(const ROOT::Math::XYZVector& p) const
  {
    if (!m_ler || !m_her) return false;
    double s = m_sinBeamCrossAngle, c = m_cosBeamCrossAngle;
    ROOT::Math::XYZVector v = -p; // invert coordinates to match ANSYS one
    double xc = v.X() * c, zs = v.Z() * s, zc = v.Z() * c, xs = v.X() * s;
    ROOT::Math::XYZVector hv{xc - zs, v.Y(), zc + xs};
    ROOT::Math::XYZVector lv{xc + zs, v.Y(), zc - xs};
    return m_ler->inRange(lv) || m_her->inRange(hv);
  }
 
  ROOT::Math::XYZVector BFieldComponentBeamline::calculate(const ROOT::Math::XYZVector& p) const
  {
    ROOT::Math::XYZVector res;
    double s = m_sinBeamCrossAngle, c = m_cosBeamCrossAngle;
    ROOT::Math::XYZVector v = -p; // invert coordinates to match ANSYS one
    double xc = v.X() * c, zs = v.Z() * s, zc = v.Z() * c, xs = v.X() * s;
    ROOT::Math::XYZVector hv{xc - zs, v.Y(), zc + xs};
    ROOT::Math::XYZVector lv{xc + zs, v.Y(), zc - xs};
    ROOT::Math::XYZVector hb = m_her->interpolateField(hv);
    ROOT::Math::XYZVector lb = m_ler->interpolateField(lv);
    ROOT::Math::XYZVector rhb{hb.X()* c + hb.Z()* s, hb.Y(),  hb.Z()* c - hb.X()* s};
    ROOT::Math::XYZVector rlb{lb.X()* c - lb.Z()* s, lb.Y(),  lb.Z()* c + lb.X()* s};
 
    double mhb = std::abs(rhb.X()) + std::abs(rhb.Y()) + std::abs(rhb.Z());
    double mlb = std::abs(rlb.X()) + std::abs(rlb.Y()) + std::abs(rlb.Z());
 
    if (mhb < 1e-10) res = rlb;
    else if (mlb < 1e-10) res = rhb;
    else {
      res = 0.5 * (rlb + rhb);
    }
    return res;
  }
 
  void BFieldComponentBeamline::terminate()
  {
  }

  BFieldComponentBeamline** GetInstancePtr()
  {
    static BFieldComponentBeamline* gInstance = nullptr;
    return &gInstance;
  }

  BFieldComponentBeamline& BFieldComponentBeamline::Instance()
  {
    auto gInstance = GetInstancePtr();
    if (*gInstance == nullptr) {
      // Constructor creates a new instance, inits gInstance.
      new BFieldComponentBeamline();
    }
    return **gInstance;
  }
 
  BFieldComponentBeamline::BFieldComponentBeamline()
  {
    auto gInstance = GetInstancePtr();
    if (*gInstance != nullptr) {
      B2WARNING("BFieldComponentBeamline: object already instantiated");
    } else {
      *gInstance = this;
    }
  }

}