Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-01-18 09:15: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 idx_1_f, idx_2_f, idx_3_f;                         // transient float indicies
0065   float dr, dx, dy, dz;                                    // deltas for interpolation
0066   std::vector<std::vector<std::array<float, 2>>> Bvals_RZ; // B map values:  {R}, {Z}, {Br,Bz}
0067   std::vector<std::vector<std::vector<std::array<float, 3>>>>
0068       Bvals_XYZ; // B map values: {X}, {Y}, {Z}, {Bx,By,Bz}
0069 };
0070 
0071 // constructor
0072 FieldMapB::FieldMapB(const std::string& field_type_str, const std::string& coord_type_str) {
0073   std::string ftype = field_type_str;
0074   for (auto& c : ftype) {
0075     c = tolower(c);
0076   }
0077 
0078   // set type
0079   if (ftype == "magnetic") {
0080     field_type = CartesianField::MAGNETIC;
0081   } else if (ftype == "electric") {
0082     field_type = CartesianField::ELECTRIC;
0083   } else {
0084     field_type = CartesianField::UNKNOWN;
0085     printout(ERROR, "FieldMap", "Unknown field type " + ftype);
0086   }
0087 
0088   if (coord_type_str.compare("BrBz") == 0) { // BrBz
0089     fieldCoord = FieldCoord::BrBz;
0090   } else { // BxByBz
0091     fieldCoord = FieldCoord::BxByBz;
0092   }
0093 }
0094 
0095 // fill field vector
0096 void FieldMapB::Configure(std::vector<xml_comp_t> dimensions) {
0097   // Fill vectors with step size, min, max for each dimension
0098   for (auto el : dimensions) {
0099     steps.push_back(el.step());
0100     mins.push_back(getAttrOrDefault<float>(el, _Unicode(min), 0));
0101     maxs.push_back(getAttrOrDefault<float>(el, _Unicode(max), 0));
0102   }
0103 
0104   if (fieldCoord == FieldCoord::BrBz) {
0105     // N bins increased by 1 beyond grid size to account for edge cases at upper limits
0106     int nr = std::roundf((maxs[0] - mins[0]) / steps[0]) + 2;
0107     int nz = std::roundf((maxs[1] - mins[1]) / steps[1]) + 2;
0108 
0109     Bvals_RZ.resize(nr);
0110     for (auto& B2 : Bvals_RZ) {
0111       B2.resize(nz);
0112     }
0113   } else {
0114     // N bins increased by 1 beyond grid size to account for edge cases at upper limits
0115     int nx = std::roundf((maxs[0] - mins[0]) / steps[0]) + 2;
0116     int ny = std::roundf((maxs[1] - mins[1]) / steps[1]) + 2;
0117     int nz = std::roundf((maxs[2] - mins[2]) / steps[2]) + 2;
0118 
0119     Bvals_XYZ.resize(nx);
0120     for (auto& B3 : Bvals_XYZ) {
0121       B3.resize(ny);
0122       for (auto& B2 : B3) {
0123         B2.resize(nz);
0124       }
0125     }
0126   }
0127 }
0128 
0129 // get RZ cell indices corresponding to point of interest
0130 bool FieldMapB::GetIndices(float R, float Z, int* idxR, int* idxZ, float* deltaR, float* deltaZ) {
0131   // boundary check
0132   if (R > maxs[0] || R < mins[0] || Z > maxs[1] || Z < mins[1]) {
0133     return false;
0134   }
0135 
0136   // get indices
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   *deltaX = std::modf((X - mins[0]) / steps[0], &idx_1_f);
0155   *deltaY = std::modf((Y - mins[1]) / steps[1], &idx_2_f);
0156   *deltaZ = std::modf((Z - mins[2]) / steps[2], &idx_3_f);
0157   *idxX   = static_cast<int>(idx_1_f);
0158   *idxY   = static_cast<int>(idx_2_f);
0159   *idxZ   = static_cast<int>(idx_3_f);
0160 
0161   return true;
0162 }
0163 
0164 // load data
0165 void FieldMapB::LoadMap(const std::string& map_file, float scale) {
0166   std::string line;
0167   std::ifstream input(map_file);
0168   if (!input) {
0169     printout(ERROR, "FieldMapB", "FieldMapB Error: file " + map_file + " cannot be read.");
0170   }
0171 
0172   std::vector<float> coord = {};
0173   std::vector<float> Bcomp = {};
0174 
0175   while (std::getline(input, line).good()) {
0176     std::istringstream iss(line);
0177 
0178     coord.clear();
0179     Bcomp.clear();
0180 
0181     if (fieldCoord == FieldCoord::BrBz) {
0182       coord.resize(2);
0183       Bcomp.resize(2);
0184       iss >> coord[0] >> coord[1] >> Bcomp[0] >> Bcomp[1];
0185 
0186       if (!GetIndices(coord[0], coord[1], &ir, &iz, &dr, &dz)) {
0187         printout(WARNING, "FieldMapB", "coordinates out of range, skipped it.");
0188       } else { // scale field
0189         Bvals_RZ[ir][iz] = {Bcomp[0] * scale * float(tesla), Bcomp[1] * scale * float(tesla)};
0190       }
0191     } else {
0192       coord.resize(3);
0193       Bcomp.resize(3);
0194       iss >> coord[0] >> coord[1] >> coord[2] >> Bcomp[0] >> Bcomp[1] >> Bcomp[2];
0195 
0196       if (!GetIndices(coord[0], coord[1], coord[2], &ix, &iy, &iz, &dx, &dy, &dz)) {
0197         printout(WARNING, "FieldMap", "coordinates out of range, skipped it.");
0198       } else { // scale and rotate B field vector
0199         auto B = ROOT::Math::XYZPoint(Bcomp[0], Bcomp[1], Bcomp[2]);
0200         B *= scale * float(tesla);
0201         B                     = fieldRot * B;
0202         Bvals_XYZ[ix][iy][iz] = {float(B.x()), float(B.y()), float(B.z())};
0203       }
0204     }
0205   }
0206 }
0207 
0208 // get field components
0209 void FieldMapB::fieldComponents(const double* pos, double* field) {
0210   // coordinate conversion
0211   auto p = coordTranslate_inv * ROOT::Math::XYZPoint(pos[0], pos[1], pos[2]);
0212 
0213   if (fieldCoord == FieldCoord::BrBz) {
0214     // coordinates conversion
0215     const float r   = sqrt(p.x() * p.x() + p.y() * p.y());
0216     const float z   = p.z();
0217     const float phi = atan2(p.y(), p.x());
0218 
0219     if (!GetIndices(r, z, &ir, &iz, &dr, &dz)) {
0220       // out of range
0221       return;
0222     }
0223 
0224     // p1    p3
0225     //    p
0226     // p0    p2
0227     auto& p0 = Bvals_RZ[ir][iz];
0228     auto& p1 = Bvals_RZ[ir][iz + 1];
0229     auto& p2 = Bvals_RZ[ir + 1][iz];
0230     auto& p3 = Bvals_RZ[ir + 1][iz + 1];
0231 
0232     // Bilinear interpolation
0233     float Br = p0[0] * (1 - dr) * (1 - dz) + p1[0] * (1 - dr) * dz + p2[0] * dr * (1 - dz) +
0234                p3[0] * dr * dz;
0235 
0236     float Bz = p0[1] * (1 - dr) * (1 - dz) + p1[1] * (1 - dr) * dz + p2[1] * dr * (1 - dz) +
0237                p3[1] * dr * dz;
0238 
0239     // convert Br Bz to Bx By Bz and rotate field
0240     auto B = fieldRot * ROOT::Math::XYZPoint(Br * cos(phi), Br * sin(phi), Bz);
0241     field[0] += B.x();
0242     field[1] += B.y();
0243     field[2] += B.z();
0244   } else { // BxByBz
0245 
0246     if (!GetIndices(p.x(), p.y(), p.z(), &ix, &iy, &iz, &dx, &dy, &dz)) {
0247       return; // out of range
0248     }
0249 
0250     float b[3] = {0};
0251     for (int comp = 0; comp < 3; comp++) { // field component loop
0252       // Trilinear interpolation
0253       // First along X, along 4 lines
0254       float b00 = Bvals_XYZ[ix][iy][iz][comp] * (1 - dx) + Bvals_XYZ[ix + 1][iy][iz][comp] * dx;
0255       float b01 =
0256           Bvals_XYZ[ix][iy][iz + 1][comp] * (1 - dx) + Bvals_XYZ[ix + 1][iy][iz + 1][comp] * dx;
0257       float b10 =
0258           Bvals_XYZ[ix][iy + 1][iz][comp] * (1 - dx) + Bvals_XYZ[ix + 1][iy + 1][iz][comp] * dx;
0259       float b11 = Bvals_XYZ[ix][iy + 1][iz + 1][comp] * (1 - dx) +
0260                   Bvals_XYZ[ix + 1][iy + 1][iz + 1][comp] * dx;
0261       // Next along Y, along 2 lines
0262       float b0 = b00 * (1 - dy) + b10 * dy;
0263       float b1 = b01 * (1 - dy) + b11 * dy;
0264       // Finally along Z
0265       b[comp] = b0 * (1 - dz) + b1 * dz;
0266     }
0267 
0268     // field rotation done in LoadMap()
0269     field[0] += b[0];
0270     field[1] += b[1];
0271     field[2] += b[2];
0272   }
0273 
0274   return;
0275 }
0276 
0277 // assign the field map to CartesianField
0278 static Ref_t create_field_map_b(Detector& /*lcdd*/, xml::Handle_t handle) {
0279   xml_comp_t x_par(handle);
0280 
0281   if (!x_par.hasAttr(_Unicode(field_map))) {
0282     throw std::runtime_error(
0283         "FieldMapB Error: must have an xml attribute \"field_map\" for the field map.");
0284   }
0285 
0286   CartesianField field;
0287   std::string field_type = x_par.attr<std::string>(_Unicode(field_type));
0288 
0289   std::string coord_type = x_par.attr<std::string>(_Unicode(coord_type));
0290 
0291   // dimensions
0292   xml_comp_t x_dim = x_par.dimensions();
0293 
0294   // vector of dimension parameters: step, min, max
0295   std::vector<xml_comp_t> dimensions;
0296 
0297   if (coord_type.compare("BrBz") == 0) {
0298     dimensions.push_back(x_dim.child(_Unicode(R)));
0299     dimensions.push_back(x_dim.child(_Unicode(Z)));
0300   } else if (coord_type.compare("BxByBz") == 0) {
0301     dimensions.push_back(x_dim.child(_Unicode(X)));
0302     dimensions.push_back(x_dim.child(_Unicode(Y)));
0303     dimensions.push_back(x_dim.child(_Unicode(Z)));
0304   } else {
0305     printout(ERROR, "FieldMapB", "Coordinate type: " + coord_type + ", is not BrBz nor BxByBz");
0306     std::_Exit(EXIT_FAILURE);
0307   }
0308 
0309   std::string field_map_file  = x_par.attr<std::string>(_Unicode(field_map));
0310   std::string field_map_url   = x_par.attr<std::string>(_Unicode(url));
0311   std::string field_map_cache = getAttrOrDefault<std::string>(x_par, _Unicode(cache), "");
0312 
0313   EnsureFileFromURLExists(field_map_url, field_map_file, field_map_cache);
0314 
0315   float field_map_scale = x_par.attr<float>(_Unicode(scale));
0316 
0317   if (!fs::exists(fs::path(field_map_file))) {
0318     printout(ERROR, "FieldMapB", "file " + field_map_file + " does not exist");
0319     printout(ERROR, "FieldMapB", "use a FileLoader plugin before the field element");
0320     std::_Exit(EXIT_FAILURE);
0321   }
0322 
0323   auto map = new FieldMapB(field_type, coord_type);
0324   map->Configure(dimensions);
0325 
0326   // translation, rotation
0327   static float deg2r = ROOT::Math::Pi() / 180.;
0328   RotationZYX rot(0., 0., 0.);
0329   if (x_dim.hasChild(_Unicode(rotationField))) {
0330     xml_comp_t rot_dim = x_dim.child(_Unicode(rotationField));
0331     rot                = RotationZYX(rot_dim.z() * deg2r, rot_dim.y() * deg2r, rot_dim.x() * deg2r);
0332   }
0333 
0334   Translation3D trans(0., 0., 0.);
0335   if (x_dim.hasChild(_Unicode(translationCoord))) {
0336     xml_comp_t trans_dim = x_dim.child(_Unicode(translationCoord));
0337     trans                = Translation3D(trans_dim.x(), trans_dim.y(), trans_dim.z());
0338   }
0339   map->SetCoordTranslation(Transform3D(trans));
0340   map->SetFieldRotation(Transform3D(rot));
0341 
0342   map->LoadMap(field_map_file, field_map_scale);
0343   field.assign(map, x_par.nameStr(), "FieldMapB");
0344 
0345   return field;
0346 }
0347 
0348 DECLARE_XMLELEMENT(epic_FieldMapB, create_field_map_b)