tensogram-anemoi 0.17.0__tar.gz

tensogram_anemoi-0.17.0/.gitignore

@@ -0,0 +1,22 @@
+.claude/*
+!.claude/commands/
+.weave/
+.sisyphus/
+.coverage
+*.dylib.dSYM/
+**/target
+**/pkg
+**/build/
+docs/book/
+python/**/dist/
+python/bindings/Cargo.lock
+rust/tensogram-grib/Cargo.lock
+rust/tensogram-netcdf/Cargo.lock
+rust/tensogram-wasm/Cargo.lock
+# Python virtualenv, caches, and maturin-installed extension modules
+.venv/
+**/__pycache__/
+*.pyc
+python/bindings/python/tensogram/tensogram*.so
+# TODO do we want to have uv.locks ignored?
+**/uv.lock

tensogram_anemoi-0.17.0/PKG-INFO

@@ -0,0 +1,24 @@
+Metadata-Version: 2.4
+Name: tensogram-anemoi
+Version: 0.17.0
+Summary: anemoi-inference output plugin for tensogram .tgm files
+Project-URL: Homepage, https://sites.ecmwf.int/docs/tensogram/main
+Project-URL: Repository, https://github.com/ecmwf/tensogram
+Project-URL: Documentation, https://sites.ecmwf.int/docs/tensogram/main
+Author-email: ECMWF <software@ecmwf.int>
+License-Expression: Apache-2.0
+Classifier: Development Status :: 4 - Beta
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Atmospheric Science
+Requires-Python: >=3.11
+Requires-Dist: anemoi-inference<0.11,>=0.4
+Requires-Dist: fsspec
+Requires-Dist: numpy
+Requires-Dist: tensogram<0.18,>=0.17.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == 'dev'
+Requires-Dist: ruff>=0.4; extra == 'dev'
+Description-Content-Type: text/plain
+
+anemoi-inference output plugin for tensogram .tgm files

tensogram_anemoi-0.17.0/pyproject.toml

@@ -0,0 +1,62 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "tensogram-anemoi"
+version = "0.17.0"
+description = "anemoi-inference output plugin for tensogram .tgm files"
+readme = {text = "anemoi-inference output plugin for tensogram .tgm files", content-type = "text/plain"}
+requires-python = ">=3.11"
+license = "Apache-2.0"
+authors = [{name = "ECMWF", email = "software@ecmwf.int"}]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering",
+    "Topic :: Scientific/Engineering :: Atmospheric Science",
+]
+dependencies = [
+    "tensogram>=0.17.0,<0.18",
+    "anemoi-inference>=0.4,<0.11",
+    "fsspec",
+    "numpy",
+]
+
+[project.urls]
+Homepage = "https://sites.ecmwf.int/docs/tensogram/main"
+Repository = "https://github.com/ecmwf/tensogram"
+Documentation = "https://sites.ecmwf.int/docs/tensogram/main"
+
+[project.optional-dependencies]
+dev = ["pytest>=7.0", "ruff>=0.4"]
+
+[project.entry-points."anemoi.inference.outputs"]
+tensogram = "tensogram_anemoi.output:TensogramOutput"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/tensogram_anemoi"]
+
+[tool.ruff]
+line-length = 99
+target-version = "py311"
+
+[tool.ruff.lint]
+select = [
+    "E",
+    "W",
+    "F",
+    "I",
+    "N",
+    "UP",
+    "B",
+    "SIM",
+    "PT",
+    "RUF",
+]
+
+[tool.ruff.lint.isort]
+known-third-party = ["numpy", "tensogram", "anemoi", "fsspec", "pytest"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

tensogram_anemoi-0.17.0/src/tensogram_anemoi/__init__.py

@@ -0,0 +1,4 @@
+# (C) Copyright 2025 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.

tensogram_anemoi-0.17.0/src/tensogram_anemoi/output.py

@@ -0,0 +1,336 @@
+# (C) Copyright 2025 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+
+import logging
+from functools import cached_property
+from pathlib import Path
+
+import numpy as np
+
+from anemoi.inference.context import Context
+from anemoi.inference.decorators import main_argument
+from anemoi.inference.output import Output
+from anemoi.inference.types import ProcessorConfig
+from anemoi.inference.types import State
+
+LOG = logging.getLogger(__name__)
+
+# Written into base[i]["name"] for lat/lon objects.  These names are
+# deliberately NOT in tensogram-xarray's KNOWN_COORD_NAMES so that lat/lon
+# objects share the same flat dimension as all field objects rather than each
+# spawning its own dimension.  The canonical names are preserved in the "anemoi"
+# namespace for downstream consumers.
+_COORD_NAME_MAP = {
+    "latitude": "grid_latitude",
+    "longitude": "grid_longitude",
+}
+
+
+@main_argument("path")
+class TensogramOutput(Output):
+    """Tensogram output plugin for anemoi-inference.
+
+    Writes each forecast step as one tensogram message to a .tgm file.
+    Each message contains lat/lon coordinate
+    objects followed by one object per field (or one stacked object per
+    pressure-level parameter when ``stack_pressure_levels=True``).
+
+    Per-object metadata is stored under the ``"anemoi"`` namespace in CBOR.
+    Dimension-name hints are written into each base entry's ``dim_names``
+    key (per-object, axis-ordered) *and* into ``_extra_["dim_names"]``
+    (message-level size-to-name dict) so readers on either convention can
+    resolve meaningful names without callers passing ``dim_names`` explicitly.
+
+    Supports local paths and remote URLs (s3://, gs://, az://, ...) via fsspec.
+    Each step is encoded and written immediately -- no full-forecast buffering.
+
+    Pressure-level stacking
+    -----------------------
+    When ``stack_pressure_levels=True``, all fields sharing the same ``param``
+    are stacked into a single 2-D object of shape ``(n_grid, n_levels)``,
+    sorted by level in ascending order.  The per-object metadata stores
+    ``"levelist": [500, 850, ...]`` in the ``"mars"`` namespace instead of
+    the scalar ``"level"`` key used for unstacked fields.
+
+    Without stacking (default), every field is a separate 1-D object and the
+    scalar ``"level"`` key is always stored when it is present in the
+    checkpoint's GRIB keys.
+    """
+
+    def __init__(
+        self,
+        context: Context,
+        path: str,
+        *,
+        encoding: str = "none",
+        bits: int | None = None,
+        compression: str = "zstd",
+        dtype: str = "float32",
+        storage_options: dict | None = None,
+        stack_pressure_levels: bool = False,
+        variables: list[str] | None = None,
+        post_processors: list[ProcessorConfig] | None = None,
+        output_frequency: int | None = None,
+        write_initial_state: bool | None = None,
+    ) -> None:
+        super().__init__(
+            context,
+            variables=variables,
+            post_processors=post_processors,
+            output_frequency=output_frequency,
+            write_initial_state=write_initial_state,
+        )
+        if encoding == "simple_packing" and bits is None:
+            raise ValueError("bits must be set when encoding='simple_packing'")
+        self.path = path
+        self.encoding = encoding
+        self.bits = bits
+        self.compression = compression
+        self.dtype = dtype
+        self.storage_options = storage_options or {}
+        self.stack_pressure_levels = stack_pressure_levels
+        self._handle = None
+
+    def __repr__(self) -> str:
+        return f"TensogramOutput({self.path})"
+
+    @cached_property
+    def _numpy_dtype(self) -> np.dtype:
+        return np.dtype(self.dtype)
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def open(self, state: State) -> None:
+        import fsspec
+
+        path_str = str(self.path)
+        if "://" not in path_str:
+            Path(path_str).parent.mkdir(parents=True, exist_ok=True)
+
+        self._handle = fsspec.open(path_str, "wb", **self.storage_options).open()
+        LOG.info("TensogramOutput: writing to %s", path_str)
+
+    def write_initial_state(self, state: State) -> None:
+        """Write the initial state, reducing multi-step fields to the last step."""
+        from anemoi.inference.state import reduce_state
+
+        state = reduce_state(state)
+        return super().write_initial_state(state)
+
+    def write_step(self, state: State) -> None:
+        """Encode one forecast step as a tensogram message and write it immediately."""
+        if self._handle is None:
+            raise RuntimeError(f"{self!r}: write_step called before open() or after close()")
+
+        global_meta = {
+            "version": 3,
+            "base": [],
+            "_extra_": {},
+        }
+        descriptors_and_data = []
+
+        step_seconds = state["step"].total_seconds()
+        step_hours = int(step_seconds / 3600) if step_seconds % 3600 == 0 else step_seconds / 3600
+        base_dt = state["date"] - state["step"]
+        mars_extra = {
+            "date": base_dt.strftime("%Y%m%d"),
+            "time": base_dt.strftime("%H%M"),
+            "step": step_hours,
+        }
+
+        for coord_name, coord_arr in [
+            ("latitude", state["latitudes"]),
+            ("longitude", state["longitudes"]),
+        ]:
+            arr = np.asarray(coord_arr, dtype=np.float64)
+            global_meta["base"].append(
+                {
+                    "name": _COORD_NAME_MAP[coord_name],
+                    "anemoi": {"variable": coord_name},
+                    "dim_names": [coord_name],
+                }
+            )
+            descriptors_and_data.append(
+                (
+                    {"type": "ntensor", "shape": list(arr.shape), "dtype": "float64"},
+                    arr,
+                )
+            )
+
+        if self.stack_pressure_levels:
+            self._add_fields_stacked(state, global_meta, descriptors_and_data, mars_extra)
+        else:
+            self._add_fields_flat(state, global_meta, descriptors_and_data, mars_extra)
+
+        n_grid = len(state["latitudes"])
+        dim_names_hint: dict[str, str] = {str(n_grid): "values"}
+        if self.stack_pressure_levels:
+            for _, arr in descriptors_and_data:
+                if arr.ndim == 2:
+                    level_size = str(arr.shape[1])
+                    if level_size not in dim_names_hint:
+                        dim_names_hint[level_size] = "level"
+        global_meta["_extra_"]["dim_names"] = dim_names_hint
+
+        import tensogram
+
+        msg_bytes = tensogram.encode(global_meta, descriptors_and_data)
+        self._handle.write(msg_bytes)
+
+    def close(self) -> None:
+        """Flush and close the output stream."""
+        if self._handle is not None:
+            try:
+                self._handle.flush()
+            except Exception:
+                LOG.warning(
+                    "TensogramOutput: failed to flush output stream before close", exc_info=True
+                )
+            self._handle.close()
+            self._handle = None
+
+    # ------------------------------------------------------------------
+    # Field object builders
+    # ------------------------------------------------------------------
+
+    def _add_fields_flat(
+        self,
+        state: State,
+        global_meta: dict,
+        descriptors_and_data: list,
+        mars_extra: dict,
+    ) -> None:
+        """Add one object per field (default, no stacking)."""
+        for name, values in state["fields"].items():
+            if self.skip_variable(name):
+                continue
+            variable = self.typed_variables.get(name)
+            if variable is None:
+                LOG.warning(
+                    "TensogramOutput: no typed variable for %r -- metadata will be incomplete",
+                    name,
+                )
+            grib = getattr(variable, "grib_keys", {}) if variable else {}
+            base_entry, descriptor, arr = self._build_field_object(name, grib, values, mars_extra)
+            global_meta["base"].append(base_entry)
+            descriptors_and_data.append((descriptor, arr))
+
+    def _add_fields_stacked(
+        self,
+        state: State,
+        global_meta: dict,
+        descriptors_and_data: list,
+        mars_extra: dict,
+    ) -> None:
+        """Group pressure-level fields by param and stack; write others flat."""
+        pl_groups: dict[str, list[tuple[int, str, dict, np.ndarray]]] = {}
+        non_pl: list[tuple[str, dict, np.ndarray]] = []
+
+        for name, values in state["fields"].items():
+            if self.skip_variable(name):
+                continue
+            variable = self.typed_variables.get(name)
+            if variable is None:
+                LOG.warning(
+                    "TensogramOutput: no typed variable for %r -- metadata will be incomplete",
+                    name,
+                )
+            grib = getattr(variable, "grib_keys", {}) if variable else {}
+            if variable is not None and variable.is_pressure_level:
+                param = variable.param
+                level = variable.level
+                pl_groups.setdefault(param, []).append((level, name, grib, values))
+            else:
+                non_pl.append((name, grib, values))
+
+        if not pl_groups:
+            LOG.warning(
+                "TensogramOutput: stack_pressure_levels=True but no pressure-level fields found"
+            )
+
+        for name, grib, values in non_pl:
+            base_entry, descriptor, arr = self._build_field_object(name, grib, values, mars_extra)
+            global_meta["base"].append(base_entry)
+            descriptors_and_data.append((descriptor, arr))
+
+        for param in sorted(pl_groups):
+            group = sorted(pl_groups[param], key=lambda x: x[0])
+            base_entry, descriptor, arr = self._build_stacked_object(param, group, mars_extra)
+            global_meta["base"].append(base_entry)
+            descriptors_and_data.append((descriptor, arr))
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _prepare_array(self, values: np.ndarray) -> np.ndarray:
+        arr = np.asarray(values, dtype=self._numpy_dtype)
+        if self.encoding == "simple_packing":
+            arr = arr.astype(np.float64)
+        return arr
+
+    def _build_descriptor(self, arr: np.ndarray) -> dict:
+        descriptor = {
+            "type": "ntensor",
+            "shape": list(arr.shape),
+            "dtype": arr.dtype.name,
+            "encoding": self.encoding,
+            "compression": self.compression,
+        }
+        if self.encoding == "simple_packing" and self.bits is not None:
+            import tensogram
+
+            sp_params = tensogram.compute_packing_params(arr.ravel(), self.bits, 0)
+            descriptor.update(sp_params)
+        return descriptor
+
+    def _build_field_object(
+        self,
+        name: str,
+        grib: dict,
+        values: np.ndarray,
+        mars_extra: dict,
+    ) -> tuple[dict, dict, np.ndarray]:
+        """Build (base_entry, descriptor, array) for a single flat field object."""
+        mars = {**mars_extra, **grib}
+        base_entry: dict = {
+            "name": name,
+            "anemoi": {"variable": name},
+            "dim_names": ["values"],
+        }
+        if mars:
+            base_entry["mars"] = mars
+        arr = self._prepare_array(values)
+        return base_entry, self._build_descriptor(arr), arr
+
+    def _build_stacked_object(
+        self,
+        param: str,
+        group: list[tuple[int, str, dict, np.ndarray]],
+        mars_extra: dict,
+    ) -> tuple[dict, dict, np.ndarray]:
+        """Build (base_entry, descriptor, array) for a stacked pressure-level object."""
+        levels = [item[0] for item in group]
+        first_grib = group[0][2]
+
+        arrays = [self._prepare_array(item[3]) for item in group]
+        stacked = np.column_stack(arrays)
+
+        mars = {
+            **mars_extra,
+            **{k: v for k, v in first_grib.items() if k != "level"},
+            "levelist": levels,
+        }
+
+        base_entry: dict = {
+            "name": param,
+            "anemoi": {"variable": param},
+            "dim_names": ["values", "level"],
+        }
+        if mars:
+            base_entry["mars"] = mars
+        return base_entry, self._build_descriptor(stacked), stacked

tensogram_anemoi-0.17.0/tests/test_tensogram_output.py

@@ -0,0 +1,571 @@
+# (C) Copyright 2025 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+
+"""Tests for TensogramOutput."""
+
+from datetime import datetime
+from datetime import timedelta
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+from tensogram_anemoi.output import TensogramOutput
+
+tensogram = pytest.importorskip("tensogram", reason="tensogram not installed")
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+N_GRID = 100
+FIELD_NAMES = ["2t", "10u", "10v"]
+
+
+def _make_variable(param=None, levtype=None, level=None):
+    grib = {"param": param} if param is not None else {}
+    if levtype is not None:
+        grib["levtype"] = levtype
+    if level is not None:
+        grib["level"] = level
+    is_pl = levtype is not None and level is not None
+    return SimpleNamespace(
+        grib_keys=grib,
+        is_computed_forcing=False,
+        is_pressure_level=is_pl,
+        param=param,
+        level=level,
+    )
+
+
+def _make_context(field_names=FIELD_NAMES):
+    typed_variables = {name: _make_variable(param=name) for name in field_names}
+
+    checkpoint = SimpleNamespace(
+        typed_variables=typed_variables,
+    )
+    context = SimpleNamespace(
+        checkpoint=checkpoint,
+        reference_date=datetime(2024, 1, 1),
+        write_initial_state=False,
+        output_frequency=None,
+        typed_variables={},
+    )
+    return context
+
+
+def _make_state(step_hours=1, field_names=FIELD_NAMES, n_grid=N_GRID, seed=0):
+    rng = np.random.default_rng(seed)
+    date = datetime(2024, 1, 1) + timedelta(hours=step_hours)
+    return {
+        "date": date,
+        "step": timedelta(hours=step_hours),
+        "latitudes": np.linspace(-90, 90, n_grid, dtype=np.float64),
+        "longitudes": np.linspace(0, 360, n_grid, dtype=np.float64),
+        "fields": {name: rng.random(n_grid).astype(np.float32) for name in field_names},
+    }
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+def test_write_step_before_open_raises(tmp_path):
+    """write_step raises RuntimeError when called before open()."""
+    context = _make_context()
+    output = TensogramOutput(context, str(tmp_path / "never.tgm"))
+    with pytest.raises(RuntimeError, match="open"):
+        output.write_step(_make_state())
+
+
+def test_write_initial_state_reduces_multistep(tmp_path):
+    """write_initial_state reduces a multi-step field array to its last step."""
+    from unittest.mock import patch
+
+    path = tmp_path / "init.tgm"
+    context = _make_context(["2t"])
+    output = TensogramOutput(context, str(path), write_initial_state=True)
+
+    rng = np.random.default_rng(1)
+    multi_step_values = rng.random((3, N_GRID)).astype(np.float32)
+    state = {
+        "date": datetime(2024, 1, 1),
+        "step": timedelta(0),
+        "latitudes": np.linspace(-90, 90, N_GRID, dtype=np.float64),
+        "longitudes": np.linspace(0, 360, N_GRID, dtype=np.float64),
+        "fields": {"2t": multi_step_values},
+    }
+
+    written_states = []
+
+    def capture_write_step(s):
+        written_states.append(s)
+
+    output.open(state)
+    with patch.object(output, "write_step", side_effect=capture_write_step):
+        output.write_initial_state(state)
+    output.close()
+
+    assert written_states, "write_step was not called by write_initial_state"
+    written_field = written_states[0]["fields"]["2t"]
+    np.testing.assert_array_equal(written_field, multi_step_values[-1])
+
+
+def test_write_and_read_local(tmp_path):
+    """Write 3 steps to a local .tgm file and verify round-trip."""
+    path = tmp_path / "forecast.tgm"
+    context = _make_context()
+    output = TensogramOutput(context, str(path))
+
+    states = [_make_state(h) for h in range(1, 4)]
+    output.open(states[0])
+    for state in states:
+        output.write_step(state)
+    output.close()
+
+    assert path.exists()
+
+    tgm_file = tensogram.TensogramFile.open(str(path))
+    assert len(tgm_file) == 3
+
+    for i, msg in enumerate(tgm_file):
+        meta, objects = msg
+
+        assert len(objects) == 5
+
+        assert meta.base[0]["name"] == "grid_latitude"
+        assert meta.base[1]["name"] == "grid_longitude"
+        assert meta.base[0]["anemoi"]["variable"] == "latitude"
+        assert meta.base[1]["anemoi"]["variable"] == "longitude"
+        lat_desc, lat_arr = objects[0]
+        lon_desc, lon_arr = objects[1]
+        assert lat_desc.dtype == "float64"
+        assert lon_desc.dtype == "float64"
+        np.testing.assert_allclose(lat_arr, states[i]["latitudes"])
+        np.testing.assert_allclose(lon_arr, states[i]["longitudes"])
+
+        expected_step = int(states[i]["step"].total_seconds() / 3600)
+        base_dt = states[i]["date"] - states[i]["step"]
+        expected_date = base_dt.strftime("%Y%m%d")
+        expected_time = base_dt.strftime("%H%M")
+        for j, name in enumerate(FIELD_NAMES):
+            obj_idx = 2 + j
+            assert meta.base[obj_idx]["name"] == name
+            assert meta.base[obj_idx]["anemoi"]["variable"] == name
+            assert meta.base[obj_idx]["mars"]["param"] == name
+            assert meta.base[obj_idx]["mars"]["step"] == expected_step
+            assert meta.base[obj_idx]["mars"]["date"] == expected_date
+            assert meta.base[obj_idx]["mars"]["time"] == expected_time
+            _, field_arr = objects[obj_idx]
+            np.testing.assert_allclose(field_arr, states[i]["fields"][name], rtol=1e-6)
+
+
+def test_variable_filter(tmp_path):
+    """Only the selected variable and coordinates are written."""
+    path = tmp_path / "filtered.tgm"
+    context = _make_context()
+    output = TensogramOutput(context, str(path), variables=["2t"])
+
+    state = _make_state()
+    output.open(state)
+    output.write_step(state)
+    output.close()
+
+    tgm_file = tensogram.TensogramFile.open(str(path))
+    msg = tgm_file[0]
+    meta, objects = msg
+
+    assert len(objects) == 3
+    assert meta.base[2]["anemoi"]["variable"] == "2t"
+
+
+def test_simple_packing_encoding(tmp_path):
+    """simple_packing round-trip stays within expected quantisation error."""
+    path = tmp_path / "packed.tgm"
+    context = _make_context(["2t"])
+    output = TensogramOutput(
+        context, str(path), encoding="simple_packing", bits=16, compression="zstd"
+    )
+
+    rng = np.random.default_rng(42)
+    values = (rng.random(N_GRID) * 50 + 250).astype(np.float32)
+    state = {
+        "date": datetime(2024, 1, 1, 1),
+        "step": timedelta(hours=1),
+        "latitudes": np.linspace(-90, 90, N_GRID, dtype=np.float64),
+        "longitudes": np.linspace(0, 360, N_GRID, dtype=np.float64),
+        "fields": {"2t": values},
+    }
+
+    output.open(state)
+    output.write_step(state)
+    output.close()
+
+    tgm_file = tensogram.TensogramFile.open(str(path))
+    meta, objects = tgm_file[0]
+    _, decoded = objects[2]
+
+    np.testing.assert_allclose(decoded, values, atol=0.002)
+
+
+def test_simple_packing_requires_bits(tmp_path):
+    """encoding='simple_packing' without bits raises ValueError immediately."""
+    context = _make_context(["2t"])
+    with pytest.raises(ValueError, match="bits must be set"):
+        TensogramOutput(context, str(tmp_path / "out.tgm"), encoding="simple_packing")
+
+
+def test_mars_metadata_forwarded(tmp_path):
+    """grib_keys appear in per-object 'mars' namespace, following tensogram-grib convention."""
+    path = tmp_path / "levs.tgm"
+    typed_variables = {
+        "t500": _make_variable(param="t", levtype="pl", level=500),
+    }
+    checkpoint = SimpleNamespace(typed_variables=typed_variables)
+    context = SimpleNamespace(
+        checkpoint=checkpoint,
+        reference_date=datetime(2024, 1, 1),
+        write_initial_state=False,
+        output_frequency=None,
+        typed_variables={},
+    )
+
+    output = TensogramOutput(context, str(path))
+    state = {
+        "date": datetime(2024, 1, 1, 6),
+        "step": timedelta(hours=6),
+        "latitudes": np.zeros(10, dtype=np.float64),
+        "longitudes": np.zeros(10, dtype=np.float64),
+        "fields": {"t500": np.ones(10, dtype=np.float32)},
+    }
+    output.open(state)
+    output.write_step(state)
+    output.close()
+
+    tgm_file = tensogram.TensogramFile.open(str(path))
+    meta, _ = tgm_file[0]
+    mars = meta.base[2]["mars"]
+    assert mars["param"] == "t"
+    assert mars["levtype"] == "pl"
+    assert mars["level"] == 500
+    assert mars["step"] == 6
+    assert mars["date"] == "20240101"
+    assert mars["time"] == "0000"
+    assert meta.base[2]["anemoi"]["variable"] == "t500"
+
+
+def test_remote_write_via_memory_fs():
+    """Write to fsspec memory filesystem and read back valid tensogram messages."""
+    import fsspec
+
+    url = "memory://test_forecast.tgm"
+    context = _make_context()
+    output = TensogramOutput(context, url)
+
+    states = [_make_state(h) for h in range(1, 3)]
+    output.open(states[0])
+    for state in states:
+        output.write_step(state)
+    output.close()
+
+    fs = fsspec.filesystem("memory")
+    raw = fs.open(url, "rb").read()
+
+    messages = tensogram.scan(raw)
+    assert len(messages) == 2
+
+    for i, (offset, length) in enumerate(messages):
+        msg_bytes = raw[offset : offset + length]
+        meta, objects = tensogram.decode(msg_bytes)
+        assert len(objects) == 5
+
+
+def test_dtype_float64(tmp_path):
+    """dtype=float64 stores field arrays as float64."""
+    path = tmp_path / "f64.tgm"
+    context = _make_context(["2t"])
+    output = TensogramOutput(context, str(path), dtype="float64")
+
+    state = _make_state(field_names=["2t"])
+    output.open(state)
+    output.write_step(state)
+    output.close()
+
+    tgm_file = tensogram.TensogramFile.open(str(path))
+    meta, objects = tgm_file[0]
+    desc, arr = objects[2]
+    assert desc.dtype == "float64"
+    assert arr.dtype == np.float64
+
+
+def test_close_is_idempotent(tmp_path):
+    """close() can be called multiple times without error."""
+    path = tmp_path / "idem.tgm"
+    context = _make_context()
+    output = TensogramOutput(context, str(path))
+    output.open(_make_state())
+    output.write_step(_make_state())
+    output.close()
+    output.close()
+
+
+def test_dim_names_in_metadata(tmp_path):
+    """Dim-name hints are written into both _extra_ (message-level) and
+    per-object base[i]['dim_names'] so xarray readers on either convention
+    resolve meaningful dim names without explicit kwargs."""
+    path = tmp_path / "dimnames.tgm"
+    context = _make_context()
+    output = TensogramOutput(context, str(path))
+    state = _make_state()
+    output.open(state)
+    output.write_step(state)
+    output.close()
+
+    tgm_file = tensogram.TensogramFile.open(str(path))
+    meta, _ = tgm_file[0]
+    dim_names = meta.extra["dim_names"]
+    assert str(N_GRID) in dim_names
+    assert dim_names[str(N_GRID)] == "values"
+
+    # Coord entries pin their axis to the canonical coord name.
+    coord_entries = {
+        entry.get("anemoi", {}).get("variable"): entry
+        for entry in meta.base
+        if entry.get("anemoi", {}).get("variable") in {"latitude", "longitude"}
+    }
+    assert coord_entries.get("latitude", {}).get("dim_names") == ["latitude"]
+    assert coord_entries.get("longitude", {}).get("dim_names") == ["longitude"]
+
+    # Flat field entries use the 'values' dim.
+    field_entries = [
+        entry
+        for entry in meta.base
+        if entry.get("anemoi", {}).get("variable") not in {"latitude", "longitude"}
+    ]
+    assert field_entries, "expected at least one field entry"
+    for entry in field_entries:
+        assert entry.get("dim_names") == ["values"], entry
+
+
+def test_stacked_dim_names_in_metadata(tmp_path):
+    """Stacked fields write size-based 'values'/'level' hints into _extra_
+    AND a ['values', 'level'] per-object list on each stacked base entry."""
+    path = tmp_path / "stacked_dims.tgm"
+    context = _make_pl_context(params=["t"], levels=[500, 850, 1000])
+    output = TensogramOutput(context, str(path), stack_pressure_levels=True)
+    state = _make_pl_state(params=["t"], levels=[500, 850, 1000])
+    output.open(state)
+    output.write_step(state)
+    output.close()
+
+    tgm_file = tensogram.TensogramFile.open(str(path))
+    meta, _ = tgm_file[0]
+    dim_names = meta.extra["dim_names"]
+    assert str(N_GRID) in dim_names
+    assert dim_names[str(N_GRID)] == "values"
+    assert str(3) in dim_names
+    assert dim_names[str(3)] == "level"
+
+    stacked_entries = [
+        entry for entry in meta.base if entry.get("anemoi", {}).get("variable") == "t"
+    ]
+    assert stacked_entries, "expected stacked pressure-level entry"
+    for entry in stacked_entries:
+        assert entry.get("dim_names") == ["values", "level"], entry
+
+
+# ---------------------------------------------------------------------------
+# Pressure-level stacking tests
+# ---------------------------------------------------------------------------
+
+PL_LEVELS = [500, 850, 1000]
+PL_PARAMS = ["t", "u"]
+
+
+def _make_pl_context(params=PL_PARAMS, levels=PL_LEVELS, extra_fields=None):
+    """Context with pressure-level variables plus optional non-PL extras."""
+    typed_variables = {}
+    for param in params:
+        for level in levels:
+            name = f"{param}{level}"
+            typed_variables[name] = _make_variable(param=param, levtype="pl", level=level)
+    for name in extra_fields or []:
+        typed_variables[name] = _make_variable(param=name)
+
+    checkpoint = SimpleNamespace(typed_variables=typed_variables)
+    return SimpleNamespace(
+        checkpoint=checkpoint,
+        reference_date=datetime(2024, 1, 1),
+        write_initial_state=False,
+        output_frequency=None,
+        typed_variables={},
+    )
+
+
+def _make_pl_state(params=PL_PARAMS, levels=PL_LEVELS, extra_fields=None, n_grid=N_GRID, seed=0):
+    rng = np.random.default_rng(seed)
+    fields = {}
+    for param in params:
+        for level in levels:
+            fields[f"{param}{level}"] = rng.random(n_grid).astype(np.float32)
+    for name in extra_fields or []:
+        fields[name] = rng.random(n_grid).astype(np.float32)
+    return {
+        "date": datetime(2024, 1, 1, 1),
+        "step": timedelta(hours=1),
+        "latitudes": np.linspace(-90, 90, n_grid, dtype=np.float64),
+        "longitudes": np.linspace(0, 360, n_grid, dtype=np.float64),
+        "fields": fields,
+    }
+
+
+def test_stack_object_count(tmp_path):
+    """With stacking: one object per param, not per level."""
+    path = tmp_path / "stacked.tgm"
+    context = _make_pl_context()
+    output = TensogramOutput(context, str(path), stack_pressure_levels=True)
+
+    state = _make_pl_state()
+    output.open(state)
+    output.write_step(state)
+    output.close()
+
+    tgm_file = tensogram.TensogramFile.open(str(path))
+    meta, objects = tgm_file[0]
+
+    assert len(objects) == 2 + len(PL_PARAMS)
+
+
+def test_stack_shape(tmp_path):
+    """Stacked objects have shape (n_grid, n_levels) -- grid axis first."""
+    path = tmp_path / "stacked.tgm"
+    context = _make_pl_context()
+    output = TensogramOutput(context, str(path), stack_pressure_levels=True)
+
+    state = _make_pl_state()
+    output.open(state)
+    output.write_step(state)
+    output.close()
+
+    tgm_file = tensogram.TensogramFile.open(str(path))
+    meta, objects = tgm_file[0]
+
+    for obj_idx in range(2, 2 + len(PL_PARAMS)):
+        desc, arr = objects[obj_idx]
+        assert arr.shape == (N_GRID, len(PL_LEVELS)), (
+            f"expected ({N_GRID}, {len(PL_LEVELS)}), got {arr.shape}"
+        )
+
+
+def test_stack_levels_metadata(tmp_path):
+    """Stacked objects store 'levelist' sorted ascending; no scalar 'level'."""
+    path = tmp_path / "stacked.tgm"
+    context = _make_pl_context()
+    output = TensogramOutput(context, str(path), stack_pressure_levels=True)
+
+    state = _make_pl_state()
+    output.open(state)
+    output.write_step(state)
+    output.close()
+
+    tgm_file = tensogram.TensogramFile.open(str(path))
+    meta, _ = tgm_file[0]
+
+    for obj_idx in range(2, 2 + len(PL_PARAMS)):
+        entry = meta.base[obj_idx]
+        anemoi = entry["anemoi"]
+        mars = entry["mars"]
+        assert "levelist" not in anemoi
+        assert "level" not in anemoi
+        assert "level" not in mars
+        assert mars["levelist"] == sorted(PL_LEVELS)
+        assert mars["levtype"] == "pl"
+        assert "name" in entry
+        assert entry["name"] == mars["param"]
+
+
+def test_stack_values_round_trip(tmp_path):
+    """Stacked values decode to the same data as the input, in level-sorted order."""
+    path = tmp_path / "stacked.tgm"
+    context = _make_pl_context(params=["t"], levels=[850, 500, 1000])
+    output = TensogramOutput(context, str(path), stack_pressure_levels=True)
+
+    rng = np.random.default_rng(7)
+    fields = {f"t{lv}": rng.random(N_GRID).astype(np.float32) for lv in [850, 500, 1000]}
+    state = {
+        "date": datetime(2024, 1, 1, 1),
+        "step": timedelta(hours=1),
+        "latitudes": np.linspace(-90, 90, N_GRID, dtype=np.float64),
+        "longitudes": np.linspace(0, 360, N_GRID, dtype=np.float64),
+        "fields": fields,
+    }
+
+    output.open(state)
+    output.write_step(state)
+    output.close()
+
+    tgm_file = tensogram.TensogramFile.open(str(path))
+    meta, objects = tgm_file[0]
+
+    mars = meta.base[2]["mars"]
+    assert mars["levelist"] == [500, 850, 1000]
+
+    _, arr = objects[2]
+    np.testing.assert_allclose(arr[:, 0], fields["t500"], rtol=1e-6)
+    np.testing.assert_allclose(arr[:, 1], fields["t850"], rtol=1e-6)
+    np.testing.assert_allclose(arr[:, 2], fields["t1000"], rtol=1e-6)
+
+
+def test_stack_non_pl_fields_written_flat(tmp_path):
+    """Non-PL fields are written as individual objects even with stacking enabled."""
+    path = tmp_path / "mixed.tgm"
+    context = _make_pl_context(params=["t"], levels=[500, 850], extra_fields=["2t"])
+    output = TensogramOutput(context, str(path), stack_pressure_levels=True)
+
+    state = _make_pl_state(params=["t"], levels=[500, 850], extra_fields=["2t"])
+    output.open(state)
+    output.write_step(state)
+    output.close()
+
+    tgm_file = tensogram.TensogramFile.open(str(path))
+    meta, objects = tgm_file[0]
+
+    assert len(objects) == 4
+
+    flat_entries = [
+        meta.base[i]["anemoi"]
+        for i in range(2, len(objects))
+        if meta.base[i]["anemoi"].get("variable") == "2t"
+    ]
+    assert len(flat_entries) == 1
+    assert "levels" not in flat_entries[0]
+    assert "level" not in flat_entries[0]
+
+
+def test_no_stack_level_metadata_correct(tmp_path):
+    """Without stacking, each PL field stores scalar 'level' and 'levtype'."""
+    path = tmp_path / "flat_pl.tgm"
+    context = _make_pl_context(params=["t"], levels=[500, 850])
+    output = TensogramOutput(context, str(path), stack_pressure_levels=False)
+
+    state = _make_pl_state(params=["t"], levels=[500, 850])
+    output.open(state)
+    output.write_step(state)
+    output.close()
+
+    tgm_file = tensogram.TensogramFile.open(str(path))
+    meta, objects = tgm_file[0]
+
+    assert len(objects) == 4
+
+    for obj_idx in range(2, 4):
+        mars = meta.base[obj_idx]["mars"]
+        anemoi = meta.base[obj_idx]["anemoi"]
+        assert "level" in mars
+        assert "levtype" in mars
+        assert mars["levtype"] == "pl"
+        assert mars["param"] == "t"
+        assert "levels" not in anemoi