From dfe096e9c39e6d99251cb7e326c6869df09de4a3 Mon Sep 17 00:00:00 2001
From: Ziwen Liu <67518483+ziw-liu@users.noreply.github.com>
Date: Fri, 6 Sep 2024 17:00:27 -0700
Subject: [PATCH] Use annotation instead of field for tagged union also removed
 non-standard axis

---
 iohub/ngff/models.py | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/iohub/ngff/models.py b/iohub/ngff/models.py
index d94fb5dc..00ed964c 100644
--- a/iohub/ngff/models.py
+++ b/iohub/ngff/models.py
@@ -17,6 +17,7 @@
     AfterValidator,
     BaseModel,
     ConfigDict,
+    Discriminator,
     Field,
     NonNegativeInt,
     PositiveInt,
@@ -168,13 +169,11 @@ class TimeAxisMeta(NamedAxisMeta):
     )
 
 
-class NonstandardAxisMeta(NamedAxisMeta):
-    type: str | None
-    unit: str | None
-
-
 """https://ngff.openmicroscopy.org/0.4/index.html#axes-md"""
-AxisMeta = TimeAxisMeta | ChannelAxisMeta | SpaceAxisMeta | NonstandardAxisMeta
+AxisMeta = Annotated[
+    TimeAxisMeta | ChannelAxisMeta | SpaceAxisMeta,
+    Discriminator("type"),
+]
 
 
 class TransformationMeta(MetaBase):
@@ -227,7 +226,7 @@ class MultiScaleMeta(VersionMeta):
     """https://ngff.openmicroscopy.org/0.4/index.html#multiscale-md"""
 
     # MUST
-    axes: list[AxisMeta] = Field(..., discriminator="type")
+    axes: list[AxisMeta]
     # MUST
     datasets: list[DatasetMeta]
     # SHOULD