Merge pull request #81 from OpenQuantumDesign/refactor_type_reflection

Refactor type reflection

Author: yhteoh

pyproject.toml

@@ -64,7 +64,7 @@ fixable = ["ALL"]
 oqd-compiler-infrastructure = { git = "https://github.com/openquantumdesign/oqd-compiler-infrastructure" }
 
 [dependency-groups]
-dev = ["jupyter>=1.1.1", "pre-commit>=4.1.0"]
+dev = ["jupyter>=1.1.1", "pre-commit>=4.1.0", "ruff>=0.13.1"]
 
 
 [project.urls]

src/oqd_core/backend/metric.py

@@ -18,7 +18,7 @@
 from pydantic.types import NonNegativeInt
 
 ########################################################################################
-from oqd_core.interface.analog.operator import OperatorSubtypes
+from oqd_core.interface.analog.operator import OperatorSubTypes
 
 ########################################################################################
 
@@ -32,7 +32,7 @@
 
 
 class Expectation(VisitableBaseModel):
-    operator: OperatorSubtypes
+    operator: OperatorSubTypes
 
 
 class EntanglementEntropyVN(VisitableBaseModel):

src/oqd_core/interface/analog/__init__.py

@@ -31,7 +31,7 @@
     OperatorMul,
     OperatorScalarMul,
     OperatorSub,
-    OperatorSubtypes,
+    OperatorSubTypes,
     OperatorTerminal,
     Pauli,
     PauliI,
@@ -43,7 +43,7 @@
 )
 
 __all__ = [
-    "OperatorSubtypes",
+    "OperatorSubTypes",
     "Operator",
     "OperatorTerminal",
     "Pauli",

src/oqd_core/interface/analog/operation.py

@@ -12,13 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Literal, Optional, Union
+from __future__ import annotations
 
-# %%
-from oqd_compiler_infrastructure import TypeReflectBaseModel, VisitableBaseModel
+from typing import Annotated, List, Optional, Union
+
+from oqd_compiler_infrastructure import TypeReflectBaseModel
+from pydantic import Discriminator
 from pydantic.types import NonNegativeInt
 
-from oqd_core.interface.analog.operator import OperatorSubtypes
+from oqd_core.interface.analog.operator import OperatorSubTypes
 
 __all__ = [
     "AnalogCircuit",
@@ -33,24 +35,26 @@
 ########################################################################################
 
 
-class AnalogGate(TypeReflectBaseModel):
+class AnalogOperation(TypeReflectBaseModel):
     """
-    Class representing an analog gate composed of Hamiltonian terms and dissipation terms
-
-    Attributes:
-        hamiltonian (Operator): Hamiltonian terms of the gate
+    Class representing an analog operation applied to the quantum system
     """
 
-    hamiltonian: OperatorSubtypes
+    pass
+
+
+########################################################################################
 
 
-# %%
-class AnalogOperation(VisitableBaseModel):
+class AnalogGate(TypeReflectBaseModel):
     """
-    Class representing an analog operation applied to the quantum system
+    Class representing an analog gate composed of Hamiltonian terms and dissipation terms
+
+    Attributes:
+        hamiltonian (Operator): Hamiltonian terms of the gate
     """
 
-    pass
+    hamiltonian: OperatorSubTypes
 
 
 class Evolve(AnalogOperation):
@@ -62,17 +66,18 @@ class Evolve(AnalogOperation):
         gate (AnalogGate): Analog gate to evolve by
     """
 
-    key: Literal["evolve"] = "evolve"
     duration: float
-    gate: Union[AnalogGate, str]
+    gate: AnalogGate
+
+
+########################################################################################
 
 
 class Measure(AnalogOperation):
     """
     Class representing a measurement in the analog circuit
     """
 
-    key: Literal["measure"] = "measure"
     targets: Optional[List[int]] = None
 
 
@@ -81,14 +86,20 @@ class Initialize(AnalogOperation):
     Class representing a initialization in the analog circuit
     """
 
-    key: Literal["initialize"] = "initialize"
     targets: Optional[List[int]] = None
 
 
+########################################################################################
+
 """
 Union of classes
 """
-Statement = Union[Measure, Evolve, Initialize]
+AnalogOperationSubTypes = Annotated[
+    Union[Measure, Evolve, Initialize], Discriminator(discriminator="class_")
+]
+
+
+########################################################################################
 
 
 class AnalogCircuit(AnalogOperation):
@@ -100,7 +111,7 @@ class AnalogCircuit(AnalogOperation):
 
     """
 
-    sequence: List[Statement] = []
+    sequence: List[AnalogOperationSubTypes] = []
 
     n_qreg: Union[NonNegativeInt, None] = None
     n_qmode: Union[NonNegativeInt, None] = None

src/oqd_core/interface/analog/operator.py

@@ -14,14 +14,15 @@
 
 from __future__ import annotations
 
-from typing import Union
+from typing import Annotated, Union
 
 from oqd_compiler_infrastructure import TypeReflectBaseModel
+from pydantic import Discriminator
 
 ########################################################################################
 from oqd_core.interface.math import (
     MathExpr,
-    MathExprSubtypes,
+    MathExprSubTypes,
     MathImag,
     MathMul,
     MathNum,
@@ -49,7 +50,7 @@
     "OperatorMul",
     "OperatorScalarMul",
     "OperatorKron",
-    "OperatorSubtypes",
+    "OperatorSubTypes",
 ]
 
 
@@ -223,8 +224,8 @@ class OperatorScalarMul(Operator):
         expr (MathExpr): [`MathExpr`][oqd_core.interface.math.MathExpr] to multiply by
     """
 
-    op: OperatorSubtypes
-    expr: MathExprSubtypes
+    op: OperatorSubTypes
+    expr: MathExprSubTypes
 
 
 class OperatorBinaryOp(Operator):
@@ -244,8 +245,8 @@ class OperatorAdd(OperatorBinaryOp):
         op2 (Operator): Right hand side [`Operator`][oqd_core.interface.analog.operator.Operator]
     """
 
-    op1: OperatorSubtypes
-    op2: OperatorSubtypes
+    op1: OperatorSubTypes
+    op2: OperatorSubTypes
 
 
 class OperatorSub(OperatorBinaryOp):
@@ -257,8 +258,8 @@ class OperatorSub(OperatorBinaryOp):
         op2 (Operator): Right hand side [`Operator`][oqd_core.interface.analog.operator.Operator]
     """
 
-    op1: OperatorSubtypes
-    op2: OperatorSubtypes
+    op1: OperatorSubTypes
+    op2: OperatorSubTypes
 
 
 class OperatorMul(OperatorBinaryOp):
@@ -270,8 +271,8 @@ class OperatorMul(OperatorBinaryOp):
         op2 (Operator): Right hand side [`Operator`][oqd_core.interface.analog.operator.Operator]
     """
 
-    op1: OperatorSubtypes
-    op2: OperatorSubtypes
+    op1: OperatorSubTypes
+    op2: OperatorSubTypes
 
 
 class OperatorKron(OperatorBinaryOp):
@@ -283,25 +284,28 @@ class OperatorKron(OperatorBinaryOp):
         op2 (Operator): Right hand side [`Operator`][oqd_core.interface.analog.operator.Operator]
     """
 
-    op1: OperatorSubtypes
-    op2: OperatorSubtypes
+    op1: OperatorSubTypes
+    op2: OperatorSubTypes
 
 
 ########################################################################################
 
-OperatorSubtypes = Union[
-    PauliI,
-    PauliX,
-    PauliY,
-    PauliZ,
-    Creation,
-    Annihilation,
-    Identity,
-    OperatorAdd,
-    OperatorSub,
-    OperatorMul,
-    OperatorScalarMul,
-    OperatorKron,
+OperatorSubTypes = Annotated[
+    Union[
+        PauliI,
+        PauliX,
+        PauliY,
+        PauliZ,
+        Creation,
+        Annihilation,
+        Identity,
+        OperatorAdd,
+        OperatorSub,
+        OperatorMul,
+        OperatorScalarMul,
+        OperatorKron,
+    ],
+    Discriminator(discriminator="class_"),
 ]
 """
 Alias for the union of concrete Operator subtypes

src/oqd_core/interface/atomic/protocol.py

@@ -14,10 +14,10 @@
 
 from __future__ import annotations
 
-from typing import List, Tuple, Union
+from typing import Annotated, List, Tuple, Union
 
 from oqd_compiler_infrastructure import TypeReflectBaseModel
-from pydantic import conlist
+from pydantic import Discriminator, conlist
 
 from oqd_core.interface.atomic.system import Transition
 from oqd_core.interface.math import CastMathExpr, ConstantMathExpr
@@ -103,7 +103,7 @@ class ParallelProtocol(Protocol):
         sequence: List of pulses or subprotocols to compose together in a parallel fashion.
     """
 
-    sequence: List[Union[PulseSubTypes, ProtocolSubTypes]]
+    sequence: List[ProtocolPulseSubTypes]
 
 
 class SequentialProtocol(Protocol):
@@ -114,8 +114,16 @@ class SequentialProtocol(Protocol):
         sequence: List of pulses or subprotocols to compose together in a sequntial fashion.
     """
 
-    sequence: List[Union[PulseSubTypes, ProtocolSubTypes]]
+    sequence: List[ProtocolPulseSubTypes]
 
 
-ProtocolSubTypes = Union[SequentialProtocol, ParallelProtocol]
-PulseSubTypes = Union[Pulse, MeasurePulse]
+ProtocolSubTypes = Annotated[
+    Union[SequentialProtocol, ParallelProtocol], Discriminator(discriminator="class_")
+]
+PulseSubTypes = Annotated[
+    Union[Pulse, MeasurePulse], Discriminator(discriminator="class_")
+]
+ProtocolPulseSubTypes = Annotated[
+    Union[SequentialProtocol, ParallelProtocol, Pulse, MeasurePulse],
+    Discriminator(discriminator="class_"),
+]

src/oqd_core/interface/math.py

@@ -48,7 +48,7 @@
     "MathMul",
     "MathDiv",
     "MathPow",
-    "MathExprSubtypes",
+    "MathExprSubTypes",
     "ConstantMathExpr",
     "CastMathExpr",
 ]
@@ -339,27 +339,25 @@ class MathPow(MathBinaryOp):
 
 ########################################################################################
 
-MathExprSubtypes = Annotated[
+MathExprSubTypes = Annotated[
     Union[
-        Annotated[MathNum, Tag("MathNum")],
-        Annotated[MathVar, Tag("MathVar")],
-        Annotated[MathImag, Tag("MathImag")],
-        Annotated[MathFunc, Tag("MathFunc")],
-        Annotated[MathAdd, Tag("MathAdd")],
-        Annotated[MathSub, Tag("MathSub")],
-        Annotated[MathMul, Tag("MathMul")],
-        Annotated[MathDiv, Tag("MathDiv")],
-        Annotated[MathPow, Tag("MathPow")],
+        MathNum,
+        MathVar,
+        MathImag,
+        MathFunc,
+        MathAdd,
+        MathSub,
+        MathMul,
+        MathDiv,
+        MathPow,
     ],
-    Discriminator(
-        lambda v: v["class_"] if isinstance(v, dict) else getattr(v, "class_")
-    ),
+    Discriminator(discriminator="class_"),
 ]
 """
 Alias for the union of concrete MathExpr subtypes
 """
 
-CastMathExpr = Annotated[MathExprSubtypes, BeforeValidator(MathExpr.cast)]
+CastMathExpr = Annotated[MathExprSubTypes, BeforeValidator(MathExpr.cast)]
 """
 Annotated type that cast typical numeric python types to MathExpr
 """