imstkUniformSpatialGrid.h

/*
** This file is part of the Interactive Medical Simulation Toolkit (iMSTK)
** iMSTK is distributed under the Apache License, Version 2.0.
** See accompanying NOTICE for details.
*/

#pragma once

#include "imstkMath.h"
#include "imstkLogger.h"

#include <array>

namespace imstk
{
template<class CellData>
class UniformSpatialGrid
{
public:
    UniformSpatialGrid() : UniformSpatialGrid(Vec3d::Zero(), Vec3d(1.0, 1.0, 1.0), double(1.0))
    {}

    UniformSpatialGrid(const Vec3d& lowerCorner, const Vec3d& upperCorner, double cellSize)
    {
        initialize(lowerCorner, upperCorner, cellSize);
    }

    void initialize(const Vec3d& lowerCorner, const Vec3d& upperCorner, const double cellSize)
    {
        CHECK(cellSize > 0) << "Invalid cell size";

        m_LowerCorner = lowerCorner;
        m_UpperCorner = upperCorner;

        m_CellSize    = cellSize;
        m_InvCellSize = 1.0 / m_CellSize;

        m_NTotalCells = 1u;
        for (int i = 0; i < 3; ++i)
        {
            m_Resolution[i] = static_cast<unsigned int>(std::ceil((m_UpperCorner[i] - m_LowerCorner[i]) / m_CellSize));
            m_NTotalCells  *= m_Resolution[i];
        }

        CHECK(m_NTotalCells != 0) << "Invalid grid size: [" +
            std::to_string(m_LowerCorner[0]) + ", " + std::to_string(m_LowerCorner[1]) + ", " + std::to_string(m_LowerCorner[2]) + "] => " +
            std::to_string(m_UpperCorner[0]) + ", " + std::to_string(m_UpperCorner[1]) + ", " + std::to_string(m_UpperCorner[2]) + "], " +
            "cellSize = " + std::to_string(m_CellSize);

        // cell data must be resized to match with the number of cells
        m_CellData.resize(m_NTotalCells);
    }

    std::array<unsigned int, 3> getResolution() const { return m_Resolution; }

    unsigned int getNumTotalCells() const { return m_NTotalCells; }

    template<int d>
    bool isValidCellIndex(const int idx) const
    { return idx >= 0 && static_cast<unsigned int>(idx) < m_Resolution[d]; }

    bool isValidCellIndices(const int i, const int j, const int k) const
    { return isValidCellIndex<0>(i) && isValidCellIndex<1>(j) && isValidCellIndex<2>(k); }

    template<class IndexType>
    std::array<IndexType, 3> getCell3DIndices(const Vec3d& ppos) const
    {
        std::array<IndexType, 3> cellIdx;
        for (int d = 0; d < 3; ++d)
        {
            cellIdx[d] = static_cast<IndexType>((ppos[d] - m_LowerCorner[d]) * m_InvCellSize);
        }
        return cellIdx;
    }

    std::vector<CellData>& getAllCellData() { return m_CellData; }

    const std::vector<CellData>& getAllCellData() const { return m_CellData; }

    CellData& getCellData(const Vec3d& ppos)
    {
        return m_CellData[getCellLinearizedIndex<unsigned int>(ppos)];
    }

    const CellData& getCellData(const Vec3d& ppos) const
    { return m_CellData[getCellLinearizedIndex<unsigned int>(ppos)]; }

    CellData& getCellData(size_t linearizedIdx)
    {
        assert(linearizedIdx < m_CellData.size());
        return m_CellData[linearizedIdx];
    }

    const CellData& getCellData(size_t linearizedIdx) const
    {
        assert(linearizedIdx < m_CellData.size());
        return m_CellData[linearizedIdx];
    }

    template<class IndexType>
    CellData& getCellData(const std::array<IndexType, 3>& cellIdx)
    {
        return m_CellData[getCellLinearizedIndex(cellIdx[0], cellIdx[1], cellIdx[2])];
    }

    template<class IndexType>
    const CellData& getCellData(const std::array<IndexType, 3>& cellIdx) const
    {
        return m_CellData[getCellLinearizedIndex(cellIdx[0], cellIdx[1], cellIdx[2])];
    }

    template<class IndexType>
    CellData& getCellData(const IndexType i, const IndexType j, const IndexType k)
    {
        return m_CellData[getCellLinearizedIndex(i, j, k)];
    }

    template<class IndexType>
    const CellData& getCellData(const IndexType i, const IndexType j, const IndexType k) const
    {
        return m_CellData[getCellLinearizedIndex(i, j, k)];
    }

    template<class Function>
    void loopAllCellData(Function&& func)
    {
        for (auto& cellData: m_CellData)
        {
            func(cellData);
        }
    }

    template<class IndexType>
    IndexType getCellLinearizedIndex(const IndexType i, const IndexType j, const IndexType k) const
    {
        auto flatIndex = (k * static_cast<IndexType>(m_Resolution[1]) + j) * static_cast<IndexType>(m_Resolution[0]) + i;
        assert(flatIndex < static_cast<IndexType>(m_NTotalCells));
        return flatIndex;
    }

    template<class IndexType>
    IndexType getCellLinearizedIndex(const Vec3d& ppos) const
    {
        auto cellIdx = getCell3DIndices<IndexType>(ppos);
#if defined(DEBUG) || defined(_DEBUG) || !defined(NDEBUG)
        LOG_IF(FATAL, (!isValidCellIndices(cellIdx[0], cellIdx[1], cellIdx[2]))) <<
            "Invalid cell indices: " +
            std::to_string(cellIdx[0]) + "/" + std::to_string(m_Resolution[0]) + ", " +
            std::to_string(cellIdx[1]) + "/" + std::to_string(m_Resolution[1]) + ", " +
            std::to_string(cellIdx[2]) + "/" + std::to_string(m_Resolution[2]);
#endif
        return getCellLinearizedIndex<IndexType>(cellIdx[0], cellIdx[1], cellIdx[2]);
    }

private:
    Vec3d  m_LowerCorner;                     
    Vec3d  m_UpperCorner;                     
    double m_CellSize;                        
    double m_InvCellSize;                     

    unsigned int m_NTotalCells;               
    std::array<unsigned int, 3> m_Resolution; 

    std::vector<CellData> m_CellData;         
};
} // namespace imstk