include/root/TGeoVoxelGrid.h

0001 /// \author Sandro Wenzel <sandro.wenzel@cern.ch>
0002 /// \date 2024-02-22
0003
0004 /*************************************************************************
0005  * Copyright (C) 1995-2024, Rene Brun and Fons Rademakers.               *
0006  * All rights reserved.                                                  *
0007  *                                                                       *
0008  * For the licensing terms see $ROOTSYS/LICENSE.                         *
0009  * For the list of contributors see $ROOTSYS/README/CREDITS.             *
0010  *************************************************************************/
0011
0012 #ifndef ROOT_TGeoVoxelGrid
0013 #define ROOT_TGeoVoxelGrid
0014
0015 #include <array>
0016 #include <cmath>
0017 #include <limits>
0018
0019 // a simple structure to encode voxel indices, to address
0020 // individual voxels in the 3D grid.
0021 struct TGeoVoxelGridIndex {
0022    int ix{-1};
0023    int iy{-1};
0024    int iz{-1};
0025    size_t idx{std::numeric_limits<size_t>::max()};
0026    bool isValid() const { return idx != std::numeric_limits<size_t>::max(); }
0027 };
0028
0029 /// A finite 3D grid structure, mapping/binning arbitrary 3D cartesian points
0030 /// onto discrete "voxels". Each such voxel can store an object of type T.
0031 /// The precision of the voxel binning is done with S (float or double).
0032 template <typename T, typename S = float>
0033 class TGeoVoxelGrid {
0034 public:
0035    TGeoVoxelGrid(S xmin, S ymin, S zmin, S xmax, S ymax, S zmax, S Lx_, S Ly_, S Lz_)
0036       : fMinBound{xmin, ymin, zmin}, fMaxBound{xmax, ymax, zmax}, fLx(Lx_), fLy(Ly_), fLz(Lz_)
0037    {
0038
0039       // Calculate the number of voxels in each dimension
0040       fNx = static_cast<int>((fMaxBound[0] - fMinBound[0]) / fLx);
0041       fNy = static_cast<int>((fMaxBound[1] - fMinBound[1]) / fLy);
0042       fNz = static_cast<int>((fMaxBound[2] - fMinBound[2]) / fLz);
0043
0044       finvLx = 1. / fLx;
0045       finvLy = 1. / fLy;
0046       finvLz = 1. / fLz;
0047
0048       fHalfDiag = std::sqrt(fLx / 2. * fLx / 2. + fLy / 2. * fLy / 2. + fLz / 2. * fLz / 2.);
0049
0050       // Resize the grid to hold the voxels
0051       fGrid.resize(fNx * fNy * fNz);
0052    }
0053
0054    T &at(int i, int j, int k) { return fGrid[index(i, j, k)]; }
0055
0056    // check if point is covered by voxel structure
0057    bool inside(std::array<S, 3> const &p) const
0058    {
0059       for (int i = 0; i < 3; ++i) {
0060          if (p[i] < fMinBound[i] || p[i] > fMaxBound[i]) {
0061             return false;
0062          }
0063       }
0064       return true;
0065    }
0066
0067    // Access a voxel given a 3D point P
0068    T &at(std::array<S, 3> const &P)
0069    {
0070       int i, j, k;
0071       pointToVoxelIndex(P, i, j, k); // Convert point to voxel index
0072       return fGrid[index(i, j, k)];  // Return reference to voxel's data
0073    }
0074
0075    T *at(TGeoVoxelGridIndex const &vi)
0076    {
0077       if (!vi.isValid()) {
0078          return nullptr;
0079       }
0080       return &fGrid[vi.idx];
0081    }
0082
0083    // Set the data of a voxel at point P
0084    void set(std::array<S, 3> const &p, const T &value)
0085    {
0086       int i, j, k;
0087       pointToVoxelIndex(p, i, j, k); // Convert point to voxel index
0088       fGrid[index(i, j, k)] = value; // Set the value at the voxel
0089    }
0090
0091    // Set the data of a voxel at point P
0092    void set(int i, int j, int k, const T &value)
0093    {
0094       fGrid[index(i, j, k)] = value; // Set the value at the voxel
0095    }
0096
0097    void set(TGeoVoxelGridIndex const &vi, const T &value) { fGrid[vi.idx] = value; }
0098
0099    // Get voxel dimensions
0100    int getVoxelCountX() const { return fNx; }
0101    int getVoxelCountY() const { return fNy; }
0102    int getVoxelCountZ() const { return fNz; }
0103
0104    // returns the cartesian mid-point coordinates of a voxel given by a VoxelIndex
0105    std::array<S, 3> getVoxelMidpoint(TGeoVoxelGridIndex const &vi) const
0106    {
0107       const S midX = fMinBound[0] + (vi.ix + 0.5) * fLx;
0108       const S midY = fMinBound[1] + (vi.iy + 0.5) * fLy;
0109       const S midZ = fMinBound[2] + (vi.iz + 0.5) * fLz;
0110
0111       return {midX, midY, midZ};
0112    }
0113
0114    S getDiagonalLength() const { return fHalfDiag; }
0115
0116    // Convert a point p(x, y, z) to voxel indices (i, j, k)
0117    // if point is outside set indices i,j,k to -1
0118    void pointToVoxelIndex(std::array<S, 3> const &p, int &i, int &j, int &k) const
0119    {
0120       if (!inside(p)) {
0121          i = -1;
0122          j = -1;
0123          k = -1;
0124       }
0125
0126       i = static_cast<int>((p[0] - fMinBound[0]) * finvLx);
0127       j = static_cast<int>((p[1] - fMinBound[1]) * finvLy);
0128       k = static_cast<int>((p[2] - fMinBound[2]) * finvLz);
0129
0130       // Clamp the indices to valid ranges
0131       i = std::min(i, fNx - 1);
0132       j = std::min(j, fNy - 1);
0133       k = std::min(k, fNz - 1);
0134    }
0135
0136    // Convert a point p(x, y, z) to voxel index object
0137    // if outside, an invalid index object will be returned
0138    TGeoVoxelGridIndex pointToVoxelIndex(std::array<S, 3> const &p) const
0139    {
0140       if (!inside(p)) {
0141          return TGeoVoxelGridIndex(); // invalid voxel index
0142       }
0143
0144       int i = static_cast<int>((p[0] - fMinBound[0]) * finvLx);
0145       int j = static_cast<int>((p[1] - fMinBound[1]) * finvLy);
0146       int k = static_cast<int>((p[2] - fMinBound[2]) * finvLz);
0147
0148       // Clamp the indices to valid ranges
0149       i = std::min(i, fNz - 1);
0150       j = std::min(j, fNy - 1);
0151       k = std::min(k, fNz - 1);
0152
0153       return TGeoVoxelGridIndex{i, j, k, index(i, j, k)};
0154    }
0155
0156    TGeoVoxelGridIndex pointToVoxelIndex(S x, S y, S z) const { return pointToVoxelIndex(std::array<S, 3>{x, y, z}); }
0157
0158    // Convert voxel indices (i, j, k) to a linear index in the grid array
0159    size_t index(int i, int j, int k) const { return i + fNx * (j + fNy * k); }
0160
0161    void indexToIndices(size_t idx, int &i, int &j, int &k) const
0162    {
0163       k = idx % fNz;
0164       j = (idx / fNz) % fNy;
0165       i = idx / (fNy * fNz);
0166    }
0167
0168    // member data
0169
0170    std::array<S, 3> fMinBound;
0171    std::array<S, 3> fMaxBound; // 3D bounds for grid structure
0172    S fLx, fLy, fLz;            // Voxel dimensions
0173    S finvLx, finvLy, finvLz;   // inverse voxel dimensions
0174    S fHalfDiag;                // cached value for voxel half diagonal length
0175
0176    int fNx, fNy, fNz;    // Number of voxels in each dimension
0177    std::vector<T> fGrid; // The actual voxel grid data container
0178 };
0179
0180 #endif