WIP things to make Microscope Calibration work (#4)

* Default with correct type

* Add more tests

...in particular for `run_iter()` and more components.

Test descan error implementation for correctly handling the names, not only indices.

Drive-by flake8 fixes

* Introduce type for single coordinates or pixels; allow float pixels

The types work well to make clear what x and y are, and to distinguish
pixel vs physical coordinates in code.

* Introduce types for single values as opposed to arrays, which correspond well to single rays etc
* Make Grid center a single type
* Allow floats for pixel coordinates since conversion to discrete values
  should happen as late as possible to avoid rounding issues
* Test for conversion helpers that reduce boilerplate

* Set float64 support everywhere

---------

Co-authored-by: Matthew Bryan <78845903+matbryan52@users.noreply.github.com>

Author: uellue

conftest.py

@@ -1,3 +1,4 @@
 import jax
 
+jax.config.update("jax_enable_x64", True)  # noqa: E702
 jax.config.update("jax_platform_name", "cpu")

src/temgym_core/__init__.py

@@ -1,5 +1,8 @@
 from typing_extensions import TypeAlias
-from typing import NamedTuple
+from typing import NamedTuple, Union
+
+import jax; jax.config.update("jax_enable_x64", True)  # noqa: E702
+import jax.numpy as jnp
 import numpy as np
 from numpy.typing import NDArray
 
@@ -46,6 +49,26 @@ class ScaleYX(NamedTuple):
     x: float
 
 
+class CoordXY(NamedTuple):
+    """Continuous coordinates in the optical frame.
+
+    Parameters
+    ----------
+    x : float
+        X position, metres.
+    y : float
+        Y position, metres.
+    """
+    x: float
+    y: float
+
+    def to_coords(self) -> 'CoordsXY':
+        return CoordsXY(
+            x=jnp.array((self.x,)),
+            y=jnp.array((self.y,))
+        )
+
+
 class CoordsXY(NamedTuple):
     """Continuous coordinates in the optical frame.
 
@@ -60,22 +83,46 @@ class CoordsXY(NamedTuple):
     y: NDArray[np.floating]
 
 
+class PixelYX(NamedTuple):
+    """Pixel coordinates for images.
+
+    Parameters
+    ----------
+    y : Union[int, float]
+        Pixel row indices
+    x : Union[int, float]
+        Pixel column indices
+
+    Notes
+    -----
+    Pixel indices are 0-based.
+    """
+    y: Union[int, float]
+    x: Union[int, float]
+
+    def to_pixels(self) -> 'PixelsYX':
+        return PixelsYX(
+            x=jnp.array((self.x,)),
+            y=jnp.array((self.y,))
+        )
+
+
 class PixelsYX(NamedTuple):
-    """Discrete pixel coordinates for images.
+    """Pixel coordinates for images.
 
     Parameters
     ----------
     y : numpy.ndarray
-        Pixel row indices. Integer dtype.
+        Pixel row indices. Integer or floating dtype.
     x : numpy.ndarray
-        Pixel column indices. Integer dtype.
+        Pixel column indices. Integer or floating dtype.
 
     Notes
     -----
     Pixel indices are 0-based.
     """
-    y: NDArray[np.integer]
-    x: NDArray[np.integer]
+    y: NDArray[Union[np.integer, np.floating]]
+    x: NDArray[Union[np.integer, np.floating]]
 
 
 # Convenience re-exports

src/temgym_core/components.py

@@ -1,4 +1,5 @@
 from typing import NamedTuple, Dict
+import jax; jax.config.update("jax_enable_x64", True)  # noqa: E702
 import jax_dataclasses as jdc
 import jax.numpy as jnp
 

src/temgym_core/coordinate_transforms.py

@@ -1,3 +1,4 @@
+import jax; jax.config.update("jax_enable_x64", True)  # noqa: E702
 import jax.numpy as jnp
 import jax.lax as lax
 from . import Degrees, Radians, ShapeYX, CoordsXY, ScaleYX

src/temgym_core/gaussian.py

@@ -1,12 +1,13 @@
+import jax; jax.config.update("jax_enable_x64", True)  # noqa: E702
 import jax.numpy as jnp
-import jax
+import jax_dataclasses as jdc
+from jax._src.lax.control_flow.loops import _batch_and_remainder
+from jax import lax
+
 from .grid import Grid
 from .run import run_to_end
 from .utils import custom_jacobian_matrix
 from .ray import Ray
-import jax_dataclasses as jdc
-from jax._src.lax.control_flow.loops import _batch_and_remainder
-from jax import lax
 
 
 def w_z(w0, z, z_r):

src/temgym_core/grid.py

@@ -1,8 +1,9 @@
 from typing import Union
 import numpy as np
+import jax; jax.config.update("jax_enable_x64", True)  # noqa: E702
 import jax.numpy as jnp
 
-from . import Degrees, ShapeYX, CoordsXY, ScaleYX, PixelsYX
+from . import Degrees, ShapeYX, CoordXY, CoordsXY, ScaleYX, PixelsYX
 from .ray import Ray
 from .utils import inplace_sum, try_ravel, try_reshape
 from .coordinate_transforms import pixels_to_metres_transform, apply_transformation
@@ -15,7 +16,7 @@ class Grid:
     ----------
     z : float
         Axial position in metres.
-    centre : CoordsXY
+    centre : CoordXY
         Grid centre in metres (x, y).
     shape : ShapeYX
         Grid shape (y, x) in pixels.
@@ -27,7 +28,7 @@ class Grid:
         If True, apply an additional vertical flip.
     """
     z: float
-    centre: CoordsXY
+    centre: CoordXY
     shape: ShapeYX
     pixel_size: ScaleYX
     rotation: Degrees
@@ -146,7 +147,10 @@ def metres_to_pixels(self, coords: CoordsXY, cast: bool = True) -> PixelsYX:
         if cast:
             pixels_y = jnp.round(pixels_y).astype(jnp.int32)
             pixels_x = jnp.round(pixels_x).astype(jnp.int32)
-        return try_reshape(pixels_y, coords_y), try_reshape(pixels_x, coords_x)
+        return PixelsYX(
+            y=try_reshape(pixels_y, coords_y),
+            x=try_reshape(pixels_x, coords_x)
+        )
 
     def pixels_to_metres(self, pixels: PixelsYX) -> CoordsXY:
         """Convert pixel indices to metric coordinates.
@@ -172,7 +176,10 @@ def pixels_to_metres(self, pixels: PixelsYX) -> CoordsXY:
         metres_y, metres_x = apply_transformation(
             try_ravel(pixels_y), try_ravel(pixels_x), pixels_to_metres_mat
         )
-        return try_reshape(metres_x, pixels_x), try_reshape(metres_y, pixels_y)
+        return CoordsXY(
+            x=try_reshape(metres_x, pixels_x),
+            y=try_reshape(metres_y, pixels_y)
+        )
 
     def ray_at_grid(
         self, px_y: float, px_x: float, dx: float = 0., dy: float = 0., z: float | None = None

src/temgym_core/ray.py

@@ -1,4 +1,5 @@
 import dataclasses
+import jax; jax.config.update("jax_enable_x64", True)  # noqa: E702
 import jax_dataclasses as jdc
 import jax.numpy as jnp
 from .tree_utils import HasParamsMixin

src/temgym_core/run.py

@@ -2,7 +2,7 @@
 import dataclasses
 from typing import TYPE_CHECKING, Sequence, Union, Any, Callable, Generator
 
-import jax
+import jax; jax.config.update("jax_enable_x64", True)  # noqa: E702
 import jax.numpy as jnp
 from .utils import custom_jacobian_matrix
 from .propagator import FreeSpaceParaxial, BasePropagator, Propagator

src/temgym_core/source.py

@@ -1,4 +1,5 @@
 import numpy as np
+import jax; jax.config.update("jax_enable_x64", True)  # noqa: E702
 import jax_dataclasses as jdc
 
 from .tree_utils import HasParamsMixin
@@ -98,7 +99,7 @@ class PointSource(Source):
     """
     z: float
     semi_conv: float
-    offset_xy: CoordsXY = (0.0, 0.0)
+    offset_xy: CoordsXY = CoordsXY(x=0.0, y=0.0)
 
     def generate_array(self, num: int, random: bool = False) -> np.ndarray:
         """Generate rays with varying slopes within a cone of semi-convergence.

src/temgym_core/transfer.py

@@ -1,3 +1,4 @@
+import jax; jax.config.update("jax_enable_x64", True)  # noqa: E702
 import jax.numpy as jnp
 import numpy as np