Merge pull request #39 from ATTPC/gwm17_dev

Fix small bug in simulator.py, add type hint for Numba dicts

Author: sigmanotation

docs/api/detector/pairing.md

@@ -0,0 +1,3 @@
+# pairing Module
+
+::: attpc_engine.detector.pairing

docs/api/detector/typed_dict.md

@@ -0,0 +1,3 @@
+# typed_dict Module
+
+::: attpc_engine.detector.typed_dict

mkdocs.yml

@@ -19,11 +19,13 @@ nav:
       - About detector: api/detector/index.md
       - beam_pads: api/detector/beam_pads.md
       - constants: api/detector/constants.md
+      - pairing: api/detector/pairing.md
       - parameters: api/detector/parameters.md
       - response: api/detector/response.md
       - simulator: api/detector/simulator.md
       - solver: api/detector/solver.md
       - transporter: api/detector/transporter.md
+      - typed_dict: api/detector/typed_dict.md
       - writer: api/detector/writer.md
     - kinematics:
       - About kinematics: api/kinematics/index.md

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "attpc_engine"
-version = "0.2.0"
+version = "0.3.0"
 description = "AT-TPC Monte-Carlo simulation engine"
 authors = [
     {name = "Gordon McCann", email = "gordonmccann215@gmail.com"},

src/attpc_engine/detector/simulator.py

@@ -11,6 +11,7 @@
 from .constants import NUM_TB
 from .. import nuclear_map
 from .pairing import unpair
+from .typed_dict import NumbaTypedDict
 
 import numpy as np
 import h5py as h5
@@ -123,7 +124,7 @@ def digitize(self, config: Config, rng: Generator) -> np.ndarray:
 
 
 @njit
-def dict_to_points(points: Dict[int, int]) -> np.ndarray:
+def dict_to_points(points: NumbaTypedDict[int, int]) -> np.ndarray:
     """
     Converts dictionary of N pad,tb keys with corresponding number of electrons
     to Nx3 array where each row is [pad, tb, e], now combined over pad/tb combos.
@@ -294,7 +295,7 @@ def generate_electrons(
         return electrons
 
     def generate_point_cloud(
-        self, config: Config, rng: Generator, points: Dict[int, int]
+        self, config: Config, rng: Generator, points: NumbaTypedDict
     ):
         """Create the point cloud
 
@@ -331,7 +332,7 @@ def generate_point_cloud(
         if config.pad_grid_edges is None or config.pad_grid is None:
             raise ValueError("Pad grid is not loaded at SimParticle.generate_hits!")
 
-        points = transport_track(
+        transport_track(
             config.pad_grid,
             config.pad_grid_edges,
             config.det_params.diffusion,

src/attpc_engine/detector/transporter.py

@@ -1,9 +1,9 @@
 from .pairing import pair
 from .beam_pads import BEAM_PADS_ARRAY
+from .typed_dict import NumbaTypedDict
 
 import numpy as np
 from numba import njit
-from numba.typed import Dict
 
 STEPS = 10
 
@@ -124,7 +124,7 @@ def point_transport(
     time: float,
     center: tuple[float, float],
     electrons: int,
-    points: Dict[int, int],
+    points: NumbaTypedDict[int, int],
 ):
     """
     Transports all electrons created at a point in a simulated nucleus' track
@@ -170,7 +170,7 @@ def transverse_transport(
     center: tuple[float, float],
     electrons: int,
     sigma_t: float,
-    points: Dict[int, int],
+    points: NumbaTypedDict[int, int],
 ):
     """
     Transports all electrons created at a point in a simulated nucleus'
@@ -246,7 +246,7 @@ def find_pads_hit(
     center: tuple[float, float],
     electrons: int,
     sigma_t: float,
-    points: Dict[int, int],
+    points: NumbaTypedDict[int, int],
 ):
     """
     Finds the pads hit by transporting the electrons created at a point in
@@ -298,7 +298,7 @@ def transport_track(
     dv: float,
     track: np.ndarray,
     electrons: np.ndarray,
-    points: Dict[int, int],
+    points: NumbaTypedDict[int, int],
 ):
     """
     High-level function that transports each point in a nucleus' trajectory

src/attpc_engine/detector/typed_dict.py

@@ -0,0 +1,36 @@
+from typing import TypeVar, Generic, Iterable
+
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+class NumbaTypedDict(Generic[K, V]):
+    """This is simply a type hint interface for Numba typed dictionaries
+
+    Use this as a type hint wherever you would be using numba.typed.Dict
+    and magically all of your linters will be happy.
+
+    Example:
+    ```py
+    my_dict = numba.typed.Dict.empty(
+        key_type=numba.core.types.int64,
+        value_type=numba.core.types.int64
+    )
+    ```
+    is type-hinted as
+    ```py
+    my_dict: NumbaTypedDict[int, int]
+    ```
+
+    Do not attempt to instantiate this object! It won't work!
+    """
+
+    def __getitem__(self, x: K) -> V: ...
+
+    def __setitem__(self, x: K, v: V): ...
+
+    def __len__(self) -> int: ...
+
+    def items(self) -> Iterable[tuple[K, V]]: ...
+
+    def get(self, x: K, default: V | None = None) -> V: ...