Turn nbnxn_grid_t into a class
[alexxy/gromacs.git] / src / gromacs / nbnxm / grid.cpp
index 3e91161f2ccb8533dc3ab7a019241b571c25d747..e971e0f34b306d7c6bd0a08c84df122ee9e573f4 100644 (file)
  * the research papers on the package. Check out http://www.gromacs.org.
  */
 
+/*! \internal \file
+ *
+ * \brief
+ * Implements the Grid class.
+ *
+ * \author Berk Hess <hess@kth.se>
+ * \ingroup module_nbnxm
+ */
+
 #include "gmxpre.h"
 
 #include "grid.h"
@@ -51,6 +60,8 @@
 #include "gromacs/nbnxm/atomdata.h"
 #include "gromacs/nbnxm/nbnxm.h"
 #include "gromacs/nbnxm/nbnxm_geometry.h"
+#include "gromacs/nbnxm/pairlist.h"
+#include "gromacs/nbnxm/pairlistset.h"
 #include "gromacs/simd/simd.h"
 #include "gromacs/simd/vector_operations.h"
 #include "gromacs/utility/exceptions.h"
 
 struct gmx_domdec_zones_t;
 
-static real grid_atom_density(int        numAtoms,
-                              const rvec lowerCorner,
-                              const rvec upperCorner)
+namespace Nbnxm
+{
+
+Grid::Geometry::Geometry(const PairlistType pairlistType) :
+    isSimple(pairlistType != PairlistType::Hierarchical8x8),
+    numAtomsICluster(IClusterSizePerListType[pairlistType]),
+    numAtomsJCluster(JClusterSizePerListType[pairlistType]),
+    numAtomsPerCell((isSimple ? 1 : c_gpuNumClusterPerCell)*numAtomsICluster),
+    numAtomsICluster2Log(get_2log(numAtomsICluster))
+{
+}
+
+Grid::Grid(const PairlistType pairlistType) :
+    geometry_(pairlistType)
+{
+}
+
+/*! \brief Returns the atom density (> 0) of a rectangular grid */
+static real gridAtomDensity(int        numAtoms,
+                            const rvec lowerCorner,
+                            const rvec upperCorner)
 {
     rvec size;
 
@@ -77,60 +106,70 @@ static real grid_atom_density(int        numAtoms,
     return numAtoms/(size[XX]*size[YY]*size[ZZ]);
 }
 
-static void set_grid_size_xy(const nbnxn_search   *nbs,
-                             nbnxn_grid_t         *grid,
-                             int                   ddZone,
-                             int                   numAtoms,
-                             const rvec            lowerCorner,
-                             const rvec            upperCorner,
-                             real                  atomDensity)
+void Grid::setDimensions(const nbnxn_search *nbs,
+                         const int           ddZone,
+                         const int           numAtoms,
+                         const rvec          lowerCorner,
+                         const rvec          upperCorner,
+                         real                atomDensity,
+                         const real          maxAtomGroupRadius)
 {
-    rvec size;
-    real tlen, tlen_x, tlen_y;
+    /* For the home zone we compute the density when not set (=-1) or when =0 */
+    if (ddZone == 0 && atomDensity <= 0)
+    {
+        atomDensity = gridAtomDensity(numAtoms, lowerCorner, upperCorner);
+    }
 
+    dimensions_.atomDensity        = atomDensity;
+    dimensions_.maxAtomGroupRadius = maxAtomGroupRadius;
+
+    rvec size;
     rvec_sub(upperCorner, lowerCorner, size);
 
-    if (numAtoms > grid->na_sc)
+    if (numAtoms > geometry_.numAtomsPerCell)
     {
         GMX_ASSERT(atomDensity > 0, "With one or more atoms, the density should be positive");
 
         /* target cell length */
-        if (grid->bSimple)
+        real tlen_x;
+        real tlen_y;
+        if (geometry_.isSimple)
         {
             /* To minimize the zero interactions, we should make
              * the largest of the i/j cell cubic.
              */
-            int numAtomsInCell = std::max(grid->na_c, grid->na_cj);
+            int numAtomsInCell = std::max(geometry_.numAtomsICluster,
+                                          geometry_.numAtomsJCluster);
 
             /* Approximately cubic cells */
-            tlen   = std::cbrt(numAtomsInCell/atomDensity);
-            tlen_x = tlen;
-            tlen_y = tlen;
+            real tlen = std::cbrt(numAtomsInCell/atomDensity);
+            tlen_x    = tlen;
+            tlen_y    = tlen;
         }
         else
         {
             /* Approximately cubic sub cells */
-            tlen   = std::cbrt(grid->na_c/atomDensity);
-            tlen_x = tlen*c_gpuNumClusterPerCellX;
-            tlen_y = tlen*c_gpuNumClusterPerCellY;
+            real tlen = std::cbrt(geometry_.numAtomsICluster/atomDensity);
+            tlen_x    = tlen*c_gpuNumClusterPerCellX;
+            tlen_y    = tlen*c_gpuNumClusterPerCellY;
         }
         /* We round ncx and ncy down, because we get less cell pairs
          * in the nbsist when the fixed cell dimensions (x,y) are
          * larger than the variable one (z) than the other way around.
          */
-        grid->numCells[XX] = std::max(1, static_cast<int>(size[XX]/tlen_x));
-        grid->numCells[YY] = std::max(1, static_cast<int>(size[YY]/tlen_y));
+        dimensions_.numCells[XX] = std::max(1, static_cast<int>(size[XX]/tlen_x));
+        dimensions_.numCells[YY] = std::max(1, static_cast<int>(size[YY]/tlen_y));
     }
     else
     {
-        grid->numCells[XX] = 1;
-        grid->numCells[YY] = 1;
+        dimensions_.numCells[XX] = 1;
+        dimensions_.numCells[YY] = 1;
     }
 
     for (int d = 0; d < DIM - 1; d++)
     {
-        grid->cellSize[d]    = size[d]/grid->numCells[d];
-        grid->invCellSize[d] = 1/grid->cellSize[d];
+        dimensions_.cellSize[d]    = size[d]/dimensions_.numCells[d];
+        dimensions_.invCellSize[d] = 1/dimensions_.cellSize[d];
     }
 
     if (ddZone > 0)
@@ -141,76 +180,77 @@ static void set_grid_size_xy(const nbnxn_search   *nbs,
          * they end up on the grid, but for performance it's better
          * if they don't end up in cells that can be within cut-off range.
          */
-        grid->numCells[XX]++;
-        grid->numCells[YY]++;
+        dimensions_.numCells[XX]++;
+        dimensions_.numCells[YY]++;
     }
 
     /* We need one additional cell entry for particles moved by DD */
-    int numCellsXY = grid->numCells[XX]*grid->numCells[YY];
-    grid->cxy_na.resize(numCellsXY + 1);
-    grid->cxy_ind.resize(numCellsXY + 2);
+    cxy_na_.resize(numColumns() + 1);
+    cxy_ind_.resize(numColumns() + 2);
 
     for (nbnxn_search_work_t &work : nbs->work)
     {
-        work.cxy_na.resize(numCellsXY + 1);
+        work.cxy_na.resize(numColumns() + 1);
     }
 
     /* Worst case scenario of 1 atom in each last cell */
     int maxNumCells;
-    if (grid->na_cj <= grid->na_c)
+    if (geometry_.numAtomsJCluster <= geometry_.numAtomsICluster)
     {
-        maxNumCells = numAtoms/grid->na_sc + numCellsXY;
+        maxNumCells = numAtoms/geometry_.numAtomsPerCell + numColumns();
     }
     else
     {
-        maxNumCells = numAtoms/grid->na_sc + numCellsXY*grid->na_cj/grid->na_c;
+        maxNumCells = numAtoms/geometry_.numAtomsPerCell + numColumns()*geometry_.numAtomsJCluster/geometry_.numAtomsICluster;
     }
 
-    grid->nsubc.resize(maxNumCells);
-    grid->bbcz.resize(maxNumCells*NNBSBB_D);
+    if (!geometry_.isSimple)
+    {
+        numClusters_.resize(maxNumCells);
+    }
+    bbcz_.resize(maxNumCells*NNBSBB_D);
 
     /* This resize also zeros the contents, this avoid possible
      * floating exceptions in SIMD with the unused bb elements.
      */
-    if (grid->bSimple)
+    if (geometry_.isSimple)
     {
-        grid->bb.resize(maxNumCells);
+        bb_.resize(maxNumCells);
     }
     else
     {
 #if NBNXN_BBXXXX
-        grid->pbb.resize(maxNumCells*c_gpuNumClusterPerCell/STRIDE_PBB*NNBSBB_XXXX);
+        pbb_.resize(maxNumCells*c_gpuNumClusterPerCell/STRIDE_PBB*NNBSBB_XXXX);
 #else
-        grid->bb.resize(maxNumCells*c_gpuNumClusterPerCell);
+        bb_.resize(maxNumCells*c_gpuNumClusterPerCell);
 #endif
     }
 
-    if (grid->bSimple)
+    if (geometry_.numAtomsJCluster == geometry_.numAtomsICluster)
     {
-        if (grid->na_cj == grid->na_c)
-        {
-            grid->bbj = grid->bb;
-        }
-        else
-        {
-            grid->bbjStorage.resize(maxNumCells*grid->na_c/grid->na_cj);
-            grid->bbj = grid->bbjStorage;
-        }
+        bbj_ = bb_;
+    }
+    else
+    {
+        GMX_ASSERT(geometry_.isSimple, "Only CPU lists should have different i/j cluster sizes");
+
+        bbjStorage_.resize(maxNumCells*geometry_.numAtomsICluster/geometry_.numAtomsJCluster);
+        bbj_ = bbjStorage_;
     }
 
-    grid->flags.resize(maxNumCells);
+    flags_.resize(maxNumCells);
     if (nbs->bFEP)
     {
-        grid->fep.resize(maxNumCells*grid->na_sc/grid->na_c);
+        fep_.resize(maxNumCells*geometry_.numAtomsPerCell/geometry_.numAtomsICluster);
     }
 
-    copy_rvec(lowerCorner, grid->c0);
-    copy_rvec(upperCorner, grid->c1);
-    copy_rvec(size,        grid->size);
+    copy_rvec(lowerCorner, dimensions_.lowerCorner);
+    copy_rvec(upperCorner, dimensions_.upperCorner);
+    copy_rvec(size,        dimensions_.gridSize);
 }
 
 /* We need to sort paricles in grid columns on z-coordinate.
- * As particle are very often distributed homogeneously, we a sorting
+ * As particle are very often distributed homogeneously, we use a sorting
  * algorithm similar to pigeonhole sort. We multiply the z-coordinate
  * by a factor, cast to an int and try to store in that hole. If the hole
  * is full, we move this or another particle. A second pass is needed to make
@@ -219,20 +259,23 @@ static void set_grid_size_xy(const nbnxn_search   *nbs,
  * for an O(#particles) sort up till distributions were all particles are
  * concentrated in 1/4 of the space. No NlogN fallback is implemented,
  * as it can be expensive to detect imhomogeneous particle distributions.
- * SGSF is the maximum ratio of holes used, in the worst case all particles
- * end up in the last hole and we need #particles extra holes at the end.
  */
-#define SORT_GRID_OVERSIZE 4
-#define SGSF (SORT_GRID_OVERSIZE + 1)
+/*! \brief Ratio of grid cells to atoms */
+static constexpr int c_sortGridRatio         = 4;
+/*! \brief Maximum ratio of holes used, in the worst case all particles
+ * end up in the last hole and we need num. atoms extra holes at the end.
+ */
+static constexpr int c_sortGridMaxSizeFactor = c_sortGridRatio + 1;
 
-/* Sort particle index a on coordinates x along dim.
+/*! \brief Sorts particle index a on coordinates x along dim.
+ *
  * Backwards tells if we want decreasing iso increasing coordinates.
  * h0 is the minimum of the coordinate range.
  * invh is the 1/length of the sorting range.
  * n_per_h (>=n) is the expected average number of particles per 1/invh
  * sort is the sorting work array.
- * sort should have a size of at least n_per_h*SORT_GRID_OVERSIZE + n,
- * or easier, allocate at least n*SGSF elements.
+ * sort should have a size of at least n_per_h*c_sortGridRatio + n,
+ * or easier, allocate at least n*c_sortGridMaxSizeFactor elements.
  */
 static void sort_atoms(int dim, gmx_bool Backwards,
                        int gmx_unused dd_zone,
@@ -251,12 +294,12 @@ static void sort_atoms(int dim, gmx_bool Backwards,
     GMX_ASSERT(n <= n_per_h, "We require n <= n_per_h");
 
     /* Transform the inverse range height into the inverse hole height */
-    invh *= n_per_h*SORT_GRID_OVERSIZE;
+    invh *= n_per_h*c_sortGridRatio;
 
     /* Set nsort to the maximum possible number of holes used.
      * In worst case all n elements end up in the last bin.
      */
-    int nsort = n_per_h*SORT_GRID_OVERSIZE + n;
+    int nsort = n_per_h*c_sortGridRatio + n;
 
     /* Determine the index range used, so we can limit it for the second pass */
     int zi_min = INT_MAX;
@@ -274,11 +317,11 @@ static void sort_atoms(int dim, gmx_bool Backwards,
 #ifndef NDEBUG
         /* As we can have rounding effect, we use > iso >= here */
         if (relevantAtomsAreWithinGridBounds &&
-            (zi < 0 || (dd_zone == 0 && zi > n_per_h*SORT_GRID_OVERSIZE)))
+            (zi < 0 || (dd_zone == 0 && zi > n_per_h*c_sortGridRatio)))
         {
             gmx_fatal(FARGS, "(int)((x[%d][%c]=%f - %f)*%f) = %d, not in 0 - %d*%d\n",
                       a[i], 'x'+dim, x[a[i]][dim], h0, invh, zi,
-                      n_per_h, SORT_GRID_OVERSIZE);
+                      n_per_h, c_sortGridRatio);
         }
 #endif
         if (zi < 0)
@@ -290,9 +333,9 @@ static void sort_atoms(int dim, gmx_bool Backwards,
          * can be far beyond the grid size, which is set by the non-bonded
          * cut-off distance. We sort such particles into the last cell.
          */
-        if (zi > n_per_h*SORT_GRID_OVERSIZE)
+        if (zi > n_per_h*c_sortGridRatio)
         {
-            zi = n_per_h*SORT_GRID_OVERSIZE;
+            zi = n_per_h*c_sortGridRatio;
         }
 
         /* Ideally this particle should go in sort cell zi,
@@ -370,15 +413,32 @@ static void sort_atoms(int dim, gmx_bool Backwards,
 }
 
 #if GMX_DOUBLE
-#define R2F_D(x) ((float)((x) >= 0 ? ((1-GMX_FLOAT_EPS)*(x)) : ((1+GMX_FLOAT_EPS)*(x))))
-#define R2F_U(x) ((float)((x) >= 0 ? ((1+GMX_FLOAT_EPS)*(x)) : ((1-GMX_FLOAT_EPS)*(x))))
+//! Returns double up to one least significant float bit smaller than x
+static double R2F_D(const float x)
+{
+    return static_cast<float>(x >= 0 ? (1 - GMX_FLOAT_EPS)*x : (1 + GMX_FLOAT_EPS)*x);
+}
+//! Returns double up to one least significant float bit larger than x
+static double R2F_U(const float x)
+{
+    return static_cast<float>(x >= 0 ? (1 + GMX_FLOAT_EPS)*x : (1 - GMX_FLOAT_EPS)*x);
+}
 #else
-#define R2F_D(x) (x)
-#define R2F_U(x) (x)
+//! Returns x
+static float R2F_D(const float x)
+{
+    return x;
+}
+//! Returns x
+static float R2F_U(const float x)
+{
+    return x;
+}
 #endif
 
-/* Coordinate order x,y,z, bb order xyz0 */
-static void calc_bounding_box(int na, int stride, const real *x, nbnxn_bb_t *bb)
+//! Computes the bounding box for na coordinates in order x,y,z, bb order xyz0
+static void calc_bounding_box(int na, int stride, const real *x,
+                              BoundingBox *bb)
 {
     int  i;
     real xl, xh, yl, yh, zl, zh;
@@ -410,8 +470,9 @@ static void calc_bounding_box(int na, int stride, const real *x, nbnxn_bb_t *bb)
     bb->upper[BB_Z] = R2F_U(zh);
 }
 
-/* Packed coordinates, bb order xyz0 */
-static void calc_bounding_box_x_x4(int na, const real *x, nbnxn_bb_t *bb)
+/*! \brief Computes the bounding box for na packed coordinates, bb order xyz0 */
+static void calc_bounding_box_x_x4(int na, const real *x,
+                                   BoundingBox *bb)
 {
     real xl, xh, yl, yh, zl, zh;
 
@@ -439,8 +500,9 @@ static void calc_bounding_box_x_x4(int na, const real *x, nbnxn_bb_t *bb)
     bb->upper[BB_Z] = R2F_U(zh);
 }
 
-/* Packed coordinates, bb order xyz0 */
-static void calc_bounding_box_x_x8(int na, const real *x, nbnxn_bb_t *bb)
+/*! \brief Computes the bounding box for na coordinates, bb order xyz0 */
+static void calc_bounding_box_x_x8(int na, const real *x,
+                                   BoundingBox *bb)
 {
     real xl, xh, yl, yh, zl, zh;
 
@@ -468,9 +530,10 @@ static void calc_bounding_box_x_x8(int na, const real *x, nbnxn_bb_t *bb)
     bb->upper[BB_Z] = R2F_U(zh);
 }
 
-/* Packed coordinates, bb order xyz0 */
+/*! \brief Computes the bounding box for na packed coordinates, bb order xyz0 */
 gmx_unused static void calc_bounding_box_x_x4_halves(int na, const real *x,
-                                                     nbnxn_bb_t *bb, nbnxn_bb_t *bbj)
+                                                     BoundingBox *bb,
+                                                     BoundingBox *bbj)
 {
     // TODO: During SIMDv2 transition only some archs use namespace (remove when done)
     using namespace gmx;
@@ -512,7 +575,7 @@ gmx_unused static void calc_bounding_box_x_x4_halves(int na, const real *x,
 
 #if NBNXN_SEARCH_BB_SIMD4
 
-/* Coordinate order xyz, bb order xxxxyyyyzzzz */
+/*! \brief Computes the bounding box for na coordinates in order xyz, bb order xxxxyyyyzzzz */
 static void calc_bounding_box_xxxx(int na, int stride, const real *x, float *bb)
 {
     int  i;
@@ -549,8 +612,9 @@ static void calc_bounding_box_xxxx(int na, int stride, const real *x, float *bb)
 
 #if NBNXN_SEARCH_SIMD4_FLOAT_X_BB
 
-/* Coordinate order xyz?, bb order xyz0 */
-static void calc_bounding_box_simd4(int na, const float *x, nbnxn_bb_t *bb)
+/*! \brief Computes the bounding box for na coordinates in order xyz?, bb order xyz0 */
+static void calc_bounding_box_simd4(int na, const float *x,
+                                    BoundingBox *bb)
 {
     // TODO: During SIMDv2 transition only some archs use namespace (remove when done)
     using namespace gmx;
@@ -572,9 +636,9 @@ static void calc_bounding_box_simd4(int na, const float *x, nbnxn_bb_t *bb)
     store4(&bb->upper[0], bb_1_S);
 }
 
-/* Coordinate order xyz?, bb order xxxxyyyyzzzz */
+/*! \brief Computes the bounding box for na coordinates in order xyz?, bb order xxxxyyyyzzzz */
 static void calc_bounding_box_xxxx_simd4(int na, const float *x,
-                                         nbnxn_bb_t *bb_work_aligned,
+                                         BoundingBox *bb_work_aligned,
                                          real *bb)
 {
     calc_bounding_box_simd4(na, x, bb_work_aligned);
@@ -590,21 +654,22 @@ static void calc_bounding_box_xxxx_simd4(int na, const float *x,
 #endif /* NBNXN_SEARCH_SIMD4_FLOAT_X_BB */
 
 
-/* Combines pairs of consecutive bounding boxes */
-static void combine_bounding_box_pairs(nbnxn_grid_t                    *grid,
-                                       gmx::ArrayRef<const nbnxn_bb_t>  bb)
+/*! \brief Combines pairs of consecutive bounding boxes */
+static void combine_bounding_box_pairs(const Grid                       &grid,
+                                       gmx::ArrayRef<const BoundingBox>  bb,
+                                       gmx::ArrayRef<BoundingBox>        bbj)
 {
     // TODO: During SIMDv2 transition only some archs use namespace (remove when done)
     using namespace gmx;
 
-    for (int i = 0; i < grid->numCells[XX]*grid->numCells[YY]; i++)
+    for (int i = 0; i < grid.numColumns(); i++)
     {
         /* Starting bb in a column is expected to be 2-aligned */
-        int sc2 = grid->cxy_ind[i]>>1;
+        const int sc2 = grid.firstCellInColumn(i) >> 1;
         /* For odd numbers skip the last bb here */
-        int nc2 = (grid->cxy_na[i]+3)>>(2+1);
-        int c2;
-        for (c2 = sc2; c2 < sc2+nc2; c2++)
+        const int nc2 = (grid.numAtomsInColumn(i) + 3) >> (2 + 1);
+        int       c2;
+        for (c2 = sc2; c2 < sc2 + nc2; c2++)
         {
 #if NBNXN_SEARCH_BB_SIMD4
             Simd4Float min_S, max_S;
@@ -613,72 +678,73 @@ static void combine_bounding_box_pairs(nbnxn_grid_t                    *grid,
                         load4(&bb[c2*2+1].lower[0]));
             max_S = max(load4(&bb[c2*2+0].upper[0]),
                         load4(&bb[c2*2+1].upper[0]));
-            store4(&grid->bbj[c2].lower[0], min_S);
-            store4(&grid->bbj[c2].upper[0], max_S);
+            store4(&bbj[c2].lower[0], min_S);
+            store4(&bbj[c2].upper[0], max_S);
 #else
             for (int j = 0; j < NNBSBB_C; j++)
             {
-                grid->bbj[c2].lower[j] = std::min(bb[c2*2+0].lower[j],
-                                                  bb[c2*2+1].lower[j]);
-                grid->bbj[c2].upper[j] = std::max(bb[c2*2+0].upper[j],
-                                                  bb[c2*2+1].upper[j]);
+                bbj[c2].lower[j] = std::min(bb[c2*2 + 0].lower[j],
+                                            bb[c2*2 + 1].lower[j]);
+                bbj[c2].upper[j] = std::max(bb[c2*2 + 0].upper[j],
+                                            bb[c2*2 + 1].upper[j]);
             }
 #endif
         }
-        if (((grid->cxy_na[i]+3)>>2) & 1)
+        if (((grid.numAtomsInColumn(i) + 3) >> 2) & 1)
         {
             /* The bb count in this column is odd: duplicate the last bb */
             for (int j = 0; j < NNBSBB_C; j++)
             {
-                grid->bbj[c2].lower[j] = bb[c2*2].lower[j];
-                grid->bbj[c2].upper[j] = bb[c2*2].upper[j];
+                bbj[c2].lower[j] = bb[c2*2].lower[j];
+                bbj[c2].upper[j] = bb[c2*2].upper[j];
             }
         }
     }
 }
 
 
-/* Prints the average bb size, used for debug output */
-static void print_bbsizes_simple(FILE                *fp,
-                                 const nbnxn_grid_t  *grid)
+/*! \brief Prints the average bb size, used for debug output */
+static void print_bbsizes_simple(FILE       *fp,
+                                 const Grid &grid)
 {
     dvec ba;
 
     clear_dvec(ba);
-    for (int c = 0; c < grid->nc; c++)
+    for (int c = 0; c < grid.numCells(); c++)
     {
         for (int d = 0; d < DIM; d++)
         {
-            ba[d] += grid->bb[c].upper[d] - grid->bb[c].lower[d];
+            ba[d] += grid.iBoundingBoxes()[c].upper[d] - grid.iBoundingBoxes()[c].lower[d];
         }
     }
-    dsvmul(1.0/grid->nc, ba, ba);
+    dsvmul(1.0/grid.numCells(), ba, ba);
 
-    real avgCellSizeZ = (grid->atom_density > 0 ? grid->na_c/(grid->atom_density*grid->cellSize[XX]*grid->cellSize[YY]) : 0.0);
+    const Grid::Dimensions &dims         = grid.dimensions();
+    real                    avgCellSizeZ = (dims.atomDensity > 0 ? grid.geometry().numAtomsICluster/(dims.atomDensity*dims.cellSize[XX]*dims.cellSize[YY]) : 0.0);
 
     fprintf(fp, "ns bb: grid %4.2f %4.2f %4.2f abs %4.2f %4.2f %4.2f rel %4.2f %4.2f %4.2f\n",
-            grid->cellSize[XX],
-            grid->cellSize[YY],
+            dims.cellSize[XX],
+            dims.cellSize[YY],
             avgCellSizeZ,
             ba[XX], ba[YY], ba[ZZ],
-            ba[XX]*grid->invCellSize[XX],
-            ba[YY]*grid->invCellSize[YY],
-            grid->atom_density > 0 ? ba[ZZ]/avgCellSizeZ : 0.0);
+            ba[XX]*dims.invCellSize[XX],
+            ba[YY]*dims.invCellSize[YY],
+            dims.atomDensity > 0 ? ba[ZZ]/avgCellSizeZ : 0.0);
 }
 
-/* Prints the average bb size, used for debug output */
-static void print_bbsizes_supersub(FILE                *fp,
-                                   const nbnxn_grid_t  *grid)
+/*! \brief Prints the average bb size, used for debug output */
+static void print_bbsizes_supersub(FILE       *fp,
+                                   const Grid &grid)
 {
     int  ns;
     dvec ba;
 
     clear_dvec(ba);
     ns = 0;
-    for (int c = 0; c < grid->nc; c++)
+    for (int c = 0; c < grid.numCells(); c++)
     {
 #if NBNXN_BBXXXX
-        for (int s = 0; s < grid->nsubc[c]; s += STRIDE_PBB)
+        for (int s = 0; s < grid.numClustersPerCell()[c]; s += STRIDE_PBB)
         {
             int cs_w = (c*c_gpuNumClusterPerCell + s)/STRIDE_PBB;
             for (int i = 0; i < STRIDE_PBB; i++)
@@ -686,38 +752,41 @@ static void print_bbsizes_supersub(FILE                *fp,
                 for (int d = 0; d < DIM; d++)
                 {
                     ba[d] +=
-                        grid->pbb[cs_w*NNBSBB_XXXX+(DIM+d)*STRIDE_PBB+i] -
-                        grid->pbb[cs_w*NNBSBB_XXXX+     d *STRIDE_PBB+i];
+                        grid.packedBoundingBoxes()[cs_w*NNBSBB_XXXX + (DIM + d)*STRIDE_PBB + i] -
+                        grid.packedBoundingBoxes()[cs_w*NNBSBB_XXXX +        d *STRIDE_PBB + i];
                 }
             }
         }
 #else
-        for (int s = 0; s < grid->nsubc[c]; s++)
+        for (int s = 0; s < grid.numClustersPerCell()[c]; s++)
         {
             int cs = c*c_gpuNumClusterPerCell + s;
             for (int d = 0; d < DIM; d++)
             {
-                ba[d] += grid->bb[cs].upper[d] - grid->bb[cs].lower[d];
+                ba[d] += grid.iBoundingBoxes()[cs].upper[d] - grid.iBoundingBoxes()[cs].lower[d];
             }
         }
 #endif
-        ns += grid->nsubc[c];
+        ns += grid.numClustersPerCell()[c];
     }
     dsvmul(1.0/ns, ba, ba);
 
-    real avgClusterSizeZ = (grid->atom_density > 0 ? grid->na_sc/(grid->atom_density*grid->cellSize[XX]*grid->cellSize[YY]*c_gpuNumClusterPerCellZ) : 0.0);
+    const Grid::Dimensions &dims            = grid.dimensions();
+    const real              avgClusterSizeZ =
+        (dims.atomDensity > 0 ? grid.geometry().numAtomsPerCell/(dims.atomDensity*dims.cellSize[XX]*dims.cellSize[YY]*c_gpuNumClusterPerCellZ) : 0.0);
 
     fprintf(fp, "ns bb: grid %4.2f %4.2f %4.2f abs %4.2f %4.2f %4.2f rel %4.2f %4.2f %4.2f\n",
-            grid->cellSize[XX]/c_gpuNumClusterPerCellX,
-            grid->cellSize[YY]/c_gpuNumClusterPerCellY,
+            dims.cellSize[XX]/c_gpuNumClusterPerCellX,
+            dims.cellSize[YY]/c_gpuNumClusterPerCellY,
             avgClusterSizeZ,
             ba[XX], ba[YY], ba[ZZ],
-            ba[XX]*c_gpuNumClusterPerCellX*grid->invCellSize[XX],
-            ba[YY]*c_gpuNumClusterPerCellY*grid->invCellSize[YY],
-            grid->atom_density > 0 ? ba[ZZ]/avgClusterSizeZ : 0.0);
+            ba[XX]*c_gpuNumClusterPerCellX*dims.invCellSize[XX],
+            ba[YY]*c_gpuNumClusterPerCellY*dims.invCellSize[YY],
+            dims.atomDensity > 0 ? ba[ZZ]/avgClusterSizeZ : 0.0);
 }
 
-/* Set non-bonded interaction flags for the current cluster.
+/*!\brief Set non-bonded interaction flags for the current cluster.
+ *
  * Sorts atoms on LJ coefficients: !=0 first, ==0 at the end.
  */
 static void sort_cluster_on_flag(int                 numAtomsInCluster,
@@ -792,28 +861,28 @@ static void sort_cluster_on_flag(int                 numAtomsInCluster,
     }
 }
 
-/* Fill a pair search cell with atoms.
+/*! \brief Fill a pair search cell with atoms.
+ *
  * Potentially sorts atoms and sets the interaction flags.
  */
-static void fill_cell(nbnxn_search                  *nbs,
-                      nbnxn_grid_t                  *grid,
-                      nbnxn_atomdata_t              *nbat,
-                      int                            atomStart,
-                      int                            atomEnd,
-                      const int                     *atinfo,
-                      gmx::ArrayRef<const gmx::RVec> x,
-                      nbnxn_bb_t gmx_unused         *bb_work_aligned)
+void Grid::fillCell(nbnxn_search                   *nbs,
+                    nbnxn_atomdata_t               *nbat,
+                    int                             atomStart,
+                    int                             atomEnd,
+                    const int                      *atinfo,
+                    gmx::ArrayRef<const gmx::RVec>  x,
+                    BoundingBox gmx_unused         *bb_work_aligned)
 {
     const int numAtoms = atomEnd - atomStart;
 
-    if (grid->bSimple)
+    if (geometry_.isSimple)
     {
         /* Note that non-local grids are already sorted.
          * Then sort_cluster_on_flag will only set the flags and the sorting
          * will not affect the atom order.
          */
-        sort_cluster_on_flag(grid->na_c, atomStart, atomEnd, atinfo, nbs->a,
-                             grid->flags.data() + (atomStart >> grid->na_c_2log) - grid->cell0);
+        sort_cluster_on_flag(geometry_.numAtomsICluster, atomStart, atomEnd, atinfo, nbs->a,
+                             flags_.data() + atomToCluster(atomStart) - cellOffset_);
     }
 
     if (nbs->bFEP)
@@ -821,13 +890,13 @@ static void fill_cell(nbnxn_search                  *nbs,
         /* Set the fep flag for perturbed atoms in this (sub-)cell */
 
         /* The grid-local cluster/(sub-)cell index */
-        int cell = (atomStart >> grid->na_c_2log) - grid->cell0*(grid->bSimple ? 1 : c_gpuNumClusterPerCell);
-        grid->fep[cell] = 0;
+        int cell = atomToCluster(atomStart) - cellOffset_*(geometry_.isSimple ? 1 : c_gpuNumClusterPerCell);
+        fep_[cell] = 0;
         for (int at = atomStart; at < atomEnd; at++)
         {
             if (nbs->a[at] >= 0 && GET_CGINFO_FEP(atinfo[nbs->a[at]]))
             {
-                grid->fep[cell] |= (1 << (at - atomStart));
+                fep_[cell] |= (1 << (at - atomStart));
             }
         }
     }
@@ -838,21 +907,21 @@ static void fill_cell(nbnxn_search                  *nbs,
         nbs->cell[nbs->a[at]] = at;
     }
 
-    copy_rvec_to_nbat_real(nbs->a.data() + atomStart, numAtoms, grid->na_c,
+    copy_rvec_to_nbat_real(nbs->a.data() + atomStart, numAtoms, geometry_.numAtomsICluster,
                            as_rvec_array(x.data()),
                            nbat->XFormat, nbat->x().data(), atomStart);
 
     if (nbat->XFormat == nbatX4)
     {
         /* Store the bounding boxes as xyz.xyz. */
-        size_t      offset = (atomStart - grid->cell0*grid->na_sc) >> grid->na_c_2log;
-        nbnxn_bb_t *bb_ptr = grid->bb.data() + offset;
+        size_t       offset = atomToCluster(atomStart - cellOffset_*geometry_.numAtomsICluster);
+        BoundingBox *bb_ptr = bb_.data() + offset;
 
 #if GMX_SIMD && GMX_SIMD_REAL_WIDTH == 2
-        if (2*grid->na_cj == grid->na_c)
+        if (2*geometry_.numAtomsJCluster == geometry_.numAtomsICluster)
         {
             calc_bounding_box_x_x4_halves(numAtoms, nbat->x().data() + atom_to_x_index<c_packX4>(atomStart), bb_ptr,
-                                          grid->bbj.data() + offset*2);
+                                          bbj_.data() + offset*2);
         }
         else
 #endif
@@ -863,21 +932,21 @@ static void fill_cell(nbnxn_search                  *nbs,
     else if (nbat->XFormat == nbatX8)
     {
         /* Store the bounding boxes as xyz.xyz. */
-        size_t      offset = (atomStart - grid->cell0*grid->na_sc) >> grid->na_c_2log;
-        nbnxn_bb_t *bb_ptr = grid->bb.data() + offset;
+        size_t       offset = atomToCluster(atomStart - cellOffset_*geometry_.numAtomsICluster);
+        BoundingBox *bb_ptr = bb_.data() + offset;
 
         calc_bounding_box_x_x8(numAtoms, nbat->x().data() + atom_to_x_index<c_packX8>(atomStart), bb_ptr);
     }
 #if NBNXN_BBXXXX
-    else if (!grid->bSimple)
+    else if (!geometry_.isSimple)
     {
         /* Store the bounding boxes in a format convenient
          * for SIMD4 calculations: xxxxyyyyzzzz...
          */
         float *pbb_ptr =
-            grid->pbb.data() +
-            ((atomStart - grid->cell0*grid->na_sc) >> (grid->na_c_2log + STRIDE_PBB_2LOG))*NNBSBB_XXXX +
-            (((atomStart - grid->cell0*grid->na_sc) >> grid->na_c_2log) & (STRIDE_PBB - 1));
+            pbb_.data() +
+            ((atomStart - cellOffset_*geometry_.numAtomsPerCell) >> (geometry_.numAtomsICluster2Log + STRIDE_PBB_2LOG))*NNBSBB_XXXX +
+            (((atomStart - cellOffset_*geometry_.numAtomsPerCell) >> geometry_.numAtomsICluster2Log) & (STRIDE_PBB - 1));
 
 #if NBNXN_SEARCH_SIMD4_FLOAT_X_BB
         if (nbat->XFormat == nbatXYZQ)
@@ -894,7 +963,7 @@ static void fill_cell(nbnxn_search                  *nbs,
         if (gmx_debug_at)
         {
             fprintf(debug, "cell %4d bb %5.2f %5.2f %5.2f %5.2f %5.2f %5.2f\n",
-                    atomStart >> grid->na_c_2log,
+                    atomToCluster(atomStart),
                     pbb_ptr[0*STRIDE_PBB], pbb_ptr[3*STRIDE_PBB],
                     pbb_ptr[1*STRIDE_PBB], pbb_ptr[4*STRIDE_PBB],
                     pbb_ptr[2*STRIDE_PBB], pbb_ptr[5*STRIDE_PBB]);
@@ -904,97 +973,97 @@ static void fill_cell(nbnxn_search                  *nbs,
     else
     {
         /* Store the bounding boxes as xyz.xyz. */
-        nbnxn_bb_t *bb_ptr = grid->bb.data() + ((atomStart - grid->cell0*grid->na_sc) >> grid->na_c_2log);
+        BoundingBox *bb_ptr = bb_.data() + atomToCluster(atomStart - cellOffset_*geometry_.numAtomsPerCell);
 
         calc_bounding_box(numAtoms, nbat->xstride, nbat->x().data() + atomStart*nbat->xstride,
                           bb_ptr);
 
         if (gmx_debug_at)
         {
-            int bbo = (atomStart - grid->cell0*grid->na_sc)/grid->na_c;
+            int bbo = atomToCluster(atomStart - cellOffset_*geometry_.numAtomsPerCell);
             fprintf(debug, "cell %4d bb %5.2f %5.2f %5.2f %5.2f %5.2f %5.2f\n",
-                    atomStart >> grid->na_c_2log,
-                    grid->bb[bbo].lower[BB_X],
-                    grid->bb[bbo].lower[BB_Y],
-                    grid->bb[bbo].lower[BB_Z],
-                    grid->bb[bbo].upper[BB_X],
-                    grid->bb[bbo].upper[BB_Y],
-                    grid->bb[bbo].upper[BB_Z]);
+                    atomToCluster(atomStart),
+                    bb_[bbo].lower[BB_X],
+                    bb_[bbo].lower[BB_Y],
+                    bb_[bbo].lower[BB_Z],
+                    bb_[bbo].upper[BB_X],
+                    bb_[bbo].upper[BB_Y],
+                    bb_[bbo].upper[BB_Z]);
         }
     }
 }
 
-/* Spatially sort the atoms within one grid column */
-static void sort_columns_simple(nbnxn_search *nbs,
-                                int dd_zone,
-                                nbnxn_grid_t *grid,
-                                int atomStart, int atomEnd,
-                                const int *atinfo,
-                                gmx::ArrayRef<const gmx::RVec> x,
-                                nbnxn_atomdata_t *nbat,
-                                int cxy_start, int cxy_end,
-                                gmx::ArrayRef<int> sort_work)
+void Grid::sortColumnsCpuGeometry(nbnxn_search *nbs,
+                                  int dd_zone,
+                                  int atomStart, int atomEnd,
+                                  const int *atinfo,
+                                  gmx::ArrayRef<const gmx::RVec> x,
+                                  nbnxn_atomdata_t *nbat,
+                                  int cxy_start, int cxy_end,
+                                  gmx::ArrayRef<int> sort_work)
 {
     if (debug)
     {
-        fprintf(debug, "cell0 %d sorting columns %d - %d, atoms %d - %d\n",
-                grid->cell0, cxy_start, cxy_end, atomStart, atomEnd);
+        fprintf(debug, "cell_offset %d sorting columns %d - %d, atoms %d - %d\n",
+                cellOffset_, cxy_start, cxy_end, atomStart, atomEnd);
     }
 
-    const bool relevantAtomsAreWithinGridBounds = (grid->maxAtomGroupRadius == 0);
+    const bool relevantAtomsAreWithinGridBounds = (dimensions_.maxAtomGroupRadius == 0);
+
+    const int  numAtomsPerCell = geometry_.numAtomsPerCell;
 
     /* Sort the atoms within each x,y column in 3 dimensions */
     for (int cxy = cxy_start; cxy < cxy_end; cxy++)
     {
-        int na  = grid->cxy_na[cxy];
-        int ncz = grid->cxy_ind[cxy+1] - grid->cxy_ind[cxy];
-        int ash = (grid->cell0 + grid->cxy_ind[cxy])*grid->na_sc;
+        const int numAtoms   = numAtomsInColumn(cxy);
+        const int numCellsZ  = cxy_ind_[cxy + 1] - cxy_ind_[cxy];
+        const int atomOffset = firstAtomInColumn(cxy);
 
         /* Sort the atoms within each x,y column on z coordinate */
         sort_atoms(ZZ, FALSE, dd_zone,
                    relevantAtomsAreWithinGridBounds,
-                   nbs->a.data() + ash, na, x,
-                   grid->c0[ZZ],
-                   1.0/grid->size[ZZ], ncz*grid->na_sc,
+                   nbs->a.data() + atomOffset, numAtoms, x,
+                   dimensions_.lowerCorner[ZZ],
+                   1.0/dimensions_.gridSize[ZZ], numCellsZ*numAtomsPerCell,
                    sort_work);
 
         /* Fill the ncz cells in this column */
-        int cfilled = grid->cxy_ind[cxy];
-        for (int cz = 0; cz < ncz; cz++)
+        const int firstCell  = firstCellInColumn(cxy);
+        int       cellFilled = firstCell;
+        for (int cellZ = 0; cellZ < numCellsZ; cellZ++)
         {
-            int c     = grid->cxy_ind[cxy] + cz;
+            const int cell           = firstCell + cellZ;
 
-            int ash_c = ash + cz*grid->na_sc;
-            int na_c  = std::min(grid->na_sc, na-(ash_c-ash));
+            const int atomOffsetCell = atomOffset + cellZ*numAtomsPerCell;
+            const int numAtomsCell   = std::min(numAtomsPerCell, numAtoms - (atomOffsetCell - atomOffset));
 
-            fill_cell(nbs, grid, nbat,
-                      ash_c, ash_c+na_c, atinfo, x,
-                      nullptr);
+            fillCell(nbs, nbat,
+                     atomOffsetCell, atomOffsetCell + numAtomsCell, atinfo, x,
+                     nullptr);
 
             /* This copy to bbcz is not really necessary.
              * But it allows to use the same grid search code
              * for the simple and supersub cell setups.
              */
-            if (na_c > 0)
+            if (numAtomsCell > 0)
             {
-                cfilled = c;
+                cellFilled = cell;
             }
-            grid->bbcz[c*NNBSBB_D  ] = grid->bb[cfilled].lower[BB_Z];
-            grid->bbcz[c*NNBSBB_D+1] = grid->bb[cfilled].upper[BB_Z];
+            bbcz_[cell*NNBSBB_D    ] = bb_[cellFilled].lower[BB_Z];
+            bbcz_[cell*NNBSBB_D + 1] = bb_[cellFilled].upper[BB_Z];
         }
 
         /* Set the unused atom indices to -1 */
-        for (int ind = na; ind < ncz*grid->na_sc; ind++)
+        for (int ind = numAtoms; ind < numCellsZ*numAtomsPerCell; ind++)
         {
-            nbs->a[ash+ind] = -1;
+            nbs->a[atomOffset + ind] = -1;
         }
     }
 }
 
 /* Spatially sort the atoms within one grid column */
-static void sort_columns_supersub(nbnxn_search *nbs,
+void Grid::sortColumnsGpuGeometry(nbnxn_search *nbs,
                                   int dd_zone,
-                                  nbnxn_grid_t *grid,
                                   int atomStart, int atomEnd,
                                   const int *atinfo,
                                   gmx::ArrayRef<const gmx::RVec> x,
@@ -1002,63 +1071,65 @@ static void sort_columns_supersub(nbnxn_search *nbs,
                                   int cxy_start, int cxy_end,
                                   gmx::ArrayRef<int> sort_work)
 {
-    nbnxn_bb_t  bb_work_array[2];
-    nbnxn_bb_t *bb_work_aligned = reinterpret_cast<nbnxn_bb_t *>((reinterpret_cast<std::size_t>(bb_work_array + 1)) & (~(static_cast<std::size_t>(15))));
+    BoundingBox  bb_work_array[2];
+    BoundingBox *bb_work_aligned = reinterpret_cast<BoundingBox *>((reinterpret_cast<std::size_t>(bb_work_array + 1)) & (~(static_cast<std::size_t>(15))));
 
     if (debug)
     {
-        fprintf(debug, "cell0 %d sorting columns %d - %d, atoms %d - %d\n",
-                grid->cell0, cxy_start, cxy_end, atomStart, atomEnd);
+        fprintf(debug, "cell_offset %d sorting columns %d - %d, atoms %d - %d\n",
+                cellOffset_, cxy_start, cxy_end, atomStart, atomEnd);
     }
 
-    const bool relevantAtomsAreWithinGridBounds = (grid->maxAtomGroupRadius == 0);
+    const bool relevantAtomsAreWithinGridBounds = (dimensions_.maxAtomGroupRadius == 0);
 
-    int        subdiv_x = grid->na_c;
-    int        subdiv_y = c_gpuNumClusterPerCellX*subdiv_x;
-    int        subdiv_z = c_gpuNumClusterPerCellY*subdiv_y;
+    const int  numAtomsPerCell = geometry_.numAtomsPerCell;
+
+    const int  subdiv_x = geometry_.numAtomsICluster;
+    const int  subdiv_y = c_gpuNumClusterPerCellX*subdiv_x;
+    const int  subdiv_z = c_gpuNumClusterPerCellY*subdiv_y;
 
     /* Sort the atoms within each x,y column in 3 dimensions.
      * Loop over all columns on the x/y grid.
      */
     for (int cxy = cxy_start; cxy < cxy_end; cxy++)
     {
-        int gridX            = cxy/grid->numCells[YY];
-        int gridY            = cxy - gridX*grid->numCells[YY];
+        const int gridX            = cxy/dimensions_.numCells[YY];
+        const int gridY            = cxy - gridX*dimensions_.numCells[YY];
 
-        int numAtomsInColumn = grid->cxy_na[cxy];
-        int numCellsInColumn = grid->cxy_ind[cxy + 1] - grid->cxy_ind[cxy];
-        int ash              = (grid->cell0 + grid->cxy_ind[cxy])*grid->na_sc;
+        const int numAtomsInColumn = cxy_na_[cxy];
+        const int numCellsInColumn = cxy_ind_[cxy + 1] - cxy_ind_[cxy];
+        const int atomOffset       = firstAtomInColumn(cxy);
 
         /* Sort the atoms within each x,y column on z coordinate */
         sort_atoms(ZZ, FALSE, dd_zone,
                    relevantAtomsAreWithinGridBounds,
-                   nbs->a.data() + ash, numAtomsInColumn, x,
-                   grid->c0[ZZ],
-                   1.0/grid->size[ZZ], numCellsInColumn*grid->na_sc,
+                   nbs->a.data() + atomOffset, numAtomsInColumn, x,
+                   dimensions_.lowerCorner[ZZ],
+                   1.0/dimensions_.gridSize[ZZ], numCellsInColumn*numAtomsPerCell,
                    sort_work);
 
         /* This loop goes over the cells and clusters along z at once */
         for (int sub_z = 0; sub_z < numCellsInColumn*c_gpuNumClusterPerCellZ; sub_z++)
         {
-            int ash_z = ash + sub_z*subdiv_z;
-            int na_z  = std::min(subdiv_z, numAtomsInColumn - (ash_z - ash));
-            int cz    = -1;
+            const int atomOffsetZ = atomOffset + sub_z*subdiv_z;
+            const int numAtomsZ   = std::min(subdiv_z, numAtomsInColumn - (atomOffsetZ - atomOffset));
+            int       cz          = -1;
             /* We have already sorted on z */
 
             if (sub_z % c_gpuNumClusterPerCellZ == 0)
             {
                 cz = sub_z/c_gpuNumClusterPerCellZ;
-                int cell = grid->cxy_ind[cxy] + cz;
+                const int cell = cxy_ind_[cxy] + cz;
 
                 /* The number of atoms in this cell/super-cluster */
-                int numAtoms = std::min(grid->na_sc, numAtomsInColumn - (ash_z - ash));
+                const int numAtoms = std::min(numAtomsPerCell, numAtomsInColumn - (atomOffsetZ - atomOffset));
 
-                grid->nsubc[cell] = std::min(c_gpuNumClusterPerCell,
-                                             (numAtoms + grid->na_c - 1)/grid->na_c);
+                numClusters_[cell] = std::min(c_gpuNumClusterPerCell,
+                                              (numAtoms + geometry_.numAtomsICluster - 1)/ geometry_.numAtomsICluster);
 
                 /* Store the z-boundaries of the bounding box of the cell */
-                grid->bbcz[cell*NNBSBB_D  ] = x[nbs->a[ash_z]][ZZ];
-                grid->bbcz[cell*NNBSBB_D+1] = x[nbs->a[ash_z + numAtoms - 1]][ZZ];
+                bbcz_[cell*NNBSBB_D  ] = x[nbs->a[atomOffsetZ]][ZZ];
+                bbcz_[cell*NNBSBB_D+1] = x[nbs->a[atomOffsetZ + numAtoms - 1]][ZZ];
             }
 
             if (c_gpuNumClusterPerCellY > 1)
@@ -1066,49 +1137,49 @@ static void sort_columns_supersub(nbnxn_search *nbs,
                 /* Sort the atoms along y */
                 sort_atoms(YY, (sub_z & 1) != 0, dd_zone,
                            relevantAtomsAreWithinGridBounds,
-                           nbs->a.data() + ash_z, na_z, x,
-                           grid->c0[YY] + gridY*grid->cellSize[YY],
-                           grid->invCellSize[YY], subdiv_z,
+                           nbs->a.data() + atomOffsetZ, numAtomsZ, x,
+                           dimensions_.lowerCorner[YY] + gridY*dimensions_.cellSize[YY],
+                           dimensions_.invCellSize[YY], subdiv_z,
                            sort_work);
             }
 
             for (int sub_y = 0; sub_y < c_gpuNumClusterPerCellY; sub_y++)
             {
-                int ash_y = ash_z + sub_y*subdiv_y;
-                int na_y  = std::min(subdiv_y, numAtomsInColumn - (ash_y - ash));
+                const int atomOffsetY = atomOffsetZ + sub_y*subdiv_y;
+                const int numAtomsY   = std::min(subdiv_y, numAtomsInColumn - (atomOffsetY - atomOffset));
 
                 if (c_gpuNumClusterPerCellX > 1)
                 {
                     /* Sort the atoms along x */
                     sort_atoms(XX, ((cz*c_gpuNumClusterPerCellY + sub_y) & 1) != 0, dd_zone,
                                relevantAtomsAreWithinGridBounds,
-                               nbs->a.data() + ash_y, na_y, x,
-                               grid->c0[XX] + gridX*grid->cellSize[XX],
-                               grid->invCellSize[XX], subdiv_y,
+                               nbs->a.data() + atomOffsetY, numAtomsY, x,
+                               dimensions_.lowerCorner[XX] + gridX*dimensions_.cellSize[XX],
+                               dimensions_.invCellSize[XX], subdiv_y,
                                sort_work);
                 }
 
                 for (int sub_x = 0; sub_x < c_gpuNumClusterPerCellX; sub_x++)
                 {
-                    int ash_x = ash_y + sub_x*subdiv_x;
-                    int na_x  = std::min(subdiv_x, numAtomsInColumn - (ash_x - ash));
+                    const int atomOffsetX = atomOffsetY + sub_x*subdiv_x;
+                    const int numAtomsX   = std::min(subdiv_x, numAtomsInColumn - (atomOffsetX - atomOffset));
 
-                    fill_cell(nbs, grid, nbat,
-                              ash_x, ash_x + na_x, atinfo, x,
-                              bb_work_aligned);
+                    fillCell(nbs, nbat,
+                             atomOffsetX, atomOffsetX + numAtomsX, atinfo, x,
+                             bb_work_aligned);
                 }
             }
         }
 
         /* Set the unused atom indices to -1 */
-        for (int ind = numAtomsInColumn; ind < numCellsInColumn*grid->na_sc; ind++)
+        for (int ind = numAtomsInColumn; ind < numCellsInColumn*numAtomsPerCell; ind++)
         {
-            nbs->a[ash + ind] = -1;
+            nbs->a[atomOffset + ind] = -1;
         }
     }
 }
 
-/* Sets the cell index in the cell array for atom \p atomIndex and increments the atom count for the grid column */
+/*! \brief Sets the cell index in the cell array for atom \p atomIndex and increments the atom count for the grid column */
 static void setCellAndAtomCount(gmx::ArrayRef<int>  cell,
                                 int                 cellIndex,
                                 gmx::ArrayRef<int>  cxy_na,
@@ -1118,8 +1189,8 @@ static void setCellAndAtomCount(gmx::ArrayRef<int>  cell,
     cxy_na[cellIndex] += 1;
 }
 
-/* Determine in which grid column atoms should go */
-static void calc_column_indices(nbnxn_grid_t *grid,
+/*! \brief Determine in which grid column atoms should go */
+static void calc_column_indices(const Grid::Dimensions &gridDims,
                                 const gmx::UpdateGroupsCog *updateGroupsCog,
                                 int atomStart, int atomEnd,
                                 gmx::ArrayRef<const gmx::RVec> x,
@@ -1128,14 +1199,17 @@ static void calc_column_indices(nbnxn_grid_t *grid,
                                 gmx::ArrayRef<int> cell,
                                 gmx::ArrayRef<int> cxy_na)
 {
+    const int numColumns = gridDims.numCells[XX]*gridDims.numCells[YY];
+
     /* We add one extra cell for particles which moved during DD */
-    for (int i = 0; i < grid->numCells[XX]*grid->numCells[YY] + 1; i++)
+    for (int i = 0; i < numColumns; i++)
     {
         cxy_na[i] = 0;
     }
 
     int taskAtomStart = atomStart + static_cast<int>((thread + 0)*(atomEnd - atomStart))/nthread;
     int taskAtomEnd   = atomStart + static_cast<int>((thread + 1)*(atomEnd - atomStart))/nthread;
+
     if (dd_zone == 0)
     {
         /* Home zone */
@@ -1151,35 +1225,39 @@ static void calc_column_indices(nbnxn_grid_t *grid,
                  * The int cast takes care of the lower bound,
                  * we will explicitly take care of the upper bound.
                  */
-                int cx = static_cast<int>((coord[XX] - grid->c0[XX])*grid->invCellSize[XX]);
-                int cy = static_cast<int>((coord[YY] - grid->c0[YY])*grid->invCellSize[YY]);
+                int cx = static_cast<int>((coord[XX] - gridDims.lowerCorner[XX])*gridDims.invCellSize[XX]);
+                int cy = static_cast<int>((coord[YY] - gridDims.lowerCorner[YY])*gridDims.invCellSize[YY]);
 
 #ifndef NDEBUG
-                if (cx < 0 || cx > grid->numCells[XX] ||
-                    cy < 0 || cy > grid->numCells[YY])
+                if (cx < 0 || cx > gridDims.numCells[XX] ||
+                    cy < 0 || cy > gridDims.numCells[YY])
                 {
                     gmx_fatal(FARGS,
                               "grid cell cx %d cy %d out of range (max %d %d)\n"
                               "atom %f %f %f, grid->c0 %f %f",
-                              cx, cy, grid->numCells[XX], grid->numCells[YY],
-                              x[i][XX], x[i][YY], x[i][ZZ], grid->c0[XX], grid->c0[YY]);
+                              cx, cy,
+                              gridDims.numCells[XX],
+                              gridDims.numCells[YY],
+                              x[i][XX], x[i][YY], x[i][ZZ],
+                              gridDims.lowerCorner[XX],
+                              gridDims.lowerCorner[YY]);
                 }
 #endif
                 /* Take care of potential rouding issues */
-                cx = std::min(cx, grid->numCells[XX] - 1);
-                cy = std::min(cy, grid->numCells[YY] - 1);
+                cx = std::min(cx, gridDims.numCells[XX] - 1);
+                cy = std::min(cy, gridDims.numCells[YY] - 1);
 
                 /* For the moment cell will contain only the, grid local,
                  * x and y indices, not z.
                  */
-                setCellAndAtomCount(cell, cx*grid->numCells[YY] + cy, cxy_na, i);
+                setCellAndAtomCount(cell, cx*gridDims.numCells[YY] + cy, cxy_na, i);
             }
             else
             {
                 /* Put this moved particle after the end of the grid,
                  * so we can process it later without using conditionals.
                  */
-                setCellAndAtomCount(cell, grid->numCells[XX]*grid->numCells[YY], cxy_na, i);
+                setCellAndAtomCount(cell, numColumns, cxy_na, i);
             }
         }
     }
@@ -1188,8 +1266,8 @@ static void calc_column_indices(nbnxn_grid_t *grid,
         /* Non-home zone */
         for (int i = taskAtomStart; i < taskAtomEnd; i++)
         {
-            int cx = static_cast<int>((x[i][XX] - grid->c0[XX])*grid->invCellSize[XX]);
-            int cy = static_cast<int>((x[i][YY] - grid->c0[YY])*grid->invCellSize[YY]);
+            int cx = static_cast<int>((x[i][XX] - gridDims.lowerCorner[XX])*gridDims.invCellSize[XX]);
+            int cy = static_cast<int>((x[i][YY] - gridDims.lowerCorner[YY])*gridDims.invCellSize[YY]);
 
             /* For non-home zones there could be particles outside
              * the non-bonded cut-off range, which have been communicated
@@ -1198,26 +1276,24 @@ static void calc_column_indices(nbnxn_grid_t *grid,
              * we put them in an extra row at the border.
              */
             cx = std::max(cx, 0);
-            cx = std::min(cx, grid->numCells[XX] - 1);
+            cx = std::min(cx, gridDims.numCells[XX] - 1);
             cy = std::max(cy, 0);
-            cy = std::min(cy, grid->numCells[YY] - 1);
+            cy = std::min(cy, gridDims.numCells[YY] - 1);
 
             /* For the moment cell will contain only the, grid local,
              * x and y indices, not z.
              */
-            setCellAndAtomCount(cell, cx*grid->numCells[YY] + cy, cxy_na, i);
+            setCellAndAtomCount(cell, cx*gridDims.numCells[YY] + cy, cxy_na, i);
         }
     }
 }
 
-/* Resizes grid and atom data which depend on the number of cells */
-static void resizeForNumberOfCells(const nbnxn_grid_t &grid,
-                                   int                 numAtomsMoved,
+/*! \brief Resizes grid and atom data which depend on the number of cells */
+static void resizeForNumberOfCells(const int           numNbnxnAtoms,
+                                   const int           numAtomsMoved,
                                    nbnxn_search       *nbs,
                                    nbnxn_atomdata_t   *nbat)
 {
-    int numNbnxnAtoms = (grid.cell0 + grid.nc)*grid.na_sc;
-
     /* Note: nbs->cell was already resized before */
 
     /* To avoid conditionals we store the moved particles at the end of a,
@@ -1229,20 +1305,20 @@ static void resizeForNumberOfCells(const nbnxn_grid_t &grid,
     nbat->resizeCoordinateBuffer(numNbnxnAtoms);
 }
 
-/* Determine in which grid cells the atoms should go */
-static void
-calc_cell_indices(nbnxn_search                   *nbs,
-                  int                             ddZone,
-                  nbnxn_grid_t                   *grid,
-                  const gmx::UpdateGroupsCog     *updateGroupsCog,
-                  int                             atomStart,
-                  int                             atomEnd,
-                  const int                      *atinfo,
-                  gmx::ArrayRef<const gmx::RVec>  x,
-                  int                             numAtomsMoved,
-                  const int                      *move,
-                  nbnxn_atomdata_t               *nbat)
+void Grid::calcCellIndices(nbnxn_search                   *nbs,
+                           int                             ddZone,
+                           int                             cellOffset,
+                           const gmx::UpdateGroupsCog     *updateGroupsCog,
+                           int                             atomStart,
+                           int                             atomEnd,
+                           const int                      *atinfo,
+                           gmx::ArrayRef<const gmx::RVec>  x,
+                           int                             numAtomsMoved,
+                           const int                      *move,
+                           nbnxn_atomdata_t               *nbat)
 {
+    cellOffset_ = cellOffset;
+
     /* First compute all grid/column indices and store them in nbs->cell */
     nbs->cell.resize(atomEnd);
 
@@ -1253,7 +1329,8 @@ calc_cell_indices(nbnxn_search                   *nbs,
     {
         try
         {
-            calc_column_indices(grid, updateGroupsCog,
+            calc_column_indices(dimensions_,
+                                updateGroupsCog,
                                 atomStart, atomEnd, x,
                                 ddZone, move, thread, nthread,
                                 nbs->cell, nbs->work[thread].cxy_na);
@@ -1261,11 +1338,13 @@ calc_cell_indices(nbnxn_search                   *nbs,
         GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR;
     }
 
+    const int numAtomsPerCell = geometry_.numAtomsPerCell;
+
     /* Make the cell index as a function of x and y */
     int ncz_max      = 0;
     int ncz          = 0;
-    grid->cxy_ind[0] = 0;
-    for (int i = 0; i < grid->numCells[XX]*grid->numCells[YY] + 1; i++)
+    cxy_ind_[0]      = 0;
+    for (int i = 0; i < numColumns() + 1; i++)
     {
         /* We set ncz_max at the beginning of the loop iso at the end
          * to skip i=grid->ncx*grid->numCells[YY] which are moved particles
@@ -1280,34 +1359,35 @@ calc_cell_indices(nbnxn_search                   *nbs,
         {
             cxy_na_i += nbs->work[thread].cxy_na[i];
         }
-        ncz = (cxy_na_i + grid->na_sc - 1)/grid->na_sc;
+        ncz = (cxy_na_i + numAtomsPerCell - 1)/numAtomsPerCell;
         if (nbat->XFormat == nbatX8)
         {
             /* Make the number of cell a multiple of 2 */
             ncz = (ncz + 1) & ~1;
         }
-        grid->cxy_ind[i+1] = grid->cxy_ind[i] + ncz;
-        /* Clear cxy_na, so we can reuse the array below */
-        grid->cxy_na[i] = 0;
+        cxy_ind_[i+1] = cxy_ind_[i] + ncz;
+        /* Clear cxy_na_, so we can reuse the array below */
+        cxy_na_[i] = 0;
     }
-    grid->nc = grid->cxy_ind[grid->numCells[XX]*grid->numCells[YY]] - grid->cxy_ind[0];
+    numCellsTotal_ = cxy_ind_[numColumns()] - cxy_ind_[0];
 
-    resizeForNumberOfCells(*grid, numAtomsMoved, nbs, nbat);
+    resizeForNumberOfCells(atomIndexEnd(), numAtomsMoved, nbs, nbat);
 
     if (debug)
     {
         fprintf(debug, "ns na_sc %d na_c %d super-cells: %d x %d y %d z %.1f maxz %d\n",
-                grid->na_sc, grid->na_c, grid->nc,
-                grid->numCells[XX], grid->numCells[YY], grid->nc/(static_cast<double>(grid->numCells[XX]*grid->numCells[YY])),
+                numAtomsPerCell, geometry_.numAtomsICluster, numCellsTotal_,
+                dimensions_.numCells[XX], dimensions_.numCells[YY],
+                numCellsTotal_/(static_cast<double>(numColumns())),
                 ncz_max);
         if (gmx_debug_at)
         {
             int i = 0;
-            for (int cy = 0; cy < grid->numCells[YY]; cy++)
+            for (int cy = 0; cy < dimensions_.numCells[YY]; cy++)
             {
-                for (int cx = 0; cx < grid->numCells[XX]; cx++)
+                for (int cx = 0; cx < dimensions_.numCells[XX]; cx++)
                 {
-                    fprintf(debug, " %2d", grid->cxy_ind[i+1]-grid->cxy_ind[i]);
+                    fprintf(debug, " %2d", cxy_ind_[i + 1] - cxy_ind_[i]);
                     i++;
                 }
                 fprintf(debug, "\n");
@@ -1316,12 +1396,13 @@ calc_cell_indices(nbnxn_search                   *nbs,
     }
 
     /* Make sure the work array for sorting is large enough */
-    if (ncz_max*grid->na_sc*SGSF > gmx::index(nbs->work[0].sortBuffer.size()))
+    const int worstCaseSortBufferSize = ncz_max*numAtomsPerCell*c_sortGridMaxSizeFactor;
+    if (worstCaseSortBufferSize > gmx::index(nbs->work[0].sortBuffer.size()))
     {
         for (nbnxn_search_work_t &work : nbs->work)
         {
             /* Elements not in use should be -1 */
-            work.sortBuffer.resize(ncz_max*grid->na_sc*SGSF, -1);
+            work.sortBuffer.resize(worstCaseSortBufferSize, -1);
         }
     }
 
@@ -1332,14 +1413,14 @@ calc_cell_indices(nbnxn_search                   *nbs,
     {
         /* At this point nbs->cell contains the local grid x,y indices */
         int cxy = nbs->cell[i];
-        nbs->a[(grid->cell0 + grid->cxy_ind[cxy])*grid->na_sc + grid->cxy_na[cxy]++] = i;
+        nbs->a[firstAtomInColumn(cxy) + cxy_na_[cxy]++] = i;
     }
 
     if (ddZone == 0)
     {
         /* Set the cell indices for the moved particles */
-        int n0 = grid->nc*grid->na_sc;
-        int n1 = grid->nc*grid->na_sc+grid->cxy_na[grid->numCells[XX]*grid->numCells[YY]];
+        int n0 = numCellsTotal_*numAtomsPerCell;
+        int n1 = numCellsTotal_*numAtomsPerCell + cxy_na_[numColumns()];
         if (ddZone == 0)
         {
             for (int i = n0; i < n1; i++)
@@ -1355,54 +1436,59 @@ calc_cell_indices(nbnxn_search                   *nbs,
     {
         try
         {
-            int columnStart = ((thread + 0)*grid->numCells[XX]*grid->numCells[YY])/nthread;
-            int columnEnd   = ((thread + 1)*grid->numCells[XX]*grid->numCells[YY])/nthread;
-            if (grid->bSimple)
+            int columnStart = ((thread + 0)*numColumns())/nthread;
+            int columnEnd   = ((thread + 1)*numColumns())/nthread;
+            if (geometry_.isSimple)
             {
-                sort_columns_simple(nbs, ddZone, grid, atomStart, atomEnd, atinfo, x, nbat,
-                                    columnStart, columnEnd,
-                                    nbs->work[thread].sortBuffer);
+                sortColumnsCpuGeometry(nbs, ddZone, atomStart, atomEnd, atinfo, x, nbat,
+                                       columnStart, columnEnd,
+                                       nbs->work[thread].sortBuffer);
             }
             else
             {
-                sort_columns_supersub(nbs, ddZone, grid, atomStart, atomEnd, atinfo, x, nbat,
-                                      columnStart, columnEnd,
-                                      nbs->work[thread].sortBuffer);
+                sortColumnsGpuGeometry(nbs, ddZone, atomStart, atomEnd, atinfo, x, nbat,
+                                       columnStart, columnEnd,
+                                       nbs->work[thread].sortBuffer);
             }
         }
         GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR;
     }
 
-    if (grid->bSimple && nbat->XFormat == nbatX8)
+    if (geometry_.isSimple && nbat->XFormat == nbatX8)
     {
-        combine_bounding_box_pairs(grid, grid->bb);
+        combine_bounding_box_pairs(*this, bb_, bbj_);
     }
 
-    if (!grid->bSimple)
+    if (!geometry_.isSimple)
     {
-        grid->nsubc_tot = 0;
-        for (int i = 0; i < grid->nc; i++)
+        numClustersTotal_ = 0;
+        for (int i = 0; i < numCellsTotal_; i++)
         {
-            grid->nsubc_tot += grid->nsubc[i];
+            numClustersTotal_ += numClusters_[i];
         }
     }
 
     if (debug)
     {
-        if (grid->bSimple)
+        if (geometry_.isSimple)
         {
-            print_bbsizes_simple(debug, grid);
+            print_bbsizes_simple(debug, *this);
         }
         else
         {
             fprintf(debug, "ns non-zero sub-cells: %d average atoms %.2f\n",
-                    grid->nsubc_tot, (atomEnd - atomStart)/static_cast<double>(grid->nsubc_tot));
+                    numClustersTotal_,
+                    (atomEnd - atomStart)/static_cast<double>(numClustersTotal_));
 
-            print_bbsizes_supersub(debug, grid);
+            print_bbsizes_supersub(debug, *this);
         }
     }
 }
 
+} // namespace Nbnxm
+
+// TODO: Move the functions below to nbnxn.cpp
+
 /* Sets up a grid and puts the atoms on the grid.
  * This function only operates on one domain of the domain decompostion.
  * Note that without domain decomposition there is only one domain.
@@ -1422,82 +1508,63 @@ void nbnxn_put_on_grid(nonbonded_verlet_t             *nbv,
                        const int                      *move)
 {
     nbnxn_search *nbs  = nbv->nbs.get();
-    nbnxn_grid_t *grid = &nbs->grid[ddZone];
+    Nbnxm::Grid  &grid = nbs->grid[ddZone];
 
     nbs_cycle_start(&nbs->cc[enbsCCgrid]);
 
-    grid->bSimple = nbv->pairlistIsSimple();
-
-    grid->na_c      = IClusterSizePerListType[nbv->pairlistSets().params().pairlistType];
-    grid->na_cj     = JClusterSizePerListType[nbv->pairlistSets().params().pairlistType];
-    grid->na_sc     = (grid->bSimple ? 1 : c_gpuNumClusterPerCell)*grid->na_c;
-    grid->na_c_2log = get_2log(grid->na_c);
-
+    int cellOffset;
     if (ddZone == 0)
     {
-        grid->cell0 = 0;
+        cellOffset = 0;
     }
     else
     {
-        grid->cell0 =
-            (nbs->grid[ddZone - 1].cell0 + nbs->grid[ddZone - 1].nc)*
-            nbs->grid[ddZone- 1].na_sc/grid->na_sc;
+        const Nbnxm::Grid &previousGrid = nbs->grid[ddZone - 1];
+        cellOffset = previousGrid.atomIndexEnd()/previousGrid.geometry().numAtomsPerCell;
     }
 
     const int n = atomEnd - atomStart;
 
+    real      maxAtomGroupRadius;
     if (ddZone == 0)
     {
         copy_mat(box, nbs->box);
 
-        /* Avoid zero density */
-        if (atomDensity > 0)
-        {
-            grid->atom_density = atomDensity;
-        }
-        else
-        {
-            grid->atom_density = grid_atom_density(n - numAtomsMoved, lowerCorner, upperCorner);
-        }
-
-        grid->cell0 = 0;
-
         nbs->natoms_local    = atomEnd - numAtomsMoved;
         /* We assume that nbnxn_put_on_grid is called first
          * for the local atoms (ddZone=0).
          */
         nbs->natoms_nonlocal = atomEnd - numAtomsMoved;
 
-        /* When not using atom groups, all atoms should be within the grid bounds */
-        grid->maxAtomGroupRadius = (updateGroupsCog ? updateGroupsCog->maxUpdateGroupRadius() : 0);
-        /* For the non-local grids the situation is the same as for the local */
-        for (size_t g = 1; g < nbs->grid.size(); g++)
-        {
-            nbs->grid[g].maxAtomGroupRadius = grid->maxAtomGroupRadius;
-        }
+        maxAtomGroupRadius = (updateGroupsCog ? updateGroupsCog->maxUpdateGroupRadius() : 0);
 
         if (debug)
         {
             fprintf(debug, "natoms_local = %5d atom_density = %5.1f\n",
-                    nbs->natoms_local, grid->atom_density);
+                    nbs->natoms_local, atomDensity);
         }
     }
     else
     {
+        const Nbnxm::Grid::Dimensions &dimsGrid0 = nbs->grid[0].dimensions();
+        atomDensity        = dimsGrid0.atomDensity;
+        maxAtomGroupRadius = dimsGrid0.maxAtomGroupRadius;
+
         nbs->natoms_nonlocal = std::max(nbs->natoms_nonlocal, atomEnd);
     }
 
     /* We always use the home zone (grid[0]) for setting the cell size,
      * since determining densities for non-local zones is difficult.
      */
-    set_grid_size_xy(nbs, grid,
-                     ddZone, n - numAtomsMoved,
-                     lowerCorner, upperCorner,
-                     nbs->grid[0].atom_density);
+    grid.setDimensions(nbs,
+                       ddZone, n - numAtomsMoved,
+                       lowerCorner, upperCorner,
+                       atomDensity,
+                       maxAtomGroupRadius);
 
     nbnxn_atomdata_t *nbat = nbv->nbat.get();
 
-    calc_cell_indices(nbs, ddZone, grid, updateGroupsCog, atomStart, atomEnd, atinfo, x, numAtomsMoved, move, nbat);
+    grid.calcCellIndices(nbs, ddZone, cellOffset, updateGroupsCog, atomStart, atomEnd, atinfo, x, numAtomsMoved, move, nbat);
 
     if (ddZone == 0)
     {
@@ -1541,16 +1608,16 @@ void nbnxn_put_on_grid_nonlocal(nonbonded_verlet_t              *nbv,
 
 void nbnxn_get_ncells(const nbnxn_search *nbs, int *ncx, int *ncy)
 {
-    *ncx = nbs->grid[0].numCells[XX];
-    *ncy = nbs->grid[0].numCells[YY];
+    *ncx = nbs->grid[0].dimensions().numCells[XX];
+    *ncy = nbs->grid[0].dimensions().numCells[YY];
 }
 
 gmx::ArrayRef<const int> nbnxn_get_atomorder(const nbnxn_search *nbs)
 {
     /* Return the atom order for the home cell (index 0) */
-    const nbnxn_grid_t &grid       = nbs->grid[0];
+    const Nbnxm::Grid &grid       = nbs->grid[0];
 
-    int                 numIndices = grid.cxy_ind[grid.numCells[XX]*grid.numCells[YY]]*grid.na_sc;
+    const int          numIndices = grid.atomIndexEnd() - grid.firstAtomInColumn(0);
 
     return gmx::constArrayRefFromArray(nbs->a.data(), numIndices);
 }
@@ -1558,22 +1625,19 @@ gmx::ArrayRef<const int> nbnxn_get_atomorder(const nbnxn_search *nbs)
 void nbnxn_set_atomorder(nbnxn_search *nbs)
 {
     /* Set the atom order for the home cell (index 0) */
-    nbnxn_grid_t *grid = &nbs->grid[0];
+    const Nbnxm::Grid &grid = nbs->grid[0];
 
-    int           ao = 0;
-    for (int cx = 0; cx < grid->numCells[XX]; cx++)
+    int                atomIndex = 0;
+    for (int cxy = 0; cxy < grid.numColumns(); cxy++)
     {
-        for (int cy = 0; cy < grid->numCells[YY]; cy++)
+        const int numAtoms  = grid.numAtomsInColumn(cxy);
+        int       cellIndex = grid.firstCellInColumn(cxy)*grid.geometry().numAtomsPerCell;
+        for (int i = 0; i < numAtoms; i++)
         {
-            int cxy = cx*grid->numCells[YY] + cy;
-            int j   = grid->cxy_ind[cxy]*grid->na_sc;
-            for (int cz = 0; cz < grid->cxy_na[cxy]; cz++)
-            {
-                nbs->a[j]     = ao;
-                nbs->cell[ao] = j;
-                ao++;
-                j++;
-            }
+            nbs->a[cellIndex]    = atomIndex;
+            nbs->cell[atomIndex] = cellIndex;
+            atomIndex++;
+            cellIndex++;
         }
     }
 }