|
| 1 | +import json |
| 2 | +from datetime import datetime |
| 3 | +from typing import Any |
| 4 | + |
| 5 | +import requests |
| 6 | +from pydantic import AnyHttpUrl, BaseModel |
| 7 | +from pydantic_settings import BaseSettings, SettingsConfigDict |
| 8 | +from requests.exceptions import HTTPError, RequestException |
| 9 | + |
| 10 | +from interactem.core.logger import get_logger |
| 11 | +from interactem.core.models.messages import BytesMessage |
| 12 | +from interactem.operators.operator import operator |
| 13 | + |
| 14 | + |
| 15 | +class Settings(BaseSettings): |
| 16 | + """Settings for communicating with the Distiller API. ENV |
| 17 | + variables will contain the necessary secrets.""" |
| 18 | + |
| 19 | + model_config = SettingsConfigDict(case_sensitive=True) |
| 20 | + DISTILLER_API_URL: AnyHttpUrl |
| 21 | + DISTILLER_API_KEY_NAME: str |
| 22 | + DISTILLER_API_KEY: str |
| 23 | + |
| 24 | + |
| 25 | +class Location(BaseModel): |
| 26 | + """The location of a data set. If host is perlmutter then |
| 27 | + the streaming operation is finished. The path is the path to the data. |
| 28 | +
|
| 29 | + Attributes |
| 30 | + ---------- |
| 31 | + host : str |
| 32 | + The name of the host where the data is located. e.g. Perlmutter |
| 33 | + path : str |
| 34 | + The file path on the host to where the data is located. |
| 35 | + """ |
| 36 | + |
| 37 | + host: str # you'll look for host = perlmutter |
| 38 | + path: str |
| 39 | + |
| 40 | + |
| 41 | +class Scan(BaseModel): |
| 42 | + """Information about a scan captured in the Distiller system. |
| 43 | +
|
| 44 | + A Scan contains metadata and location data from a 4D Camera dataset, |
| 45 | + including identifiers for both the Distiller system. Each scan tracks |
| 46 | + multiple location points and maintains creation timestamp information. |
| 47 | +
|
| 48 | + Attributes |
| 49 | + ---------- |
| 50 | + id : int |
| 51 | + Unique identifier for the scan within the Distiller system. |
| 52 | + scan_id : int, optional |
| 53 | + Identifier from the 4D camera system. May be None if the scan |
| 54 | + was not captured using 4D camera equipment. |
| 55 | + locations : list[Location] |
| 56 | + List of Location objects representing spatial data points |
| 57 | + captured during the scan. |
| 58 | + created : datetime |
| 59 | + Timestamp indicating when the scan was created. |
| 60 | + image_path : str, optional |
| 61 | + File path to an associated image. Defaults to None if no image |
| 62 | + is stored with the scan. |
| 63 | + metadata : dict, optional |
| 64 | + The metadata for the scan. |
| 65 | + notes : str, optional |
| 66 | + Notes from Distiller |
| 67 | + """ |
| 68 | + |
| 69 | + id: int # distiller id |
| 70 | + scan_id: int | None # scan id from 4d camera |
| 71 | + locations: list[Location] |
| 72 | + created: datetime |
| 73 | + image_path: str | None = None |
| 74 | + notes: str | None |
| 75 | + metadata: dict[str, Any] | None # = Field(alias="metadata_") |
| 76 | + |
| 77 | + |
| 78 | +def get_scans( |
| 79 | + skip: int = 0, |
| 80 | + limit: int = 100, |
| 81 | + scan_id: int = -1, |
| 82 | + start: datetime | None = None, |
| 83 | + end: datetime | None = None, |
| 84 | + job_id: int | None = None, |
| 85 | +) -> list[Scan]: |
| 86 | + """ |
| 87 | + Fetch a list of scans with various filter options. |
| 88 | +
|
| 89 | + Parameters: |
| 90 | + skip (int): Number of records to skip. |
| 91 | + limit (int): Maximum number of records to return. |
| 92 | + scan_id (int): Specific scan ID to filter by. |
| 93 | + start (Optional[datetime]): Start of the date range for creation time. |
| 94 | + end (Optional[datetime]): End of the date range for creation time. |
| 95 | + job_id (Optional[int]): Job ID to filter by. |
| 96 | +
|
| 97 | + Returns: |
| 98 | + List[Scan]: A list of Scan objects matching the criteria. |
| 99 | +
|
| 100 | + Raises: |
| 101 | + HTTPError: If the request fails due to an HTTP error. |
| 102 | + RequestException: For any other request-related errors. |
| 103 | + """ |
| 104 | + headers = { |
| 105 | + settings.DISTILLER_API_KEY_NAME: settings.DISTILLER_API_KEY, |
| 106 | + "Content-Type": "application/json", |
| 107 | + } |
| 108 | + |
| 109 | + params: dict[str, Any] = { |
| 110 | + "skip": skip, |
| 111 | + "limit": limit, |
| 112 | + } |
| 113 | + |
| 114 | + if scan_id != -1: |
| 115 | + params["scan_id"] = scan_id |
| 116 | + if start is not None: |
| 117 | + params["start"] = start.isoformat() |
| 118 | + if end is not None: |
| 119 | + params["end"] = end.isoformat() |
| 120 | + if job_id is not None: |
| 121 | + params["job_id"] = job_id |
| 122 | + |
| 123 | + url = f"{settings.DISTILLER_API_URL}/scans" |
| 124 | + |
| 125 | + try: |
| 126 | + response = requests.get(url, headers=headers, params=params) |
| 127 | + response.raise_for_status() # Raise an HTTPError for bad responses |
| 128 | + json_data = response.json() |
| 129 | + return [Scan(**scan_data) for scan_data in json_data] |
| 130 | + except HTTPError as http_err: |
| 131 | + raise HTTPError(f"HTTP error occurred: {http_err}") |
| 132 | + except RequestException as req_err: |
| 133 | + raise RequestException(f"Request exception occurred: {req_err}") |
| 134 | + |
| 135 | + |
| 136 | +def add_metadata(distiller_scan_id: int, metadata: dict[str, Any]): |
| 137 | + """Update the metadata field in Distiller. This can be used to |
| 138 | + store updated or supplemental metadata for a 4D STEM scan. |
| 139 | +
|
| 140 | + Parameters |
| 141 | + ---------- |
| 142 | + distiller_scan_id : int |
| 143 | + The Distiller scan id. This is a unique ID attached to each data set in the database. |
| 144 | + metadata : dict |
| 145 | + The metadata to merge into the existing metadata field in Distiller. |
| 146 | +
|
| 147 | + Returns |
| 148 | + ------- |
| 149 | + : Scan |
| 150 | + A Scan class object with information about the scan that was changed |
| 151 | + """ |
| 152 | + |
| 153 | + headers = { |
| 154 | + settings.DISTILLER_API_KEY_NAME: settings.DISTILLER_API_KEY, |
| 155 | + "Content-Type": "application/json", |
| 156 | + } |
| 157 | + |
| 158 | + url = f"{settings.DISTILLER_API_URL}/scans/{distiller_scan_id}" |
| 159 | + params = {"merge": True} |
| 160 | + |
| 161 | + try: |
| 162 | + response = requests.patch( |
| 163 | + url, |
| 164 | + headers=headers, |
| 165 | + params=params, |
| 166 | + data=json.dumps({"metadata": metadata}), |
| 167 | + ) |
| 168 | + response.raise_for_status() |
| 169 | + json_data = response.json() |
| 170 | + # print(json_data) |
| 171 | + return Scan(**json_data) |
| 172 | + except HTTPError as http_err: |
| 173 | + raise HTTPError(f"HTTP error occurred: {http_err}") |
| 174 | + except RequestException as req_err: |
| 175 | + raise RequestException(f"Request exception occurred: {req_err}") |
| 176 | + |
| 177 | + |
| 178 | +logger = get_logger() |
| 179 | +global settings |
| 180 | +settings = Settings() |
| 181 | + |
| 182 | +@operator |
| 183 | +def update_distiller_metadata( |
| 184 | + inputs: BytesMessage | None, parameters: dict[str, Any] |
| 185 | +) -> BytesMessage | None: |
| 186 | + """Adds information about a scan to the Disitller metadata""" |
| 187 | + global settings |
| 188 | + if not inputs: |
| 189 | + logger.warning("No input provided to the update_distiller_metadata operator.") |
| 190 | + return None |
| 191 | + |
| 192 | + # TODO: Implement operator logic here |
| 193 | + logger.info("update_distiller_metadata operator running...") |
| 194 | + |
| 195 | + # We have to have a scan_number to identify the scan in Distiller |
| 196 | + if "scan_number" in inputs.header.meta: |
| 197 | + scan_number = inputs.header.meta["scan_number"] |
| 198 | + logger.info(f"Detector Scan Number: {scan_number}") |
| 199 | + else: |
| 200 | + logger.info("Detector Scan Number not found in metadata.") |
| 201 | + return |
| 202 | + |
| 203 | + # Extract C12 magnitude and angle from metadata if available |
| 204 | + C12_magnitude = inputs.header.meta.get("C12", None) |
| 205 | + C12_angle = inputs.header.meta.get("phi12", None) |
| 206 | + |
| 207 | + if C12_magnitude is not None and C12_angle is not None: |
| 208 | + logger.info(f"C12: {C12_magnitude} {C12_angle}") |
| 209 | + # Add metadata to Distiller via its API |
| 210 | + # We can only use the detector scan_number to identify the scan in Distiller |
| 211 | + logger.info("Querying Distiller for matching scan...") |
| 212 | + scans = get_scans(limit=3) |
| 213 | + logger.info(f"Retrieved {len(scans)} scans from Distiller.") |
| 214 | + scan_ids = [] |
| 215 | + if scans: |
| 216 | + for scan in scans: |
| 217 | + if scan.scan_id == scan_number: |
| 218 | + logger.info(f"Found matching Distiller scan ID: {scan.id}") |
| 219 | + # save this scan_id |
| 220 | + scan_ids.append(scan.id) |
| 221 | + else: |
| 222 | + logger.info( |
| 223 | + f"No matching scan found for detector scan number: {scan_number}" |
| 224 | + ) |
| 225 | + else: |
| 226 | + logger.info("No scans found in Distiller.") |
| 227 | + |
| 228 | + if len(scan_ids) > 1: |
| 229 | + logger.warning( |
| 230 | + f"Multiple Distiller scans found for detector scan number: {scan_number}. Not updating metadata." |
| 231 | + ) |
| 232 | + elif len(scan_ids) == 1: |
| 233 | + distiller_scan_id = scan_ids[0] |
| 234 | + logger.info(f"Updating metadata for Distiller scan ID: {distiller_scan_id}") |
| 235 | + |
| 236 | + # Prepare metadata payload |
| 237 | + metadata_payload = { |
| 238 | + "C12_magnitude": C12_magnitude, |
| 239 | + "C12_angle": C12_angle, |
| 240 | + } |
| 241 | + updated_scan = add_metadata(distiller_scan_id, metadata_payload) |
| 242 | + logger.info(f"Updated Distiller scan metadata: {updated_scan.metadata}") |
| 243 | + else: |
| 244 | + logger.info("C12 Magnitude or phi12 not found in metadata.") |
| 245 | + return |
| 246 | + |
| 247 | + return None |
0 commit comments