Skip to content

Commit

Permalink
Merge pull request #14 from wheelerMT/fix-indexing
Browse files Browse the repository at this point in the history
Fix indexing of meshgrids to use 'ij' instead of 'xy'
  • Loading branch information
wheelerMT authored Mar 23, 2024
2 parents b675a9e + 5bb25ef commit 79cf936
Show file tree
Hide file tree
Showing 3 changed files with 10 additions and 7 deletions.
10 changes: 6 additions & 4 deletions pygpe/shared/grid.py
Original file line number Diff line number Diff line change
Expand Up @@ -127,7 +127,7 @@ def _generate_2d_grids(
cp.arange(-self.num_points_y // 2, self.num_points_y // 2)
* self.grid_spacing_y
)
self.x_mesh, self.y_mesh = cp.meshgrid(x, y)
self.x_mesh, self.y_mesh = cp.meshgrid(x, y, indexing="ij")

# Generate Fourier space variables
self.fourier_spacing_x = cp.pi / (self.num_points_x // 2 * self.grid_spacing_x)
Expand All @@ -142,7 +142,9 @@ def _generate_2d_grids(
* self.fourier_spacing_y
)

self.fourier_x_mesh, self.fourier_y_mesh = cp.meshgrid(fourier_x, fourier_y)
self.fourier_x_mesh, self.fourier_y_mesh = cp.meshgrid(
fourier_x, fourier_y, indexing="ij"
)
self.fourier_x_mesh = cp.fft.fftshift(self.fourier_x_mesh)
self.fourier_y_mesh = cp.fft.fftshift(self.fourier_y_mesh)

Expand Down Expand Up @@ -181,7 +183,7 @@ def _generate_3d_grids(
cp.arange(-self.num_points_z // 2, self.num_points_z // 2)
* self.grid_spacing_z
)
self.x_mesh, self.y_mesh, self.z_mesh = cp.meshgrid(x, y, z)
self.x_mesh, self.y_mesh, self.z_mesh = cp.meshgrid(x, y, z, indexing="ij")

# Generate Fourier space variables
self.fourier_spacing_x = cp.pi / (self.num_points_x // 2 * self.grid_spacing_x)
Expand All @@ -205,7 +207,7 @@ def _generate_3d_grids(
self.fourier_x_mesh,
self.fourier_y_mesh,
self.fourier_z_mesh,
) = cp.meshgrid(fourier_x, fourier_y, fourier_z)
) = cp.meshgrid(fourier_x, fourier_y, fourier_z, indexing="ij")
self.fourier_x_mesh = cp.fft.fftshift(self.fourier_x_mesh)
self.fourier_y_mesh = cp.fft.fftshift(self.fourier_y_mesh)
self.fourier_z_mesh = cp.fft.fftshift(self.fourier_z_mesh)
Expand Down
3 changes: 2 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[tool.poetry]
name = "pygpe"
version = "2.0.0"
version = "2.0.1"
description = "A fast Gross-Pitaevskii equation solver for scalar, spin-1 and spin-2 BEC systems."
authors = ["Matt Wheeler <[email protected]>"]
license = "MIT"
Expand All @@ -20,6 +20,7 @@ gpu = ["cupy"]

[tool.poetry.group.dev.dependencies]
pytest = "^7.4.4"
black = "^24.3.0"

[build-system]
requires = ["poetry-core"]
Expand Down
4 changes: 2 additions & 2 deletions tests/shared/test_grids.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,9 +59,9 @@ def test_correct_fourier_shift_2d():

grid2d = grid.Grid((64, 64), (0.5, 0.5))

for element in grid2d.fourier_x_mesh[:, 0]:
for element in grid2d.fourier_x_mesh[0, :]:
assert element == 0
for element in grid2d.fourier_y_mesh[0, :]:
for element in grid2d.fourier_y_mesh[:, 0]:
assert element == 0


Expand Down

0 comments on commit 79cf936

Please sign in to comment.