Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-09-18 08:18:24

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