Skip to content

Faster grid.index() Implementation: No Automatic Wrapping #299

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
julianlitz merged 11 commits into from
Jun 22, 2026
Merged

Faster grid.index() Implementation: No Automatic Wrapping #299

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
619a59b
Faster .index() implementation. Safety wrapping no longer needed.
julianlitz Jun 20, 2026
dd9526a
Update extrapolated_prolongation.cpp
julianlitz Jun 20, 2026
643c1d6
Update fmg_interpolation.cpp
julianlitz Jun 20, 2026
a1b434a
Update prolongation.cpp
julianlitz Jun 20, 2026
082f2a7
Update fmg_interpolation.cpp
julianlitz Jun 20, 2026
dabf426
Update fmg_interpolation.cpp
julianlitz Jun 20, 2026
0486cc7
Update extrapolated_prolongation.cpp
julianlitz Jun 20, 2026
5331aa7
Update prolongation.cpp
julianlitz Jun 20, 2026
873b3e7
Update solve_tests.cpp
julianlitz Jun 21, 2026
e41012b
Test coverage
julianlitz Jun 21, 2026
decd328
Test coverage
julianlitz Jun 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 15 additions & 8 deletions include/DirectSolver/DirectSolverGive/applySymmetryShift.inl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ void DirectSolverGive<LevelCacheType>::applySymmetryShiftInnerBoundary(Vector<do

for (int i_theta = 0; i_theta < grid.ntheta(); i_theta++) {
const double theta = grid.theta(i_theta);

/* -------------------------- */
/* Node on the inner boundary */
/* -------------------------- */
Expand All @@ -37,12 +38,15 @@ void DirectSolverGive<LevelCacheType>::applySymmetryShiftInnerBoundary(Vector<do
k1 = grid.angularSpacing(i_theta - 1);
k2 = grid.angularSpacing(i_theta);

const int i_theta_M1 = grid.wrapThetaIndex(i_theta - 1);
const int i_theta_P1 = grid.wrapThetaIndex(i_theta + 1);

coeff2 = 0.5 * (k1 + k2) / h2;

/* Fill x(i+1,j) */
x(grid.index(i_r + 1, i_theta)) -= -coeff2 * arr * x(grid.index(i_r, i_theta)) /* Left */
+ 0.25 * art * x(grid.index(i_r, i_theta + 1)) /* Top Left */
- 0.25 * art * x(grid.index(i_r, i_theta - 1)); /* Bottom Left */
+ 0.25 * art * x(grid.index(i_r, i_theta_P1)) /* Top Left */
- 0.25 * art * x(grid.index(i_r, i_theta_M1)); /* Bottom Left */

/* --------------------------- */
/* Node next to inner boundary */
Expand All @@ -62,9 +66,9 @@ void DirectSolverGive<LevelCacheType>::applySymmetryShiftInnerBoundary(Vector<do
/* Fill x(i,j) */
x(grid.index(i_r, i_theta)) -= -coeff1 * arr * x(grid.index(i_r - 1, i_theta)); /* Left */
/* Fill x(i,j-1) */
x(grid.index(i_r, i_theta - 1)) -= +0.25 * art * x(grid.index(i_r - 1, i_theta)); /* Top Left */
x(grid.index(i_r, i_theta_M1)) -= +0.25 * art * x(grid.index(i_r - 1, i_theta)); /* Top Left */
/* Fill x(i,j+1) */
x(grid.index(i_r, i_theta + 1)) -= -0.25 * art * x(grid.index(i_r - 1, i_theta)); /* Bottom Left */
x(grid.index(i_r, i_theta_P1)) -= -0.25 * art * x(grid.index(i_r - 1, i_theta)); /* Bottom Left */
}
});
}
Expand Down Expand Up @@ -100,14 +104,17 @@ void DirectSolverGive<LevelCacheType>::applySymmetryShiftOuterBoundary(Vector<do
k1 = grid.angularSpacing(i_theta - 1);
k2 = grid.angularSpacing(i_theta);

const int i_theta_M1 = grid.wrapThetaIndex(i_theta - 1);
const int i_theta_P1 = grid.wrapThetaIndex(i_theta + 1);

coeff2 = 0.5 * (k1 + k2) / h2;

/* Fill result(i,j) */
x(grid.index(i_r, i_theta)) -= -coeff2 * arr * x(grid.index(i_r + 1, i_theta)); /* Right */
/* Fill result(i,j-1) */
x(grid.index(i_r, i_theta - 1)) -= -0.25 * art * x(grid.index(i_r + 1, i_theta)); /* Top Right */
x(grid.index(i_r, i_theta_M1)) -= -0.25 * art * x(grid.index(i_r + 1, i_theta)); /* Top Right */
/* Fill result(i,j+1) */
x(grid.index(i_r, i_theta + 1)) -= +0.25 * art * x(grid.index(i_r + 1, i_theta)); /* Bottom Right */
x(grid.index(i_r, i_theta_P1)) -= +0.25 * art * x(grid.index(i_r + 1, i_theta)); /* Bottom Right */

/* -------------------------- */
/* Node on the outer boundary */
Expand All @@ -126,8 +133,8 @@ void DirectSolverGive<LevelCacheType>::applySymmetryShiftOuterBoundary(Vector<do

/* Fill result(i-1,j) */
x(grid.index(i_r - 1, i_theta)) -= -coeff1 * arr * x(grid.index(i_r, i_theta)) /* Right */
- 0.25 * art * x(grid.index(i_r, i_theta + 1)) /* Top Right */
+ 0.25 * art * x(grid.index(i_r, i_theta - 1)); /* Bottom Right */
- 0.25 * art * x(grid.index(i_r, i_theta_P1)) /* Top Right */
+ 0.25 * art * x(grid.index(i_r, i_theta_M1)); /* Bottom Right */
}
});
}
Expand Down
6 changes: 4 additions & 2 deletions include/GMGPolar/utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -231,8 +231,9 @@ void GMGPolar<DomainGeometry, DensityProfileCoefficients>::writeToVTK(const std:
file << "<DataArray type=\"Int32\" Name=\"connectivity\" format=\"ascii\">\n";
for (int i_r = 0; i_r < grid.nr() - 1; i_r++) {
for (int i_theta = 0; i_theta < grid.ntheta(); i_theta++) {
const int i_theta_P1 = grid.wrapThetaIndex(i_theta + 1);
file << grid.index(i_r, i_theta) << " " << grid.index(i_r + 1, i_theta) << " "
<< grid.index(i_r + 1, i_theta + 1) << " " << grid.index(i_r, i_theta + 1) << "\n";
<< grid.index(i_r + 1, i_theta_P1) << " " << grid.index(i_r, i_theta_P1) << "\n";
}
}
file << "</DataArray>\n";
Expand Down Expand Up @@ -296,8 +297,9 @@ void GMGPolar<DomainGeometry, DensityProfileCoefficients>::writeToVTK(
file << "<DataArray type=\"Int32\" Name=\"connectivity\" format=\"ascii\">\n";
for (int i_r = 0; i_r < grid.nr() - 1; i_r++) {
for (int i_theta = 0; i_theta < grid.ntheta(); i_theta++) {
const int i_theta_P1 = grid.wrapThetaIndex(i_theta + 1);
file << grid.index(i_r, i_theta) << " " << grid.index(i_r + 1, i_theta) << " "
<< grid.index(i_r + 1, i_theta + 1) << " " << grid.index(i_r, i_theta + 1) << "\n";
<< grid.index(i_r + 1, i_theta_P1) << " " << grid.index(i_r, i_theta_P1) << "\n";
}
}

Expand Down
2 changes: 1 addition & 1 deletion include/PolarGrid/polargrid.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ class PolarGrid

// Optimized, inlined indexing.
KOKKOS_INLINE_FUNCTION int wrapThetaIndex(const int unwrapped_theta_index) const;
KOKKOS_INLINE_FUNCTION int index(const int r_index, const int unwrapped_theta_index) const;
KOKKOS_INLINE_FUNCTION int index(const int r_index, const int theta_index) const;
KOKKOS_INLINE_FUNCTION void multiIndex(const int node_index, int& r_index, int& theta_index) const;

// Grid Parameters
Expand Down
7 changes: 4 additions & 3 deletions include/PolarGrid/polargrid.inl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,15 +119,16 @@ KOKKOS_INLINE_FUNCTION int PolarGrid::wrapThetaIndex(const int unwrapped_theta_i
return theta_index;
}

KOKKOS_INLINE_FUNCTION int PolarGrid::index(const int r_index, const int unwrapped_theta_index) const
KOKKOS_INLINE_FUNCTION int PolarGrid::index(const int r_index, const int theta_index) const
{
// unwrapped_theta_index may be negative or larger than ntheta() to allow for periodicity.
assert(0 <= r_index && r_index < nr());
int theta_index = wrapThetaIndex(unwrapped_theta_index);
assert(0 <= theta_index && theta_index < ntheta());

int global_index =
r_index < numberSmootherCircles()
? theta_index + ntheta() * r_index
: numberCircularSmootherNodes() + r_index - numberSmootherCircles() + lengthRadialSmoother() * theta_index;

assert(0 <= global_index && global_index < numberOfNodes());
return global_index;
}
Expand Down
6 changes: 4 additions & 2 deletions src/Interpolation/extrapolated_prolongation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,11 +56,13 @@ static KOKKOS_INLINE_FUNCTION void fineNodeExtrapolatedProlongation(const int i_
const int i_r_coarse = i_r / 2;
const int i_theta_coarse = i_theta / 2;

const int i_theta_coarse_P1 = coarse_grid.wrapThetaIndex(i_theta_coarse + 1);

if (i_r & 1) {
if (i_theta & 1) { /* (odd, odd) -> node in center of coarse cell */
const double value =
0.5 * (coarse_values[coarse_grid.index(i_r_coarse + 1, i_theta_coarse)] + /* Bottom right */
coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse + 1)] /* Top left */
coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse_P1)] /* Top left */
);
fine_result[fine_grid.index(i_r, i_theta)] = value;
}
Expand All @@ -74,7 +76,7 @@ static KOKKOS_INLINE_FUNCTION void fineNodeExtrapolatedProlongation(const int i_
else {
if (i_theta & 1) { /* (even, odd) -> between coarse nodes in angular direction */
const double value = 0.5 * (coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse)] + /* Bottom */
coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse + 1)] /* Top */
coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse_P1)] /* Top */
);
fine_result[fine_grid.index(i_r, i_theta)] = value;
}
Expand Down
11 changes: 7 additions & 4 deletions src/Interpolation/extrapolated_restriction.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,20 +31,23 @@ static KOKKOS_INLINE_FUNCTION void coarseNodeExtrapolatedRestriction(const int i
const int i_r = i_r_coarse * 2;
const int i_theta = i_theta_coarse * 2;

const int i_theta_M1 = fine_grid.wrapThetaIndex(i_theta - 1);
const int i_theta_P1 = fine_grid.wrapThetaIndex(i_theta + 1);

/* Center + Angular contributions (always present) */
double value = fine_values[fine_grid.index(i_r, i_theta)] + 0.5 * fine_values[fine_grid.index(i_r, i_theta - 1)] +
0.5 * fine_values[fine_grid.index(i_r, i_theta + 1)];
double value = fine_values[fine_grid.index(i_r, i_theta)] + 0.5 * fine_values[fine_grid.index(i_r, i_theta_M1)] +
0.5 * fine_values[fine_grid.index(i_r, i_theta_P1)];

/* Left contributions (if not at inner boundary) */
if (i_r_coarse > 0) {
value += 0.5 * fine_values[fine_grid.index(i_r - 1, i_theta)] +
0.5 * fine_values[fine_grid.index(i_r - 1, i_theta + 1)]; /* Top-Left diagonal */
0.5 * fine_values[fine_grid.index(i_r - 1, i_theta_P1)]; /* Top-Left diagonal */
}

/* Right contributions (if not at outer boundary) */
if (i_r_coarse < coarse_grid.nr() - 1) {
value += 0.5 * fine_values[fine_grid.index(i_r + 1, i_theta)] +
0.5 * fine_values[fine_grid.index(i_r + 1, i_theta - 1)]; /* Bottom-Right diagonal */
0.5 * fine_values[fine_grid.index(i_r + 1, i_theta_M1)]; /* Bottom-Right diagonal */
}

coarse_result[coarse_grid.index(i_r_coarse, i_theta_coarse)] = value;
Expand Down
57 changes: 31 additions & 26 deletions src/Interpolation/fmg_interpolation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,10 @@ static KOKKOS_INLINE_FUNCTION void fineNodeFMGInterpolation(const int i_r, const
const int i_r_coarse = i_r / 2;
const int i_theta_coarse = i_theta / 2;

const int i_theta_coarse_M1 = coarse_grid.wrapThetaIndex(i_theta_coarse - 1);
const int i_theta_coarse_P1 = coarse_grid.wrapThetaIndex(i_theta_coarse + 1);
const int i_theta_coarse_P2 = coarse_grid.wrapThetaIndex(i_theta_coarse + 2);

/* Case 1: On the boundary */
if (i_r == 0 || i_r == fine_grid.nr() - 1) {
if (i_theta & 1) {
Expand All @@ -79,10 +83,10 @@ static KOKKOS_INLINE_FUNCTION void fineNodeFMGInterpolation(const int i_r, const
const double w_theta3 = -(k0 + k1) / (k0 + k1 + k2 + k3) * k1 / (k1 + k2 + k3) * k2 / k3;

fine_result[fine_grid.index(i_r, i_theta)] =
(w_theta0 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse - 1)] + /* (0, -3) */
(w_theta0 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse_M1)] + /* (0, -3) */
w_theta1 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse)] + /* (0, -1) */
w_theta2 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse + 1)] + /* (0, +1) */
w_theta3 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse + 2)] /* (0, +3) */
w_theta2 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse_P1)] + /* (0, +1) */
w_theta3 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse_P2)] /* (0, +3) */
);
}
else {
Expand All @@ -103,20 +107,21 @@ static KOKKOS_INLINE_FUNCTION void fineNodeFMGInterpolation(const int i_r, const
const double w_theta3 = -(k0 + k1) / (k0 + k1 + k2 + k3) * k1 / (k1 + k2 + k3) * k2 / k3;

const double left_value =
(w_theta0 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse - 1)] + /* (-1, -3) */
(w_theta0 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse_M1)] + /* (-1, -3) */
w_theta1 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse)] + /* (-1, -1) */
w_theta2 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse + 1)] + /* (-1, +1) */
w_theta3 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse + 2)] /* (-1, +3) */
w_theta2 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse_P1)] + /* (-1, +1) */
w_theta3 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse_P2)] /* (-1, +3) */
);
const double right_value =
(w_theta0 * coarse_values[coarse_grid.index(i_r_coarse + 1, i_theta_coarse - 1)] + /* (+1, -3) */
(w_theta0 * coarse_values[coarse_grid.index(i_r_coarse + 1, i_theta_coarse_M1)] + /* (+1, -3) */
w_theta1 * coarse_values[coarse_grid.index(i_r_coarse + 1, i_theta_coarse)] + /* (+1, -1) */
w_theta2 * coarse_values[coarse_grid.index(i_r_coarse + 1, i_theta_coarse + 1)] + /* (+1, +1) */
w_theta3 * coarse_values[coarse_grid.index(i_r_coarse + 1, i_theta_coarse + 2)] /* (+1, +3) */
w_theta2 * coarse_values[coarse_grid.index(i_r_coarse + 1, i_theta_coarse_P1)] + /* (+1, +1) */
w_theta3 * coarse_values[coarse_grid.index(i_r_coarse + 1, i_theta_coarse_P2)] /* (+1, +3) */
);

const double h1 = fine_grid.radialSpacing(i_r - 1);
const double h2 = fine_grid.radialSpacing(i_r);
const double h1 = fine_grid.radialSpacing(i_r - 1);
const double h2 = fine_grid.radialSpacing(i_r);

fine_result[fine_grid.index(i_r, i_theta)] = (h1 * left_value + h2 * right_value) / (h1 + h2);
}
else {
Expand All @@ -143,28 +148,28 @@ static KOKKOS_INLINE_FUNCTION void fineNodeFMGInterpolation(const int i_r, const
const double w_theta3 = -(k0 + k1) / (k0 + k1 + k2 + k3) * k1 / (k1 + k2 + k3) * k2 / k3;

const double outer_left_value =
(w_theta0 * coarse_values[coarse_grid.index(i_r_coarse - 1, i_theta_coarse - 1)] + /* (-3, -3) */
(w_theta0 * coarse_values[coarse_grid.index(i_r_coarse - 1, i_theta_coarse_M1)] + /* (-3, -3) */
w_theta1 * coarse_values[coarse_grid.index(i_r_coarse - 1, i_theta_coarse)] + /* (-3, -1) */
w_theta2 * coarse_values[coarse_grid.index(i_r_coarse - 1, i_theta_coarse + 1)] + /* (-3, +1) */
w_theta3 * coarse_values[coarse_grid.index(i_r_coarse - 1, i_theta_coarse + 2)] /* (-3, +3) */
w_theta2 * coarse_values[coarse_grid.index(i_r_coarse - 1, i_theta_coarse_P1)] + /* (-3, +1) */
w_theta3 * coarse_values[coarse_grid.index(i_r_coarse - 1, i_theta_coarse_P2)] /* (-3, +3) */
);
const double inner_left_value =
(w_theta0 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse - 1)] + /* (-1, -3) */
(w_theta0 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse_M1)] + /* (-1, -3) */
w_theta1 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse)] + /* (-1, -1) */
w_theta2 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse + 1)] + /* (-1, +1) */
w_theta3 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse + 2)] /* (-1, +3) */
w_theta2 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse_P1)] + /* (-1, +1) */
w_theta3 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse_P2)] /* (-1, +3) */
);
const double inner_right_value =
(w_theta0 * coarse_values[coarse_grid.index(i_r_coarse + 1, i_theta_coarse - 1)] + /* (+1, -3) */
(w_theta0 * coarse_values[coarse_grid.index(i_r_coarse + 1, i_theta_coarse_M1)] + /* (+1, -3) */
w_theta1 * coarse_values[coarse_grid.index(i_r_coarse + 1, i_theta_coarse)] + /* (+1, -1) */
w_theta2 * coarse_values[coarse_grid.index(i_r_coarse + 1, i_theta_coarse + 1)] + /* (+1, +1) */
w_theta3 * coarse_values[coarse_grid.index(i_r_coarse + 1, i_theta_coarse + 2)] /* (+1, +3) */
w_theta2 * coarse_values[coarse_grid.index(i_r_coarse + 1, i_theta_coarse_P1)] + /* (+1, +1) */
w_theta3 * coarse_values[coarse_grid.index(i_r_coarse + 1, i_theta_coarse_P2)] /* (+1, +3) */
);
const double outer_right_value =
(w_theta0 * coarse_values[coarse_grid.index(i_r_coarse + 2, i_theta_coarse - 1)] + /* (+3, -3) */
(w_theta0 * coarse_values[coarse_grid.index(i_r_coarse + 2, i_theta_coarse_M1)] + /* (+3, -3) */
w_theta1 * coarse_values[coarse_grid.index(i_r_coarse + 2, i_theta_coarse)] + /* (+3, -1) */
w_theta2 * coarse_values[coarse_grid.index(i_r_coarse + 2, i_theta_coarse + 1)] + /* (+3, +1) */
w_theta3 * coarse_values[coarse_grid.index(i_r_coarse + 2, i_theta_coarse + 2)] /* (+3, +3) */
w_theta2 * coarse_values[coarse_grid.index(i_r_coarse + 2, i_theta_coarse_P1)] + /* (+3, +1) */
w_theta3 * coarse_values[coarse_grid.index(i_r_coarse + 2, i_theta_coarse_P2)] /* (+3, +3) */
);

const double h0 = coarse_grid.radialSpacing(i_r_coarse - 1);
Expand Down Expand Up @@ -212,10 +217,10 @@ static KOKKOS_INLINE_FUNCTION void fineNodeFMGInterpolation(const int i_r, const
const double w_theta3 = -(k0 + k1) / (k0 + k1 + k2 + k3) * k1 / (k1 + k2 + k3) * k2 / k3;

fine_result[fine_grid.index(i_r, i_theta)] =
(w_theta0 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse - 1)] + /* (0, -3) */
(w_theta0 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse_M1)] + /* (0, -3) */
w_theta1 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse)] + /* (0, -1) */
w_theta2 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse + 1)] + /* (0, +1) */
w_theta3 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse + 2)] /* (0, +3) */
w_theta2 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse_P1)] + /* (0, +1) */
w_theta3 * coarse_values[coarse_grid.index(i_r_coarse, i_theta_coarse_P2)] /* (0, +3) */
);
}
else {
Expand Down
Loading
Loading