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.
- boulder_opal_scale_up_sdk-1.0.0/PKG-INFO +38 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/__init__.py +14 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/agent/__init__.py +29 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/agent/worker.py +244 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/common/__init__.py +12 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/common/dtypes.py +353 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/common/typeclasses.py +85 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/__init__.py +16 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/common.py +58 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/config_loader.py +88 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/controller/__init__.py +32 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/controller/base.py +18 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/controller/qblox.py +664 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/controller/quantum_machines.py +139 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/device.py +35 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/processor/__init__.py +23 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/processor/common.py +148 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/device/processor/superconducting_processor.py +291 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/experiments/__init__.py +44 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/experiments/common.py +96 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/experiments/power_rabi.py +60 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/experiments/ramsey.py +55 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/experiments/resonator_spectroscopy.py +64 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/experiments/resonator_spectroscopy_by_bias.py +76 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/experiments/resonator_spectroscopy_by_power.py +64 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/grpc_interceptors/__init__.py +16 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/grpc_interceptors/auth.py +101 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/plotting/__init__.py +24 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/plotting/dtypes.py +221 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/protobuf/v1/agent_pb2.py +48 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/protobuf/v1/agent_pb2.pyi +53 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/protobuf/v1/agent_pb2_grpc.py +138 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/protobuf/v1/device_pb2.py +71 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/protobuf/v1/device_pb2.pyi +110 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/protobuf/v1/device_pb2_grpc.py +274 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/protobuf/v1/task_pb2.py +53 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/protobuf/v1/task_pb2.pyi +118 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/protobuf/v1/task_pb2_grpc.py +119 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/py.typed +0 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/routines/__init__.py +9 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/routines/common.py +10 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/routines/resonator_mapping.py +13 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/third_party/__init__.py +14 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/third_party/quantum_machines/__init__.py +51 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/third_party/quantum_machines/config.py +597 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/third_party/quantum_machines/constants.py +20 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/utils/__init__.py +12 -0
- boulder_opal_scale_up_sdk-1.0.0/boulderopalscaleupsdk/utils/serial_utils.py +62 -0
- 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
|
+
]
|