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