boulder-opal-scale-up-sdk 1.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (49) hide show
  1. boulder_opal_scale_up_sdk-1.0.0/PKG-INFO +38 -0
  2. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/__init__.py +14 -0
  3. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/agent/__init__.py +29 -0
  4. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/agent/worker.py +244 -0
  5. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/common/__init__.py +12 -0
  6. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/common/dtypes.py +353 -0
  7. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/common/typeclasses.py +85 -0
  8. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/__init__.py +16 -0
  9. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/common.py +58 -0
  10. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/config_loader.py +88 -0
  11. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/controller/__init__.py +32 -0
  12. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/controller/base.py +18 -0
  13. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/controller/qblox.py +664 -0
  14. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/controller/quantum_machines.py +139 -0
  15. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/device.py +35 -0
  16. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/processor/__init__.py +23 -0
  17. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/processor/common.py +148 -0
  18. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/processor/superconducting_processor.py +291 -0
  19. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/experiments/__init__.py +44 -0
  20. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/experiments/common.py +96 -0
  21. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/experiments/power_rabi.py +60 -0
  22. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/experiments/ramsey.py +55 -0
  23. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/experiments/resonator_spectroscopy.py +64 -0
  24. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/experiments/resonator_spectroscopy_by_bias.py +76 -0
  25. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/experiments/resonator_spectroscopy_by_power.py +64 -0
  26. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/grpc_interceptors/__init__.py +16 -0
  27. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/grpc_interceptors/auth.py +101 -0
  28. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/plotting/__init__.py +24 -0
  29. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/plotting/dtypes.py +221 -0
  30. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/protobuf/v1/agent_pb2.py +48 -0
  31. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/protobuf/v1/agent_pb2.pyi +53 -0
  32. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/protobuf/v1/agent_pb2_grpc.py +138 -0
  33. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/protobuf/v1/device_pb2.py +71 -0
  34. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/protobuf/v1/device_pb2.pyi +110 -0
  35. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/protobuf/v1/device_pb2_grpc.py +274 -0
  36. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/protobuf/v1/task_pb2.py +53 -0
  37. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/protobuf/v1/task_pb2.pyi +118 -0
  38. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/protobuf/v1/task_pb2_grpc.py +119 -0
  39. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/py.typed +0 -0
  40. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/routines/__init__.py +9 -0
  41. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/routines/common.py +10 -0
  42. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/routines/resonator_mapping.py +13 -0
  43. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/third_party/__init__.py +14 -0
  44. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/third_party/quantum_machines/__init__.py +51 -0
  45. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/third_party/quantum_machines/config.py +597 -0
  46. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/third_party/quantum_machines/constants.py +20 -0
  47. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/utils/__init__.py +12 -0
  48. boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/utils/serial_utils.py +62 -0
  49. boulder_opal_scale_up_sdk-1.0.0/pyproject.toml +83 -0
@@ -0,0 +1,38 @@
1
+ Metadata-Version: 2.3
2
+ Name: boulder-opal-scale-up-sdk
3
+ Version: 1.0.0
4
+ Summary: Q-CTRL Boulder Opal Scale Up Python SDK
5
+ License: Apache-2.0
6
+ Keywords: black opal,boulder opal,fire opal,nisq,open controls,q control,q ctrl,q-control,q-ctrl,qcontrol,qctrl,quantum,quantum algorithms,quantum circuits,quantum coding,quantum coding software,quantum computing,quantum control,quantum control software,quantum control theory,quantum engineering,quantum error correction,quantum firmware,quantum fundamentals,quantum sensing,qubit,qudit
7
+ Author: Q-CTRL
8
+ Author-email: support@q-ctrl.com
9
+ Maintainer: Q-CTRL
10
+ Maintainer-email: support@q-ctrl.com
11
+ Requires-Python: >=3.10,<3.13
12
+ Classifier: Development Status :: 5 - Production/Stable
13
+ Classifier: Environment :: Console
14
+ Classifier: Intended Audience :: Developers
15
+ Classifier: Intended Audience :: Education
16
+ Classifier: Intended Audience :: Science/Research
17
+ Classifier: License :: OSI Approved :: Apache Software License
18
+ Classifier: Natural Language :: English
19
+ Classifier: Operating System :: OS Independent
20
+ Classifier: Programming Language :: Python :: 3
21
+ Classifier: Programming Language :: Python :: 3.10
22
+ Classifier: Programming Language :: Python :: 3.11
23
+ Classifier: Programming Language :: Python :: 3.12
24
+ Classifier: Topic :: Internet :: WWW/HTTP
25
+ Classifier: Topic :: Scientific/Engineering :: Physics
26
+ Classifier: Topic :: Scientific/Engineering :: Visualization
27
+ Classifier: Topic :: Software Development :: Embedded Systems
28
+ Classifier: Topic :: System :: Distributed Computing
29
+ Requires-Dist: async-timeout (>=4.0.3,<5.0.0) ; python_version < "3.11"
30
+ Requires-Dist: attrs (>=25.1.0,<26.0.0)
31
+ Requires-Dist: googleapis-common-protos (>=1.69.2,<2.0.0)
32
+ Requires-Dist: qctrl-client (>=10.1.0,<11.0.0)
33
+ Project-URL: Facebook, https://www.facebook.com/qctrl
34
+ Project-URL: GitHub, https://github.com/qctrl
35
+ Project-URL: Homepage, https://q-ctrl.com
36
+ Project-URL: LinkedIn, https://www.linkedin.com/company/q-ctrl/
37
+ Project-URL: X, https://x.com/qctrlHQ
38
+ Project-URL: YouTube, https://www.youtube.com/qctrl
@@ -0,0 +1,14 @@
1
+ # Copyright 2025 Q-CTRL. All rights reserved.
2
+ #
3
+ # Licensed under the Q-CTRL Terms of service (the "License"). Unauthorized
4
+ # copying or use of this file, via any medium, is strictly prohibited.
5
+ # Proprietary and confidential. You may not use this file except in compliance
6
+ # with the License. You may obtain a copy of the License at
7
+ #
8
+ # https://q-ctrl.com/terms
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS. See the
12
+ # License for the specific language.
13
+
14
+ """Boulder Opal Scale up SDK package."""
@@ -0,0 +1,29 @@
1
+ # Copyright 2025 Q-CTRL. All rights reserved.
2
+ #
3
+ # Licensed under the Q-CTRL Terms of service (the "License"). Unauthorized
4
+ # copying or use of this file, via any medium, is strictly prohibited.
5
+ # Proprietary and confidential. You may not use this file except in compliance
6
+ # with the License. You may obtain a copy of the License at
7
+ #
8
+ # https://q-ctrl.com/terms
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS. See the
12
+ # License for the specific language.
13
+
14
+ """Agent module.
15
+
16
+ Agents are responsible for communicating with Q-CTRL APIs to determine a set of tasks to
17
+ perform (generally, QPU tasks such as execution of quantum circuits / configuration).
18
+ Importantly, to both simplify and ensure the security of the QPU adjacent host that the
19
+ agent will run on, the agent must initiate all communications (as opposed to opening a
20
+ port on the host).
21
+ """
22
+
23
+ __all__ = [
24
+ "Agent",
25
+ "AgentSettings",
26
+ "TaskHandler",
27
+ ]
28
+
29
+ from boulderopalscaleupsdk.agent.worker import Agent, AgentSettings, TaskHandler
@@ -0,0 +1,244 @@
1
+ # Copyright 2025 Q-CTRL. All rights reserved.
2
+ #
3
+ # Licensed under the Q-CTRL Terms of service (the "License"). Unauthorized
4
+ # copying or use of this file, via any medium, is strictly prohibited.
5
+ # Proprietary and confidential. You may not use this file except in compliance
6
+ # with the License. You may obtain a copy of the License at
7
+ #
8
+ # https://q-ctrl.com/terms
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS. See the
12
+ # License for the specific language.
13
+
14
+ """Base agent module.
15
+
16
+ Contains logic for agent worker.
17
+ """
18
+
19
+ import logging
20
+ from abc import abstractmethod
21
+ from typing import Protocol
22
+
23
+ from google.protobuf import any_pb2
24
+ from google.protobuf.message import Message
25
+ from google.protobuf.struct_pb2 import Struct
26
+ from grpc import aio as grpc
27
+ from grpc import ssl_channel_credentials
28
+ from pydantic_settings import BaseSettings, SettingsConfigDict
29
+
30
+ from boulderopalscaleupsdk.protobuf.v1 import agent_pb2, task_pb2, task_pb2_grpc
31
+
32
+ LOG = logging.getLogger(__name__)
33
+
34
+
35
+ class AgentSettings(BaseSettings):
36
+ """Configuration for the Agent."""
37
+
38
+ model_config = SettingsConfigDict(
39
+ env_file=".env",
40
+ env_prefix="QCTRL_SCALE_UP_",
41
+ env_file_encoding="utf-8",
42
+ env_nested_delimiter="__",
43
+ extra="ignore",
44
+ )
45
+
46
+ agent_id: str
47
+ remote_url: str
48
+
49
+
50
+ class TaskHandler(Protocol):
51
+ @abstractmethod
52
+ async def handle(
53
+ self,
54
+ request: agent_pb2.RunQuaProgramRequest
55
+ | agent_pb2.RunQuantumMachinesMixerCalibrationRequest
56
+ | agent_pb2.DisplayResultsRequest,
57
+ ) -> (
58
+ agent_pb2.RunQuaProgramResponse
59
+ | agent_pb2.RunQuantumMachinesMixerCalibrationResponse
60
+ | agent_pb2.DisplayResultsResponse
61
+ | task_pb2.TaskErrorDetail
62
+ ): ...
63
+
64
+ async def run(
65
+ self,
66
+ task: task_pb2.Task,
67
+ ) -> any_pb2.Any | task_pb2.TaskErrorDetail:
68
+ request = (
69
+ _as_run_qua_program_request(
70
+ task.data,
71
+ )
72
+ or _as_run_qua_calibration_request(task.data)
73
+ or _as_display_results_request(task.data)
74
+ )
75
+ match request:
76
+ case (
77
+ agent_pb2.RunQuaProgramRequest()
78
+ | agent_pb2.RunQuantumMachinesMixerCalibrationRequest()
79
+ | agent_pb2.DisplayResultsRequest()
80
+ ):
81
+ return _as_any_message(await self.handle(request))
82
+ case None:
83
+ return task_pb2.TaskErrorDetail(
84
+ code=task_pb2.TaskError.TASK_ERROR_NOT_IMPLEMENTED,
85
+ detail="Unknown or unrecognized request.",
86
+ )
87
+
88
+
89
+ def _as_any_message(message: Message) -> any_pb2.Any:
90
+ msg = any_pb2.Any()
91
+ msg.Pack(message) # type: ignore[reportUnknownMemberType]
92
+ return msg
93
+
94
+
95
+ def _as_run_qua_program_request(
96
+ task_result: any_pb2.Any,
97
+ ) -> agent_pb2.RunQuaProgramRequest | None:
98
+ request = agent_pb2.RunQuaProgramRequest()
99
+ unpacked: bool = task_result.Unpack(request) # type: ignore[reportUnknownMemberType]
100
+ if not unpacked:
101
+ return None
102
+
103
+ return request
104
+
105
+
106
+ def _as_run_qua_calibration_request(
107
+ task_result: any_pb2.Any,
108
+ ) -> agent_pb2.RunQuantumMachinesMixerCalibrationRequest | None:
109
+ request = agent_pb2.RunQuantumMachinesMixerCalibrationRequest()
110
+ unpacked: bool = task_result.Unpack(request) # type: ignore[reportUnknownMemberType]
111
+ if not unpacked:
112
+ return None
113
+
114
+ return request
115
+
116
+
117
+ def _as_display_results_request(
118
+ task_result: any_pb2.Any,
119
+ ) -> agent_pb2.DisplayResultsRequest | None:
120
+ request = agent_pb2.DisplayResultsRequest()
121
+ unpacked: bool = task_result.Unpack(request) # type: ignore[reportUnknownMemberType]
122
+ if not unpacked:
123
+ return None
124
+
125
+ return request
126
+
127
+
128
+ class Agent:
129
+ """Agent implementation."""
130
+
131
+ def __init__(
132
+ self,
133
+ agent_settings: AgentSettings,
134
+ handler: TaskHandler,
135
+ grpc_interceptors: list[grpc.ClientInterceptor] | None = None,
136
+ ) -> None:
137
+ self._settings = agent_settings
138
+ self._handler = handler
139
+ self._channel: grpc.Channel | None = self._create_channel(
140
+ agent_settings.remote_url,
141
+ grpc_interceptors,
142
+ )
143
+ self._agent_manager = task_pb2_grpc.AgentManagerServiceStub(self._channel)
144
+ self._state: task_pb2.AgentState = task_pb2.AGENT_STATE_ACTIVE_IDLE
145
+
146
+ @property
147
+ def agent_id(self) -> str:
148
+ """The agent identifier."""
149
+ return self._settings.agent_id
150
+
151
+ def _create_channel(
152
+ self,
153
+ url: str,
154
+ interceptors: list[grpc.ClientInterceptor] | None = None,
155
+ ) -> grpc.Channel:
156
+ """
157
+ Create a gRPC channel.
158
+ """
159
+ host = url.split(":")[0]
160
+ if host in ["localhost", "127.0.0.1", "0.0.0.0", "::"]:
161
+ channel = grpc.insecure_channel(url)
162
+ else:
163
+ channel = grpc.secure_channel(
164
+ url,
165
+ credentials=ssl_channel_credentials(),
166
+ interceptors=interceptors,
167
+ )
168
+ return channel
169
+
170
+ async def start_session(
171
+ self,
172
+ app: str,
173
+ device_name: str,
174
+ routine: str,
175
+ data: Struct | None,
176
+ ) -> str:
177
+ if not self._channel:
178
+ raise RuntimeError("Agent is shutdown.")
179
+ _data = any_pb2.Any()
180
+ if data:
181
+ _data.Pack(data)
182
+
183
+ response: task_pb2.AgentTasksResponse = await self._agent_manager.StartSession(
184
+ task_pb2.StartSessionRequest(
185
+ agent_id=self.agent_id,
186
+ app_name=app,
187
+ device_name=device_name,
188
+ routine_name=routine,
189
+ data=_data,
190
+ ),
191
+ )
192
+ self._state = response.target_state
193
+ await self._resume(response)
194
+ await self.shutdown()
195
+ return response.session_id
196
+
197
+ async def _resume(self, response: task_pb2.AgentTasksResponse) -> None:
198
+ tasks = response.tasks
199
+ while self._state not in [
200
+ task_pb2.AGENT_STATE_SHUTDOWN_MANAGER_INITIATED,
201
+ task_pb2.AGENT_STATE_SHUTDOWN_FAULT,
202
+ task_pb2.AGENT_STATE_SHUTDOWN_CLIENT_INITIATED,
203
+ ]:
204
+ task_results: list[task_pb2.TaskResult] = []
205
+ for task in tasks:
206
+ err: task_pb2.TaskErrorDetail | None = None
207
+ result: any_pb2.Any | None = None
208
+
209
+ match await self._handler.run(task):
210
+ case task_pb2.TaskErrorDetail() as err:
211
+ pass
212
+ case any_pb2.Any() as result:
213
+ pass
214
+
215
+ task_results.append(
216
+ task_pb2.TaskResult(
217
+ task_id=task.task_id,
218
+ session_id=task.session_id,
219
+ result=result,
220
+ error=err,
221
+ ),
222
+ )
223
+
224
+ _resp: task_pb2.AgentTasksResponse = await self._agent_manager.UpdateSession(
225
+ task_pb2.UpdateSessionRequest(
226
+ session_id=response.session_id,
227
+ current_state=self._state,
228
+ results=task_results,
229
+ task_in_progress=[],
230
+ ),
231
+ )
232
+
233
+ tasks = _resp.tasks
234
+ self._state = _resp.target_state
235
+
236
+ async def shutdown(
237
+ self,
238
+ timeout: float | None = None,
239
+ _: BaseException | None = None,
240
+ ) -> None:
241
+ """Shutdown the agent."""
242
+ if self._channel:
243
+ await self._channel.close(grace=timeout)
244
+ self._channel = None
@@ -0,0 +1,12 @@
1
+ # Copyright 2025 Q-CTRL. All rights reserved.
2
+ #
3
+ # Licensed under the Q-CTRL Terms of service (the "License"). Unauthorized
4
+ # copying or use of this file, via any medium, is strictly prohibited.
5
+ # Proprietary and confidential. You may not use this file except in compliance
6
+ # with the License. You may obtain a copy of the License at
7
+ #
8
+ # https://q-ctrl.com/terms
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS. See the
12
+ # License for the specific language.
@@ -0,0 +1,353 @@
1
+ # Copyright 2025 Q-CTRL. All rights reserved.
2
+ #
3
+ # Licensed under the Q-CTRL Terms of service (the "License"). Unauthorized
4
+ # copying or use of this file, via any medium, is strictly prohibited.
5
+ # Proprietary and confidential. You may not use this file except in compliance
6
+ # with the License. You may obtain a copy of the License at
7
+ #
8
+ # https://q-ctrl.com/terms
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS. See the
12
+ # License for the specific language.
13
+
14
+ from __future__ import annotations
15
+
16
+ from typing import overload
17
+
18
+ __all__ = [
19
+ "Duration",
20
+ "Self",
21
+ "TimeUnit",
22
+ ]
23
+
24
+ import sys
25
+ from dataclasses import dataclass, field
26
+ from datetime import datetime, timedelta
27
+ from decimal import Decimal
28
+ from enum import Enum
29
+ from typing import Annotated, Any, Literal
30
+
31
+ import numpy as np
32
+ from dateutil.parser import isoparse
33
+ from pydantic import BeforeValidator, PlainSerializer, TypeAdapter
34
+ from pydantic.dataclasses import dataclass as pydantic_dataclass
35
+
36
+ if sys.version_info >= (3, 11):
37
+ from typing import Self
38
+ else:
39
+ from typing_extensions import Self
40
+
41
+
42
+ class BaseType: ...
43
+
44
+
45
+ class FrequencyUnit(str, Enum):
46
+ Hz = "Hz"
47
+
48
+
49
+ @pydantic_dataclass
50
+ class Frequency:
51
+ value: float
52
+ unit: FrequencyUnit # No default to guarantee clarity of units
53
+
54
+ @classmethod
55
+ def from_float_hz(cls, value: float) -> Frequency:
56
+ return cls(value, FrequencyUnit.Hz)
57
+
58
+ def to_int_hz(self) -> int:
59
+ return int(self.to_float_hz())
60
+
61
+ def to_float_hz(self) -> float:
62
+ match self.unit:
63
+ case FrequencyUnit.Hz:
64
+ return self.value
65
+
66
+ def __gt__(self, other: Frequency) -> bool:
67
+ return self.to_float_hz() > other.to_float_hz()
68
+
69
+ def __ge__(self, other: Frequency) -> bool:
70
+ return self.to_float_hz() >= other.to_float_hz()
71
+
72
+ def __lt__(self, other: Frequency) -> bool:
73
+ return self.to_float_hz() < other.to_float_hz()
74
+
75
+ def __le__(self, other: Frequency) -> bool:
76
+ return self.to_float_hz() <= other.to_float_hz()
77
+
78
+ def __sub__(self, rhs: Frequency) -> Frequency:
79
+ if self.unit == rhs.unit:
80
+ return Frequency(self.value - rhs.value, self.unit)
81
+ raise NotImplementedError
82
+
83
+ def __add__(self, rhs: Frequency) -> Frequency:
84
+ if self.unit == rhs.unit:
85
+ return Frequency(self.value + rhs.value, self.unit)
86
+ raise NotImplementedError
87
+
88
+ def __abs__(self) -> Frequency:
89
+ return Frequency(abs(self.value), self.unit)
90
+
91
+ def __str__(self):
92
+ return f"{self.value} {self.unit.value}"
93
+
94
+ @overload # Division by a scalar: e.g. 4.4 Hz // 2.0 = 2.2 Hz
95
+ def __truediv__(self, rhs: float) -> Frequency: ...
96
+
97
+ @overload
98
+ def __truediv__(self, rhs: Frequency) -> float: ...
99
+
100
+ def __truediv__(self, rhs: float | Frequency) -> Frequency | float:
101
+ if isinstance(rhs, Frequency):
102
+ return self.to_float_hz() / rhs.to_float_hz()
103
+ return Frequency(self.value / rhs, self.unit)
104
+
105
+ @overload # Floor division by a scalar: e.g. 2.2 Hz // 2.0 = 1 Hz
106
+ def __floordiv__(self, rhs: float) -> Frequency: ...
107
+
108
+ @overload
109
+ def __floordiv__(self, rhs: Frequency) -> float: ...
110
+
111
+ def __floordiv__(self, rhs: float | Frequency) -> Frequency | float:
112
+ if isinstance(rhs, Frequency):
113
+ return self.to_float_hz() // rhs.to_float_hz()
114
+ return Frequency(self.value // rhs, self.unit)
115
+
116
+ def __mul__(self, rhs: float) -> Frequency:
117
+ return Frequency(self.value * rhs, self.unit)
118
+
119
+ def __rmul__(self, lhs: float) -> Frequency:
120
+ return self.__mul__(lhs)
121
+
122
+
123
+ class TimeUnit(str, Enum):
124
+ S = "s"
125
+ MS = "ms"
126
+ US = "us"
127
+ NS = "ns"
128
+ DT = "dt"
129
+
130
+
131
+ _SI_TIME = Literal[TimeUnit.S, TimeUnit.MS, TimeUnit.US, TimeUnit.NS]
132
+
133
+
134
+ @pydantic_dataclass(order=True)
135
+ class Duration(BaseType):
136
+ """
137
+ A wrapper of _SiDuration and _DtDuration to manage the conversion.
138
+ """
139
+
140
+ value: int = field(compare=False)
141
+ unit: TimeUnit = field(compare=False)
142
+ dtype: Literal["duration"] = "duration"
143
+ _value: _SiDuration | _DtDuration = field(init=False, repr=False)
144
+
145
+ def __post_init__(self):
146
+ self._value = (
147
+ _DtDuration(self.value)
148
+ if self.unit == TimeUnit.DT
149
+ else _SiDuration(self.value, self.unit)
150
+ )
151
+
152
+ def is_si(self) -> bool:
153
+ return self.unit != TimeUnit.DT
154
+
155
+ @staticmethod
156
+ def from_si(d: Duration, name: str) -> Duration:
157
+ if d.unit == TimeUnit.DT:
158
+ raise TypeError(f"{name} must use SI time unit.")
159
+ return d
160
+
161
+ @staticmethod
162
+ def from_intlike(val: float, unit: TimeUnit) -> Duration:
163
+ if not np.double(val).is_integer():
164
+ raise ValueError("fail to create a Duration object. value must be an integer.")
165
+ return Duration(int(val), unit)
166
+
167
+ def convert(self, target: Duration | _SI_TIME) -> Duration:
168
+ """
169
+ In particular, we only allow the following conversions:
170
+
171
+ # ((1000, "ms"), "s") -> (1, "s")
172
+ (_SiDuration, _SI_TIME) -> _SiDuration
173
+
174
+ # ((4, "ns"), (2, "ns")) -> (2, "dt")
175
+ (_SiDuration, _SiDuration) -> _DtDuration
176
+
177
+ # ((2, "dt"), (2, "ns")) -> (4, "ns")
178
+ (_DtDuration, _SiDuration) _> _SiDuration
179
+ """
180
+ match self._value, getattr(target, "_value", target):
181
+ case _SiDuration(
182
+ _,
183
+ _,
184
+ ), TimeUnit.S | TimeUnit.MS | TimeUnit.US | TimeUnit.NS:
185
+ converted = self._value.convert_to_si(target) # type: ignore[arg-type] # pyright: ignore[reportArgumentType]
186
+ case _SiDuration(_, _), _SiDuration(_, _):
187
+ converted = self._value.convert_to_dt(target) # type: ignore[arg-type, assignment] # pyright: ignore[reportArgumentType]
188
+ case _DtDuration(_, _), _SiDuration(_, _):
189
+ converted = self._value.convert(target) # type: ignore[arg-type] # pyright: ignore[reportArgumentType]
190
+ case _:
191
+ raise TypeError(f"cant't convert type {self.unit} to {target}")
192
+ return Duration(converted.value, converted.unit)
193
+
194
+
195
+ @dataclass(order=True)
196
+ class _DtDuration(BaseType):
197
+ value: int
198
+ unit: Literal[TimeUnit.DT] = TimeUnit.DT
199
+
200
+ def convert(self, target: _SiDuration) -> _SiDuration:
201
+ return _SiDuration(self.value * target.value, target.unit)
202
+
203
+
204
+ @dataclass(order=True)
205
+ class _SiDuration(BaseType):
206
+ value: int = field(compare=False)
207
+ unit: _SI_TIME = field(compare=False)
208
+ _np_rep: np.timedelta64 = field(init=False, repr=False)
209
+
210
+ def __post_init__(self):
211
+ err = TypeError(
212
+ f"value must be an integer, got {self.value}{self.unit}. Choose a different unit to "
213
+ "scale it.",
214
+ )
215
+ dec = Decimal(self.value)
216
+ exponent = dec.as_tuple().exponent
217
+ if not isinstance(exponent, int) or exponent < 0: # pragma: no cover
218
+ raise err
219
+
220
+ self.value = int(self.value)
221
+ try:
222
+ self._np_rep = np.timedelta64(self.value, self.unit)
223
+ except ValueError as e:
224
+ raise err from e
225
+
226
+ def convert_to_dt(self, clock: _SiDuration) -> _DtDuration:
227
+ try:
228
+ converted_si = self.convert_to_si(clock.unit)
229
+ except TypeError as e:
230
+ raise TypeError(
231
+ "fail to convert to dt type. Consider rescaling the clock time.",
232
+ ) from e
233
+ val = np.double(converted_si.value / clock.value)
234
+ # N.B, this might be too strict. Some rounding might be necessary.
235
+ if not val.is_integer():
236
+ raise TypeError(
237
+ "fail to convert to dt type. Consider rescaling the clock time.",
238
+ )
239
+ return _DtDuration(int(val))
240
+
241
+ def convert_to_si(self, unit: _SI_TIME) -> _SiDuration:
242
+ if self._np_rep is None:
243
+ raise TypeError("`convert` only support SI time unit.")
244
+ val: np.float64 = self._np_rep / np.timedelta64(1, unit)
245
+ if not val.is_integer():
246
+ raise TypeError(
247
+ f"fail to convert to {unit} with {self.value}{self.unit}.",
248
+ )
249
+ return _SiDuration(int(val), unit)
250
+
251
+ def to_seconds(self) -> float:
252
+ return float(self._np_rep / np.timedelta64(1, "s"))
253
+
254
+
255
+ def ensure_frequency_hz(value: Any) -> Any:
256
+ match value:
257
+ case Frequency():
258
+ return value
259
+ case float() | int():
260
+ return Frequency(value, FrequencyUnit.Hz)
261
+ case dict():
262
+ return TypeAdapter(Frequency).validate_python(value)
263
+ case _:
264
+ raise ValueError("Frequency needs to be numeric.")
265
+
266
+
267
+ FrequencyHzLike = Annotated[
268
+ Frequency,
269
+ BeforeValidator(ensure_frequency_hz),
270
+ PlainSerializer(lambda x: x.to_float_hz(), return_type=float),
271
+ ]
272
+
273
+
274
+ def ensure_duration_ns(value: Any) -> Any:
275
+ match value:
276
+ case Duration():
277
+ return value.convert(TimeUnit.NS)
278
+ case float() | int():
279
+ return Duration.from_intlike(value, TimeUnit.NS)
280
+ case dict():
281
+ return TypeAdapter(Duration).validate_python(value)
282
+ case _:
283
+ raise ValueError("Duration needs to be numeric")
284
+
285
+
286
+ DurationNsLike = Annotated[Duration, BeforeValidator(ensure_duration_ns)]
287
+
288
+
289
+ @pydantic_dataclass
290
+ class ISO8601Datetime:
291
+ value: datetime
292
+
293
+ def __post_init__(self):
294
+ self.value = _validate_iso_datetime(self.value)
295
+
296
+ def __str__(self):
297
+ return _serialize_datetime(self.value)
298
+
299
+ def strftime(self, fmt: str) -> str:
300
+ """
301
+ Format the datetime value using the given format string.
302
+
303
+ Parameters
304
+ ----------
305
+ fmt : str
306
+ The format string to use for formatting.
307
+
308
+ Returns
309
+ -------
310
+ str
311
+ The formatted datetime string.
312
+ """
313
+ return self.value.strftime(fmt)
314
+
315
+
316
+ def _validate_iso_datetime(value: Any) -> datetime:
317
+ def _raise_invalid_timezone_error():
318
+ raise ValueError("Datetime must be in UTC timezone.")
319
+
320
+ if isinstance(value, ISO8601Datetime):
321
+ return value.value
322
+ if isinstance(value, datetime):
323
+ if value.tzinfo is None or value.tzinfo.utcoffset(value) != timedelta(0):
324
+ _raise_invalid_timezone_error()
325
+ else:
326
+ return value
327
+ if isinstance(value, str):
328
+ try:
329
+ parsed_datetime = isoparse(value)
330
+ if parsed_datetime.tzinfo is None or parsed_datetime.tzinfo.utcoffset(
331
+ parsed_datetime,
332
+ ) != timedelta(0):
333
+ _raise_invalid_timezone_error()
334
+ else:
335
+ return parsed_datetime
336
+ except Exception as e:
337
+ raise ValueError("Invalid ISO8601 datetime string.") from e
338
+ raise ValueError(
339
+ "Value must be a datetime object, an ISO8601Datetime instance, or a valid ISO8601 string.",
340
+ )
341
+
342
+
343
+ def _serialize_datetime(value: datetime) -> str:
344
+ if value.tzinfo is None or value.tzinfo.utcoffset(value) != timedelta(0):
345
+ raise ValueError("Datetime must be in UTC timezone.")
346
+ return value.isoformat()
347
+
348
+
349
+ ISO8601DatetimeUTCLike = Annotated[
350
+ datetime,
351
+ BeforeValidator(_validate_iso_datetime),
352
+ PlainSerializer(_serialize_datetime),
353
+ ]