|
| 1 | +"""High-level calcdocs API — minimal concepts, maximum discoverability. |
| 2 | +
|
| 3 | +This module provides a simplified interface for defining calculated data documents. |
| 4 | +It requires understanding only 3 concepts: |
| 5 | +
|
| 6 | +1. Meas — a raw measurement value from your data |
| 7 | +2. Calc — a calculated value derived from measurements or other calcs |
| 8 | +3. GroupBy — how to partition data for each calculation |
| 9 | +
|
| 10 | +Usage: |
| 11 | + from allotropy.calcdocs.api import calc_docs, Calc, GroupBy, Meas |
| 12 | +
|
| 13 | + def to_element(item: MyDataClass) -> dict: |
| 14 | + return { |
| 15 | + "uuid": item.uuid, |
| 16 | + "sample_id": item.sample_identifier, |
| 17 | + "value": item.measured_value, |
| 18 | + "mean": item.computed_mean, |
| 19 | + } |
| 20 | +
|
| 21 | + BY_SAMPLE = GroupBy("sample_id") |
| 22 | + BY_SAMPLE_UUID = GroupBy("sample_id", "uuid") |
| 23 | +
|
| 24 | + measurement = Meas("value", field="value") |
| 25 | + mean = Calc("mean", field="mean", sources=[measurement], group=BY_SAMPLE) |
| 26 | +
|
| 27 | + result = calc_docs(items, to_element, nodes=[measurement, mean]) |
| 28 | +""" |
| 29 | + |
| 30 | +from __future__ import annotations |
| 31 | + |
| 32 | +from collections import defaultdict |
| 33 | +from collections.abc import Callable |
| 34 | +from dataclasses import dataclass, field as dataclass_field |
| 35 | +from itertools import chain |
| 36 | +from typing import Any, TypeVar |
| 37 | + |
| 38 | +from allotropy.calcdocs.appbio_quantstudio_designandanalysis.config import ( |
| 39 | + CalculatedDataConfigWithOptional, |
| 40 | +) |
| 41 | +from allotropy.calcdocs.config import ( |
| 42 | + CalcDocsConfig, |
| 43 | + CalculatedDataConfig, |
| 44 | + MeasurementConfig, |
| 45 | +) |
| 46 | +from allotropy.calcdocs.extractor import Element |
| 47 | +from allotropy.calcdocs.view import Keys, View, ViewData |
| 48 | +from allotropy.parsers.utils.calculated_data_documents.definition import ( |
| 49 | + CalculatedDocument, |
| 50 | +) |
| 51 | + |
| 52 | +T = TypeVar("T") |
| 53 | + |
| 54 | + |
| 55 | +@dataclass(frozen=True) |
| 56 | +class GroupBy: |
| 57 | + """Defines how elements are partitioned for a calculation. |
| 58 | +
|
| 59 | + Each string in `fields` is a key in the element dict. Elements sharing the same |
| 60 | + values for all fields are grouped together. "uuid" is special — it groups by |
| 61 | + element identity (one element per group). |
| 62 | +
|
| 63 | + Args: |
| 64 | + *fields: Field names to group by, outermost first. |
| 65 | + filter_by: Only include elements where these field values match exactly. |
| 66 | + exclude: Exclude elements where the field value is in the given list. |
| 67 | + reference: Redirect key lookups to a fixed value (for reference sample/target). |
| 68 | + """ |
| 69 | + |
| 70 | + fields: tuple[str, ...] = () |
| 71 | + filter_by: dict[str, str] = dataclass_field(default_factory=dict) |
| 72 | + exclude: dict[str, list[str]] = dataclass_field(default_factory=dict) |
| 73 | + reference: dict[str, str] = dataclass_field(default_factory=dict) |
| 74 | + |
| 75 | + def __init__( |
| 76 | + self, |
| 77 | + *fields: str, |
| 78 | + filter_by: dict[str, str] | None = None, |
| 79 | + exclude: dict[str, list[str]] | None = None, |
| 80 | + reference: dict[str, str] | None = None, |
| 81 | + ): |
| 82 | + object.__setattr__(self, "fields", fields) |
| 83 | + object.__setattr__(self, "filter_by", filter_by or {}) |
| 84 | + object.__setattr__(self, "exclude", exclude or {}) |
| 85 | + object.__setattr__(self, "reference", reference or {}) |
| 86 | + |
| 87 | + |
| 88 | +@dataclass(frozen=True) |
| 89 | +class Meas: |
| 90 | + """A raw measurement value from element data. |
| 91 | +
|
| 92 | + Args: |
| 93 | + name: Feature name for this measurement in the output document. |
| 94 | + field: Key in the element dict containing the value. |
| 95 | + required: If True, a missing measurement causes the parent calc to be skipped. |
| 96 | + """ |
| 97 | + |
| 98 | + name: str |
| 99 | + field: str |
| 100 | + required: bool = False |
| 101 | + |
| 102 | + |
| 103 | +@dataclass(frozen=True) |
| 104 | +class Calc: |
| 105 | + """A calculated value derived from measurements or other calculations. |
| 106 | +
|
| 107 | + Args: |
| 108 | + name: Name of the output calculated document. |
| 109 | + field: Key in the element dict containing the pre-computed value. |
| 110 | + sources: Direct references to Meas or Calc nodes this depends on. |
| 111 | + group: How to partition elements for this calculation. |
| 112 | + unit: Unit string for the output value. |
| 113 | + description: Static description for the output document. |
| 114 | + description_field: Key in element dict for a dynamic description value. |
| 115 | + required: If True, a missing value causes the parent calc to be skipped. |
| 116 | + optional: If True, when this calc can't resolve, fall through to its sources. |
| 117 | + source_only: If True, this node is only a dependency — not emitted top-level. |
| 118 | + """ |
| 119 | + |
| 120 | + name: str |
| 121 | + field: str |
| 122 | + sources: list[Meas | Calc] = dataclass_field(default_factory=list) |
| 123 | + group: GroupBy = dataclass_field(default_factory=GroupBy) |
| 124 | + unit: str | None = None |
| 125 | + description: str | None = None |
| 126 | + description_field: str | None = None |
| 127 | + required: bool = False |
| 128 | + optional: bool = False |
| 129 | + source_only: bool = False |
| 130 | + |
| 131 | + |
| 132 | +CalcNode = Meas | Calc |
| 133 | + |
| 134 | + |
| 135 | +def calc_docs( |
| 136 | + items: list[T], |
| 137 | + to_element: Callable[[T], dict[str, Any]], |
| 138 | + nodes: list[CalcNode], |
| 139 | +) -> list[CalculatedDocument]: |
| 140 | + """Build calculated data documents from domain objects. |
| 141 | +
|
| 142 | + This is the single entry point. It: |
| 143 | + 1. Converts your domain objects to elements using to_element |
| 144 | + 2. Automatically builds view hierarchies from GroupBy specs |
| 145 | + 3. Resolves the dependency graph and produces CalculatedDocuments |
| 146 | +
|
| 147 | + Args: |
| 148 | + items: Your domain objects (WellItem, Measurement, etc.) |
| 149 | + to_element: Function that converts one item to a flat dict. |
| 150 | + Must include a "uuid" key for element identity. |
| 151 | + nodes: List of Meas and Calc nodes defining the dependency graph. |
| 152 | +
|
| 153 | + Returns: |
| 154 | + Flattened list of CalculatedDocuments ready for the schema mapper. |
| 155 | + """ |
| 156 | + elements = _build_elements(items, to_element) |
| 157 | + views = _build_views(nodes, elements) |
| 158 | + |
| 159 | + _validate(nodes) |
| 160 | + |
| 161 | + measurements: dict[int, MeasurementConfig] = {} |
| 162 | + calc_configs: dict[int, CalculatedDataConfig] = {} |
| 163 | + |
| 164 | + for node in nodes: |
| 165 | + if isinstance(node, Meas): |
| 166 | + measurements[id(node)] = MeasurementConfig( |
| 167 | + name=node.name, |
| 168 | + value=node.field, |
| 169 | + required=node.required, |
| 170 | + ) |
| 171 | + |
| 172 | + for node in nodes: |
| 173 | + if isinstance(node, Calc): |
| 174 | + _resolve(node, views, measurements, calc_configs) |
| 175 | + |
| 176 | + top_level_configs = [ |
| 177 | + calc_configs[id(node)] |
| 178 | + for node in nodes |
| 179 | + if isinstance(node, Calc) and id(node) in calc_configs and not node.source_only |
| 180 | + ] |
| 181 | + |
| 182 | + result = CalcDocsConfig(top_level_configs).construct() |
| 183 | + return list(chain.from_iterable(doc.iter_struct() for doc in result)) |
| 184 | + |
| 185 | + |
| 186 | +def _build_elements( |
| 187 | + items: list[T], to_element: Callable[[T], dict[str, Any]] |
| 188 | +) -> list[Element]: |
| 189 | + elements = [] |
| 190 | + for item in items: |
| 191 | + data = to_element(item) |
| 192 | + uuid = data.pop("uuid", None) |
| 193 | + if uuid is None: |
| 194 | + msg = "to_element must return a dict with a 'uuid' key" |
| 195 | + raise ValueError(msg) |
| 196 | + elements.append(Element(uuid=str(uuid), data=data)) |
| 197 | + return elements |
| 198 | + |
| 199 | + |
| 200 | +def _build_views(nodes: list[CalcNode], elements: list[Element]) -> dict[int, ViewData]: |
| 201 | + views: dict[int, ViewData] = {} |
| 202 | + for node in nodes: |
| 203 | + if isinstance(node, Calc): |
| 204 | + view = _group_by_to_view(node.group) |
| 205 | + views[id(node)] = view.apply(elements) |
| 206 | + return views |
| 207 | + |
| 208 | + |
| 209 | +def _group_by_to_view(group: GroupBy) -> View: |
| 210 | + """Convert a GroupBy spec into a View hierarchy.""" |
| 211 | + fields = list(group.fields) |
| 212 | + if not fields: |
| 213 | + msg = "GroupBy must specify at least one field" |
| 214 | + raise ValueError(msg) |
| 215 | + |
| 216 | + # Build view chain from innermost to outermost |
| 217 | + view: View | None = None |
| 218 | + for field_name in reversed(fields): |
| 219 | + view = _GroupByFieldView( |
| 220 | + field=field_name, |
| 221 | + sub_view=view, |
| 222 | + filter_spec=group.filter_by, |
| 223 | + exclude_spec=group.exclude, |
| 224 | + reference_spec=group.reference, |
| 225 | + ) |
| 226 | + if view is None: |
| 227 | + msg = "GroupBy must specify at least one field" |
| 228 | + raise ValueError(msg) |
| 229 | + return view |
| 230 | + |
| 231 | + |
| 232 | +class _GroupByFieldView(View): |
| 233 | + """Internal view that implements GroupBy semantics.""" |
| 234 | + |
| 235 | + def __init__( |
| 236 | + self, |
| 237 | + field: str, |
| 238 | + sub_view: View | None = None, |
| 239 | + filter_spec: dict[str, str] | None = None, |
| 240 | + exclude_spec: dict[str, list[str]] | None = None, |
| 241 | + reference_spec: dict[str, str] | None = None, |
| 242 | + ): |
| 243 | + super().__init__(name=field, sub_view=sub_view) |
| 244 | + self.field = field |
| 245 | + self.filter_spec = filter_spec or {} |
| 246 | + self.exclude_spec = exclude_spec or {} |
| 247 | + self.reference_spec = reference_spec or {} |
| 248 | + |
| 249 | + def sort_elements(self, elements: list[Element]) -> dict[str, list[Element]]: |
| 250 | + items: dict[str, list[Element]] = defaultdict(list) |
| 251 | + for element in elements: |
| 252 | + if not self._passes_filters(element): |
| 253 | + continue |
| 254 | + if self.field == "uuid": |
| 255 | + key: str | None = element.uuid |
| 256 | + else: |
| 257 | + key = element.get_str_or_none(self.field) |
| 258 | + if key is not None and not self._is_excluded(key): |
| 259 | + items[key].append(element) |
| 260 | + return dict(items) |
| 261 | + |
| 262 | + def _passes_filters(self, element: Element) -> bool: |
| 263 | + for field_name, required_value in self.filter_spec.items(): |
| 264 | + actual = element.get_str_or_none(field_name) |
| 265 | + if actual != required_value: |
| 266 | + return False |
| 267 | + return True |
| 268 | + |
| 269 | + def _is_excluded(self, key: str) -> bool: |
| 270 | + if self.field in self.exclude_spec: |
| 271 | + return key in self.exclude_spec[self.field] |
| 272 | + return False |
| 273 | + |
| 274 | + def filter_keys(self, keys: Keys) -> Keys: |
| 275 | + filtered_keys = self.sub_view.filter_keys(keys) if self.sub_view else Keys() |
| 276 | + if self.field in self.reference_spec: |
| 277 | + ref_value = self.reference_spec[self.field] |
| 278 | + if keys.get_or_none(self.name): |
| 279 | + return filtered_keys.overwrite(self.name, ref_value) |
| 280 | + return filtered_keys |
| 281 | + if key := keys.get_or_none(self.name): |
| 282 | + return filtered_keys.overwrite(self.name, key.value) |
| 283 | + return filtered_keys |
| 284 | + |
| 285 | + |
| 286 | +def _validate(nodes: list[CalcNode]) -> None: |
| 287 | + node_set = {id(n) for n in nodes} |
| 288 | + for node in nodes: |
| 289 | + if not isinstance(node, Calc): |
| 290 | + continue |
| 291 | + for source in node.sources: |
| 292 | + if id(source) not in node_set: |
| 293 | + msg = ( |
| 294 | + f"Calc '{node.name}' references source '{source.name}' " |
| 295 | + f"which is not in the nodes list" |
| 296 | + ) |
| 297 | + raise ValueError(msg) |
| 298 | + |
| 299 | + |
| 300 | +def _resolve( |
| 301 | + node: Calc, |
| 302 | + views: dict[int, ViewData], |
| 303 | + measurements: dict[int, MeasurementConfig], |
| 304 | + calc_configs: dict[int, CalculatedDataConfig], |
| 305 | +) -> CalculatedDataConfig: |
| 306 | + if id(node) in calc_configs: |
| 307 | + return calc_configs[id(node)] |
| 308 | + |
| 309 | + view_data = views[id(node)] |
| 310 | + |
| 311 | + source_configs: list[CalculatedDataConfig | MeasurementConfig] = [] |
| 312 | + for source in node.sources: |
| 313 | + if isinstance(source, Meas): |
| 314 | + source_configs.append(measurements[id(source)]) |
| 315 | + else: |
| 316 | + source_config = _resolve(source, views, measurements, calc_configs) |
| 317 | + source_configs.append(source_config) |
| 318 | + |
| 319 | + config: CalculatedDataConfig |
| 320 | + if node.optional: |
| 321 | + config = CalculatedDataConfigWithOptional( |
| 322 | + name=node.name, |
| 323 | + value=node.field, |
| 324 | + view_data=view_data, |
| 325 | + source_configs=tuple(source_configs), |
| 326 | + unit=node.unit, |
| 327 | + description=node.description, |
| 328 | + description_value_key=node.description_field, |
| 329 | + required=node.required, |
| 330 | + optional=True, |
| 331 | + ) |
| 332 | + else: |
| 333 | + config = CalculatedDataConfig( |
| 334 | + name=node.name, |
| 335 | + value=node.field, |
| 336 | + view_data=view_data, |
| 337 | + source_configs=tuple(source_configs), |
| 338 | + unit=node.unit, |
| 339 | + description=node.description, |
| 340 | + description_value_key=node.description_field, |
| 341 | + required=node.required, |
| 342 | + ) |
| 343 | + calc_configs[id(node)] = config |
| 344 | + return config |
0 commit comments