Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-07-02 07:54:56

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2023 Wouter Deconinck, Dhevan Gangadharan, Aranya Giri
0003 
0004 #include <DD4hep/DetFactoryHelper.h>
0005 #include <DD4hep/FieldTypes.h>
0006 #include <DD4hep/Printout.h>
0007 #include <XML/Utilities.h>
0008 
0009 #include <cstdlib>
0010 #include <filesystem>
0011 #include <fstream>
0012 #include <iostream>
0013 #include <sstream>
0014 #include <stdexcept>
0015 #include <string>
0016 #include <tuple>
0017 namespace fs = std::filesystem;
0018 
0019 #include "FileLoaderHelper.h"
0020 
0021 using namespace dd4hep;
0022 
0023 //  Two coordinate options:
0024 //  coord_type="BrBz"
0025 //    expected fields contained within <dimensions>:
0026 //          <R step="" min="" max=""/>
0027 //          <Z step="" min="" max=""/>
0028 //  coord_type="BxByBz"
0029 //    expected fields contained within <dimensions>:
0030 //          <X step="" min="" max=""/>
0031 //          <Y step="" min="" max=""/>
0032 //          <Z step="" min="" max=""/>
0033 
0034 // implementation of the field map
0035 class FieldMapB : public dd4hep::CartesianField::Object {
0036 
0037   enum FieldCoord { BrBz, BxByBz };
0038 
0039 public:
0040   FieldMapB(const std::string& field_type_str = "magnetic",
0041             const std::string& coord_type_str = "BrBz");
0042   void Configure(std::vector<xml_comp_t> dimensions);
0043   void LoadMap(const std::string& map_file, float scale);
0044   bool GetIndices(float R, float Z, int* idxR, int* idxZ, float* deltaR, float* deltaZ);
0045   bool GetIndices(float X, float Y, float Z, int* idxX, int* idxY, int* idxZ, float* deltaX,
0046                   float* deltaY, float* deltaZ);
0047   void SetCoordTranslation(const Transform3D& tr) {
0048     coordTranslate     = tr;
0049     coordTranslate_inv = tr.Inverse();
0050   }
0051   void SetFieldRotation(const Transform3D& tr) {
0052     fieldRot     = tr;
0053     fieldRot_inv = tr.Inverse();
0054   }
0055 
0056   virtual void fieldComponents(const double* pos, double* field);
0057 
0058 private:
0059   FieldCoord fieldCoord;                                   // field coordinate type
0060   Transform3D coordTranslate, coordTranslate_inv;          // coord translation
0061   Transform3D fieldRot, fieldRot_inv;                      // field rotation
0062   std::vector<float> steps, mins, maxs;                    // B map cell info
0063   int ir, ix, iy, iz;                                      // lookup indices
0064   float dr, dx, dy, dz;                                    // deltas for interpolation
0065   std::vector<std::vector<std::array<float, 2>>> Bvals_RZ; // B map values:  {R}, {Z}, {Br,Bz}
0066   std::vector<std::vector<std::vector<std::array<float, 3>>>>
0067       Bvals_XYZ; // B map values: {X}, {Y}, {Z}, {Bx,By,Bz}
0068 };
0069 
0070 // constructor
0071 FieldMapB::FieldMapB(const std::string& field_type_str, const std::string& coord_type_str) {
0072   std::string ftype = field_type_str;
0073   for (auto& c : ftype) {
0074     c = tolower(c);
0075   }
0076 
0077   // set type
0078   if (ftype == "magnetic") {
0079     field_type = CartesianField::MAGNETIC;
0080   } else if (ftype == "electric") {
0081     field_type = CartesianField::ELECTRIC;
0082   } else {
0083     field_type = CartesianField::UNKNOWN;
0084     printout(ERROR, "FieldMap", "Unknown field type " + ftype);
0085   }
0086 
0087   if (coord_type_str.compare("BrBz") == 0) { // BrBz
0088     fieldCoord = FieldCoord::BrBz;
0089   } else { // BxByBz
0090     fieldCoord = FieldCoord::BxByBz;
0091   }
0092 }
0093 
0094 // fill field vector
0095 void FieldMapB::Configure(std::vector<xml_comp_t> dimensions) {
0096   // Fill vectors with step size, min, max for each dimension
0097   for (auto el : dimensions) {
0098     steps.push_back(el.step());
0099     mins.push_back(getAttrOrDefault<float>(el, _Unicode(min), 0));
0100     maxs.push_back(getAttrOrDefault<float>(el, _Unicode(max), 0));
0101   }
0102 
0103   if (fieldCoord == FieldCoord::BrBz) {
0104     // N bins increased by 1 beyond grid size to account for edge cases at upper limits
0105     int nr = std::roundf((maxs[0] - mins[0]) / steps[0]) + 2;
0106     int nz = std::roundf((maxs[1] - mins[1]) / steps[1]) + 2;
0107 
0108     Bvals_RZ.resize(nr);
0109     for (auto& B2 : Bvals_RZ) {
0110       B2.resize(nz);
0111     }
0112   } else {
0113     // N bins increased by 1 beyond grid size to account for edge cases at upper limits
0114     int nx = std::roundf((maxs[0] - mins[0]) / steps[0]) + 2;
0115     int ny = std::roundf((maxs[1] - mins[1]) / steps[1]) + 2;
0116     int nz = std::roundf((maxs[2] - mins[2]) / steps[2]) + 2;
0117 
0118     Bvals_XYZ.resize(nx);
0119     for (auto& B3 : Bvals_XYZ) {
0120       B3.resize(ny);
0121       for (auto& B2 : B3) {
0122         B2.resize(nz);
0123       }
0124     }
0125   }
0126 }
0127 
0128 // get RZ cell indices corresponding to point of interest
0129 bool FieldMapB::GetIndices(float R, float Z, int* idxR, int* idxZ, float* deltaR, float* deltaZ) {
0130   // boundary check
0131   if (R > maxs[0] || R < mins[0] || Z > maxs[1] || Z < mins[1]) {
0132     return false;
0133   }
0134 
0135   // get indices
0136   float idx_1_f, idx_2_f;
0137   *deltaR = std::modf((R - mins[0]) / steps[0], &idx_1_f);
0138   *deltaZ = std::modf((Z - mins[1]) / steps[1], &idx_2_f);
0139   *idxR   = static_cast<int>(idx_1_f);
0140   *idxZ   = static_cast<int>(idx_2_f);
0141 
0142   return true;
0143 }
0144 
0145 // get XYZ cell indices corresponding to point of interest
0146 bool FieldMapB::GetIndices(float X, float Y, float Z, int* idxX, int* idxY, int* idxZ,
0147                            float* deltaX, float* deltaY, float* deltaZ) {
0148   // boundary check
0149   if (X > maxs[0] || X < mins[0] || Y > maxs[1] || Y < mins[1] || Z > maxs[2] || Z < mins[2]) {
0150     return false;
0151   }
0152 
0153   // get indices
0154   float idx_1_f, idx_2_f, idx_3_f;
0155   *deltaX = std::modf((X - mins[0]) / steps[0], &idx_1_f);
0156   *deltaY = std::modf((Y - mins[1]) / steps[1], &idx_2_f);
0157   *deltaZ = std::modf((Z - mins[2]) / steps[2], &idx_3_f);
0158   *idxX   = static_cast<int>(idx_1_f);
0159   *idxY   = static_cast<int>(idx_2_f);
0160   *idxZ   = static_cast<int>(idx_3_f);
0161 
0162   return true;
0163 }
0164 
0165 // load data
0166 void FieldMapB::LoadMap(const std::string& map_file, float scale) {
0167   std::string line;
0168   std::ifstream input(map_file);
0169   if (!input) {
0170     printout(ERROR, "FieldMapB", "FieldMapB Error: file " + map_file + " cannot be read.");
0171   }
0172 
0173   std::vector<float> coord = {};
0174   std::vector<float> Bcomp = {};
0175 
0176   while (std::getline(input, line).good()) {
0177     std::istringstream iss(line);
0178 
0179     coord.clear();
0180     Bcomp.clear();
0181 
0182     if (fieldCoord == FieldCoord::BrBz) {
0183       coord.resize(2);
0184       Bcomp.resize(2);
0185       iss >> coord[0] >> coord[1] >> Bcomp[0] >> Bcomp[1];
0186 
0187       if (!GetIndices(coord[0], coord[1], &ir, &iz, &dr, &dz)) {
0188         printout(WARNING, "FieldMapB", "coordinates out of range, skipped it.");
0189       } else { // scale field
0190         Bvals_RZ[ir][iz] = {Bcomp[0] * scale * float(tesla), Bcomp[1] * scale * float(tesla)};
0191       }
0192     } else {
0193       coord.resize(3);
0194       Bcomp.resize(3);
0195       iss >> coord[0] >> coord[1] >> coord[2] >> Bcomp[0] >> Bcomp[1] >> Bcomp[2];
0196 
0197       if (!GetIndices(coord[0], coord[1], coord[2], &ix, &iy, &iz, &dx, &dy, &dz)) {
0198         printout(WARNING, "FieldMap", "coordinates out of range, skipped it.");
0199       } else { // scale and rotate B field vector
0200         auto B = ROOT::Math::XYZPoint(Bcomp[0], Bcomp[1], Bcomp[2]);
0201         B *= scale * float(tesla);
0202         B                     = fieldRot * B;
0203         Bvals_XYZ[ix][iy][iz] = {float(B.x()), float(B.y()), float(B.z())};
0204       }
0205     }
0206   }
0207 }
0208 
0209 // get field components
0210 void FieldMapB::fieldComponents(const double* pos, double* field) {
0211   // coordinate conversion
0212   auto p = coordTranslate_inv * ROOT::Math::XYZPoint(pos[0], pos[1], pos[2]);
0213 
0214   if (fieldCoord == FieldCoord::BrBz) {
0215     // coordinates conversion
0216     const float r   = sqrt(p.x() * p.x() + p.y() * p.y());
0217     const float z   = p.z();
0218     const float phi = atan2(p.y(), p.x());
0219 
0220     if (!GetIndices(r, z, &ir, &iz, &dr, &dz)) {
0221       // out of range
0222       return;
0223     }
0224 
0225     // p1    p3
0226     //    p
0227     // p0    p2
0228     auto& p0 = Bvals_RZ[ir][iz];
0229     auto& p1 = Bvals_RZ[ir][iz + 1];
0230     auto& p2 = Bvals_RZ[ir + 1][iz];
0231     auto& p3 = Bvals_RZ[ir + 1][iz + 1];
0232 
0233     // Bilinear interpolation
0234     float Br = p0[0] * (1 - dr) * (1 - dz) + p1[0] * (1 - dr) * dz + p2[0] * dr * (1 - dz) +
0235                p3[0] * dr * dz;
0236 
0237     float Bz = p0[1] * (1 - dr) * (1 - dz) + p1[1] * (1 - dr) * dz + p2[1] * dr * (1 - dz) +
0238                p3[1] * dr * dz;
0239 
0240     // convert Br Bz to Bx By Bz and rotate field
0241     auto B = fieldRot * ROOT::Math::XYZPoint(Br * cos(phi), Br * sin(phi), Bz);
0242     field[0] += B.x();
0243     field[1] += B.y();
0244     field[2] += B.z();
0245   } else { // BxByBz
0246 
0247     if (!GetIndices(p.x(), p.y(), p.z(), &ix, &iy, &iz, &dx, &dy, &dz)) {
0248       return; // out of range
0249     }
0250 
0251     float b[3] = {0};
0252     for (int comp = 0; comp < 3; comp++) { // field component loop
0253       // Trilinear interpolation
0254       // First along X, along 4 lines
0255       float b00 = Bvals_XYZ[ix][iy][iz][comp] * (1 - dx) + Bvals_XYZ[ix + 1][iy][iz][comp] * dx;
0256       float b01 =
0257           Bvals_XYZ[ix][iy][iz + 1][comp] * (1 - dx) + Bvals_XYZ[ix + 1][iy][iz + 1][comp] * dx;
0258       float b10 =
0259           Bvals_XYZ[ix][iy + 1][iz][comp] * (1 - dx) + Bvals_XYZ[ix + 1][iy + 1][iz][comp] * dx;
0260       float b11 = Bvals_XYZ[ix][iy + 1][iz + 1][comp] * (1 - dx) +
0261                   Bvals_XYZ[ix + 1][iy + 1][iz + 1][comp] * dx;
0262       // Next along Y, along 2 lines
0263       float b0 = b00 * (1 - dy) + b10 * dy;
0264       float b1 = b01 * (1 - dy) + b11 * dy;
0265       // Finally along Z
0266       b[comp] = b0 * (1 - dz) + b1 * dz;
0267     }
0268 
0269     // field rotation done in LoadMap()
0270     field[0] += b[0];
0271     field[1] += b[1];
0272     field[2] += b[2];
0273   }
0274 
0275   return;
0276 }
0277 
0278 // assign the field map to CartesianField
0279 static Ref_t create_field_map_b(Detector& /*lcdd*/, xml::Handle_t handle) {
0280   xml_comp_t x_par(handle);
0281 
0282   if (!x_par.hasAttr(_Unicode(field_map))) {
0283     throw std::runtime_error(
0284         "FieldMapB Error: must have an xml attribute \"field_map\" for the field map.");
0285   }
0286 
0287   CartesianField field;
0288   std::string field_type = x_par.attr<std::string>(_Unicode(field_type));
0289 
0290   std::string coord_type = x_par.attr<std::string>(_Unicode(coord_type));
0291 
0292   // dimensions
0293   xml_comp_t x_dim = x_par.dimensions();
0294 
0295   // vector of dimension parameters: step, min, max
0296   std::vector<xml_comp_t> dimensions;
0297 
0298   if (coord_type.compare("BrBz") == 0) {
0299     dimensions.push_back(x_dim.child(_Unicode(R)));
0300     dimensions.push_back(x_dim.child(_Unicode(Z)));
0301   } else if (coord_type.compare("BxByBz") == 0) {
0302     dimensions.push_back(x_dim.child(_Unicode(X)));
0303     dimensions.push_back(x_dim.child(_Unicode(Y)));
0304     dimensions.push_back(x_dim.child(_Unicode(Z)));
0305   } else {
0306     printout(ERROR, "FieldMapB", "Coordinate type: " + coord_type + ", is not BrBz nor BxByBz");
0307     std::_Exit(EXIT_FAILURE);
0308   }
0309 
0310   std::string field_map_file  = x_par.attr<std::string>(_Unicode(field_map));
0311   std::string field_map_url   = x_par.attr<std::string>(_Unicode(url));
0312   std::string field_map_cache = getAttrOrDefault<std::string>(x_par, _Unicode(cache), "");
0313 
0314   EnsureFileFromURLExists(field_map_url, field_map_file, field_map_cache);
0315 
0316   float field_map_scale = x_par.attr<float>(_Unicode(scale));
0317 
0318   if (!fs::exists(fs::path(field_map_file))) {
0319     printout(ERROR, "FieldMapB", "file " + field_map_file + " does not exist");
0320     printout(ERROR, "FieldMapB", "use a FileLoader plugin before the field element");
0321     std::_Exit(EXIT_FAILURE);
0322   }
0323 
0324   auto map = new FieldMapB(field_type, coord_type);
0325   map->Configure(dimensions);
0326 
0327   // translation, rotation
0328   static float deg2r = ROOT::Math::Pi() / 180.;
0329   RotationZYX rot(0., 0., 0.);
0330   if (x_dim.hasChild(_Unicode(rotationField))) {
0331     xml_comp_t rot_dim = x_dim.child(_Unicode(rotationField));
0332     rot                = RotationZYX(rot_dim.z() * deg2r, rot_dim.y() * deg2r, rot_dim.x() * deg2r);
0333   }
0334 
0335   Translation3D trans(0., 0., 0.);
0336   if (x_dim.hasChild(_Unicode(translationCoord))) {
0337     xml_comp_t trans_dim = x_dim.child(_Unicode(translationCoord));
0338     trans                = Translation3D(trans_dim.x(), trans_dim.y(), trans_dim.z());
0339   }
0340   map->SetCoordTranslation(Transform3D(trans));
0341   map->SetFieldRotation(Transform3D(rot));
0342 
0343   map->LoadMap(field_map_file, field_map_scale);
0344   field.assign(map, x_par.nameStr(), "FieldMapB");
0345 
0346   return field;
0347 }
0348 
0349 DECLARE_XMLELEMENT(epic_FieldMapB, create_field_map_b)