Source code for cfdmod.core.data_source

"""Frozen :class:`DataSource` value object plus the five concrete kinds.

The :class:`DataSource` is the unit of input and output for every op.
It binds together:

- a :class:`TimeAxis` (affine, never materialised);
- a :class:`Topology` (or ``None`` for purely tabular sources);
- an :class:`ElementMeta` (per-element scalars / vectors);
- a dict of :class:`Grouping` over the element axis;
- a :class:`FieldStore` (the small-vs-large-data seam) plus per-field
  :class:`FieldMeta`;
- a free-form ``attrs`` dict for source-level metadata.

Every method returns a *new* :class:`DataSource`. Field arrays inside
the underlying :class:`FieldStore` are shared by reference unless an
op explicitly rewrites them; large datasets do not duplicate.

Five concrete kinds are exposed; they are thin subclasses that lock
the ``kind`` discriminator and constrain which :class:`Topology`
``cell_type`` is admissible.
"""

from __future__ import annotations

__all__ = [
    "DataSource",
    "SurfaceDataSource",
    "VolumeDataSource",
    "PointsDataSource",
    "GroupsDataSource",
    "ModesDataSource",
]

from typing import Any, Literal

import numpy as np
from pydantic import BaseModel, ConfigDict, Field, model_validator

from cfdmod.core.field_meta import FieldMeta
from cfdmod.core.grouping import Grouping
from cfdmod.core.protocols import FieldStore
from cfdmod.core.time_axis import TimeAxis
from cfdmod.core.topology import ElementMeta, Topology

DataSourceKind = Literal["surface", "volume", "points", "groups", "modes"]


[docs] class DataSource(BaseModel): """Base frozen value object. Subclasses lock :attr:`kind` and (optionally) the admissible :attr:`Topology.cell_type`. Methods on this base never mutate; they always return a new instance via ``model_copy(update=...)``. Attributes: kind: One of ``surface``, ``volume``, ``points``, ``groups``, ``modes``. Locked by each subclass. time: Affine time axis. Time-aggregated outputs use ``n_timesteps == 0``. topology: Mesh connectivity / coordinates, when applicable. ``None`` is permitted for some kinds (notably ``modes``). elements: Per-element scalar / vector attributes. groupings: Mapping of grouping name -> :class:`Grouping`. fields: A :class:`FieldStore`. Carries the heavy arrays. field_meta: Mapping of field name -> :class:`FieldMeta`. attrs: Free-form source-level metadata. """ model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True) kind: DataSourceKind time: TimeAxis topology: Topology | None elements: ElementMeta groupings: dict[str, Grouping] = Field(default_factory=dict) fields: FieldStore field_meta: dict[str, FieldMeta] = Field(default_factory=dict) attrs: dict[str, Any] = Field(default_factory=dict) @model_validator(mode="after") def _check_consistency(self) -> "DataSource": n = self.n_elements for col in ("position", "area", "volume", "normal"): arr = getattr(self.elements, col) if arr is not None and arr.shape[0] != n: raise ValueError( f"elements.{col} length {arr.shape[0]} does not match n_elements={n}" ) for gname, grouping in self.groupings.items(): if grouping.n_elements != n: raise ValueError( f"grouping {gname!r} has {grouping.n_elements} entries; expected {n}" ) for fname in self.fields.keys(): shape = self.fields.shape(fname) if not shape or shape[0] != n: raise ValueError( f"field {fname!r} leading axis {shape[0] if shape else None} " f"does not match n_elements={n}" ) t_axis_len = shape[1] if len(shape) > 1 else 0 if self.time.is_time_aggregated and t_axis_len != 0: raise ValueError( f"field {fname!r} has a time axis but the data source's " "time axis is time-aggregated" ) if not self.time.is_time_aggregated and t_axis_len != self.time.n_timesteps: raise ValueError( f"field {fname!r} time axis length {t_axis_len} does not " f"match TimeAxis.n_timesteps={self.time.n_timesteps}" ) return self @property def n_elements(self) -> int: """Number of elements (rows) on this data source.""" if self.topology is not None: return self.topology.n_elements if self.elements.position is not None: return int(self.elements.position.shape[0]) # Fall back to whichever scalar column is present. for col in ("area", "volume"): arr = getattr(self.elements, col) if arr is not None: return int(arr.shape[0]) # Nothing else available; ask the field store. for name in self.fields.keys(): return int(self.fields.shape(name)[0]) return 0 @property def field_names(self) -> list[str]: return list(self.fields.keys()) # ----- Functional updates ------------------------------------------------- def _copy_validated(self, **update: Any) -> "DataSource": """``model_copy`` + re-run the consistency validators. Pydantic's ``model_copy(update=...)`` does *not* re-run validators, so a functional update could otherwise build a shape-inconsistent (frozen) DataSource silently. Re-validating keeps the invariants the frozen model advertises true after every update, not just at first construction. """ updated = self.model_copy(update=update) return type(self).model_validate(dict(updated.__dict__))
[docs] def with_time(self, new_time: TimeAxis) -> "DataSource": """Return a copy with a new :class:`TimeAxis`. Field shapes must already match the new axis -- this is a metadata update only.""" return self._copy_validated(time=new_time)
def with_topology(self, new_topology: Topology) -> "DataSource": return self._copy_validated(topology=new_topology) def with_elements(self, new_elements: ElementMeta) -> "DataSource": return self._copy_validated(elements=new_elements)
[docs] def with_grouping(self, grouping: Grouping) -> "DataSource": """Add or replace a grouping. The grouping name is the key.""" new_groupings = dict(self.groupings) new_groupings[grouping.name] = grouping return self._copy_validated(groupings=new_groupings)
def without_grouping(self, name: str) -> "DataSource": new_groupings = {k: v for k, v in self.groupings.items() if k != name} return self.model_copy(update={"groupings": new_groupings})
[docs] def with_field( self, name: str, value: np.ndarray, meta: FieldMeta | None = None, ) -> "DataSource": """Add or replace a field. The :class:`FieldStore` decides whether the array is shared by reference or copied.""" new_store = self.fields.with_field(name, value) new_meta = dict(self.field_meta) new_meta[name] = meta or FieldMeta(name=name) return self._copy_validated(fields=new_store, field_meta=new_meta)
def with_attrs(self, **updates: Any) -> "DataSource": new_attrs = dict(self.attrs) new_attrs.update(updates) return self.model_copy(update={"attrs": new_attrs})
# --------------------------------------------------------------------------- # Concrete kinds # ---------------------------------------------------------------------------
[docs] class SurfaceDataSource(DataSource): """Faces (2D triangular cells) with optional timesteps. Topology cell type must be ``triangle``. Mirrors the existing cfdmod XDMF+H5 timeseries layout: ``/Triangles``, ``/Geometry``, ``/{group}/t{T}``. """ kind: Literal["surface"] = "surface" @model_validator(mode="after") def _check_surface(self) -> "SurfaceDataSource": if self.topology is None or self.topology.cell_type != "triangle": raise ValueError( "SurfaceDataSource requires a triangle Topology; got " f"{None if self.topology is None else self.topology.cell_type!r}" ) return self
[docs] class VolumeDataSource(DataSource): """3D cells with optional timesteps. Topology cell type must be ``cell``. Reserved -- not a Phase 1 target. The class is here so volume export can be added later additively rather than as a schema change. """ kind: Literal["volume"] = "volume" @model_validator(mode="after") def _check_volume(self) -> "VolumeDataSource": if self.topology is None or self.topology.cell_type != "cell": raise ValueError( "VolumeDataSource requires a cell Topology; got " f"{None if self.topology is None else self.topology.cell_type!r}" ) return self
[docs] class PointsDataSource(DataSource): """Bare points / probes / vertical profiles. Covers the existing :class:`InflowData` (probe array + per-component timeseries) and :class:`s1.profile.Profile` (1-D vertical profile, no time axis). Topology cell type is ``point``; connectivity is empty. """ kind: Literal["points"] = "points" @model_validator(mode="after") def _check_points(self) -> "PointsDataSource": if self.topology is None or self.topology.cell_type != "point": raise ValueError( "PointsDataSource requires a point Topology; got " f"{None if self.topology is None else self.topology.cell_type!r}" ) return self
[docs] class GroupsDataSource(DataSource): """One row per group: an aggregation over a parent surface. A groups data source carries fields whose leading axis is the *group* index, not the original element index. Its topology is *chained*: it borrows the parent surface's :class:`Topology` plus a :class:`Grouping` mapping each parent element to a group. The class does not own a triangulation of the groups themselves (each group is in general not a single triangle). This avoids the "non-triangular faces" trap. Attributes: parent_topology: The parent surface's triangle topology. parent_grouping: A :class:`Grouping` over the parent surface's elements that determines membership. """ kind: Literal["groups"] = "groups" parent_topology: Topology parent_grouping: Grouping @model_validator(mode="after") def _check_groups(self) -> "GroupsDataSource": if self.parent_topology.cell_type != "triangle": raise ValueError( "GroupsDataSource.parent_topology must be a triangle topology; got " f"{self.parent_topology.cell_type!r}" ) if self.parent_grouping.n_elements != self.parent_topology.n_elements: raise ValueError( "GroupsDataSource.parent_grouping must have one entry per parent element" ) # GroupsDataSource has no independent topology of its own. if self.topology is not None: raise ValueError( "GroupsDataSource must not carry an independent topology; topology is " "chained to parent_topology + parent_grouping" ) return self
[docs] class ModesDataSource(DataSource): """Modal axis: one row per mode, fields are generalised-displacement timeseries. No spatial topology; the original mesh / structural data lives alongside in the recipe context. ``elements`` typically carries an annotation column with mode labels. """ kind: Literal["modes"] = "modes" @model_validator(mode="after") def _check_modes(self) -> "ModesDataSource": if self.topology is not None: raise ValueError( f"ModesDataSource does not carry a topology; got {self.topology.cell_type!r}" ) return self