From 0da0fa7e6b9b5b4ff34ddb6338f8b7e02a204f72 Mon Sep 17 00:00:00 2001 From: nstarman Date: Sat, 9 Dec 2023 21:43:34 -0500 Subject: [PATCH] Add converter for AbstractParameter unit Signed-off-by: nstarman # Conflicts: # src/galdynamix/potential/_potential/param/core.py --- src/galdynamix/potential/_potential/param/core.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/galdynamix/potential/_potential/param/core.py b/src/galdynamix/potential/_potential/param/core.py index 0cf2033f..ba5f3397 100644 --- a/src/galdynamix/potential/_potential/param/core.py +++ b/src/galdynamix/potential/_potential/param/core.py @@ -15,6 +15,7 @@ FloatArrayAnyShape, FloatOrIntScalar, FloatScalar, + Unit, ) from galdynamix.utils import partial_jit, vectorize_method from galdynamix.utils.dataclasses import converter_float_array @@ -34,7 +35,7 @@ class AbstractParameter(eqx.Module): # type: ignore[misc] """ _: KW_ONLY - unit: u.Unit = eqx.field(static=True) # TODO: move this to an annotation? + unit: Unit = eqx.field(static=True, converter=u.Unit) @abc.abstractmethod def __call__(self, t: FloatScalar, **kwargs: Any) -> ArrayAnyShape: