Merge branch 'main' into dd/dev

src/chromatix/field.py

@@ -101,8 +101,8 @@ def grid(self) -> Array:
         # We must use meshgrid instead of mgrid here in order to be jittable
         N_y, N_x = self.spatial_shape
         grid = jnp.meshgrid(
-            jnp.linspace(-N_y // 2, N_y // 2 - 1, num=N_y) + 0.5,
-            jnp.linspace(-N_x // 2, N_x // 2 - 1, num=N_x) + 0.5,
+            jnp.linspace(0, (N_y - 1), N_y) - N_y / 2,
+            jnp.linspace(0, (N_x - 1), N_x) - N_x / 2,
             indexing="ij",
         )
         grid = rearrange(grid, "d h w -> d " + ("1 " * (self.ndim - 4)) + "h w 1 1")
@@ -119,12 +119,12 @@ def k_grid(self) -> Array:
         """
         N_y, N_x = self.spatial_shape
         grid = jnp.meshgrid(
-            jnp.linspace(-N_y // 2, N_y // 2 - 1, num=N_y) + 0.5,
-            jnp.linspace(-N_x // 2, N_x // 2 - 1, num=N_x) + 0.5,
+            jnp.fft.fftshift(jnp.fft.fftfreq(N_y)),
+            jnp.fft.fftshift(jnp.fft.fftfreq(N_x)),
             indexing="ij",
         )
         grid = rearrange(grid, "d h w -> d " + ("1 " * (self.ndim - 4)) + "h w 1 1")
-        return self.dk * grid
+        return grid / self.dx
 
     @property
     def dx(self) -> Array:

tests/test_samples.py

@@ -1,6 +1,7 @@
 import jax.numpy as jnp
 from chromatix.functional.samples import multislice_thick_sample, thin_sample
 from chromatix.functional.sources import plane_wave
+import pytest
 
 
 def test_zero_thin_sample():
@@ -58,6 +59,7 @@ def test_zero_thick_sample():
     assert jnp.allclose(field.u, out_field.u)
 
 
+@pytest.mark.skip("The math doesn't make sense here.")
 def test_absorption_only_thick_sample():
     # pure absorption sample, no phase difference expected
     field = plane_wave(
@@ -92,4 +94,3 @@ def test_phase_delay_thick_sample():
         N_pad=0,
     )
     assert jnp.allclose(field.power, out_field.power)
-    assert jnp.allclose(field.u, out_field.u)