Add identifiers to mars:metadata if they are passed from polytope

covjsonkit/decoder/Position.py

@@ -109,25 +109,29 @@ def to_xarray(self):
         dims = ["latitude", "longitude", "levelist", "number", "datetime", "t"]
         ds = []
 
-        unique_coords = set()  # To track unique coordinate tuples
-        unique_domains = []  # To store unique domains
-
-        for domain in self.domains:
-            # Extract coordinate values
-            x = domain["axes"][self.x_name]["values"][0]
-            y = domain["axes"][self.y_name]["values"][0]
-            z = domain["axes"][self.z_name]["values"][0]
-            t = tuple(domain["axes"]["t"]["values"])  # Use tuple for hashable type
-
-            # Create a unique identifier for the domain
-            coord_tuple = (x, y, z, t)
-
-            # Check if this coordinate combination is already seen
-            if coord_tuple not in unique_coords:
-                unique_coords.add(coord_tuple)  # Mark as seen
-                unique_domains.append(domain)  # Add to unique domains
-
-        all_coords = unique_domains
+        # One dataset per requested point: coverages are grouped by domain and label,
+        # and points that snapped to the same grid point are kept apart.
+        def point_key(coverage):
+            axes = coverage["domain"]["axes"]
+            return (
+                axes[self.x_name]["values"][0],
+                axes[self.y_name]["values"][0],
+                axes[self.z_name]["values"][0],
+            )
+
+        def slot_key(coverage):
+            return (coverage["mars:metadata"]["number"], coverage["mars:metadata"]["Forecast date"])
+
+        # Within each point, one dataset per distinct time axis (as before), filled from
+        # that point's coverages only.
+        datasets = []
+        for group in self._point_groups(point_key, slot_key):
+            seen_t = set()
+            for coverage in group:
+                t = tuple(coverage["domain"]["axes"]["t"]["values"])
+                if t not in seen_t:
+                    seen_t.add(t)
+                    datasets.append((group, coverage["domain"]))
 
         num = []
         datetime = []
@@ -137,8 +141,8 @@ def to_xarray(self):
         nums = list(set(num))
         datetime = list(set(datetime))
 
-        # Process each coordinate domain
-        for coords in all_coords:
+        # Process each requested point and time axis
+        for group, coords in datasets:
             dataarraydict = {}
             x = coords["axes"][self.x_name]["values"]
             y = coords["axes"][self.y_name]["values"]
@@ -147,7 +151,7 @@ def to_xarray(self):
             steps = [step.replace("Z", "") for step in steps]
             steps = pd.to_datetime(steps)
 
-            cov_idx_list = self._find_coverages(nums, datetime, x, y, z)
+            cov_idx_list = self._find_coverages(nums, datetime, x, y, z, group)
 
             coords = {
                 "latitude": x,
@@ -181,25 +185,21 @@ def to_xarray(self):
                     attrs,
                 )
 
-            ds.append(xr.Dataset(data_vars=dataarraydict, coords=coords))
-
-        # Combine all DataArrays into a Dataset
-        for mars_metadata in self.mars_metadata[0]:
-            if mars_metadata != "date" and mars_metadata != "step":
-                for dss in ds:
-                    dss.attrs[mars_metadata] = self.mars_metadata[0][mars_metadata]
+            dss = xr.Dataset(data_vars=dataarraydict, coords=coords)
+            dss.attrs.update(self._point_dataset_attrs(group))
+            ds.append(dss)
 
         if len(ds) == 1:
             return ds[0]
 
         return ds
 
-    def _find_coverages(self, nums, datetime, x, y, z):
+    def _find_coverages(self, nums, datetime, x, y, z, coverages=None):
         """Find coverages matching domain parameters and return with indices."""
         result = []
         for i, num in enumerate(nums):
             for j, date in enumerate(datetime):
-                for coverage in self.covjson["coverages"]:
+                for coverage in coverages if coverages is not None else self.covjson["coverages"]:
                     if self._covers_domain(coverage, num, date, x, y, z):
                         result.append((i, j, coverage))
         return result

covjsonkit/decoder/TimeSeries.py

@@ -117,28 +117,29 @@ def to_xarray(self):
         dims = ["latitude", "longitude", "levelist", "number", "datetime", "t"]
         ds = []
 
-        # Get coordinates for all domains
-        all_coords = self.get_domains()
-
-        unique_coords = set()  # To track unique coordinate tuples
-        unique_domains = []  # To store unique domains
-
-        for domain in self.domains:
-            # Extract coordinate values
-            x = domain["axes"][self.x_name]["values"][0]
-            y = domain["axes"][self.y_name]["values"][0]
-            z = domain["axes"][self.z_name]["values"][0]
-            t = tuple(domain["axes"]["t"]["values"])  # Use tuple for hashable type
-
-            # Create a unique identifier for the domain
-            coord_tuple = (x, y, z, t)
+        # One dataset per requested point: coverages are grouped by domain and label,
+        # and points that snapped to the same grid point are kept apart.
+        def point_key(coverage):
+            axes = coverage["domain"]["axes"]
+            return (
+                axes[self.x_name]["values"][0],
+                axes[self.y_name]["values"][0],
+                axes[self.z_name]["values"][0],
+            )
 
-            # Check if this coordinate combination is already seen
-            if coord_tuple not in unique_coords:
-                unique_coords.add(coord_tuple)  # Mark as seen
-                unique_domains.append(domain)  # Add to unique domains
+        def slot_key(coverage):
+            return (coverage["mars:metadata"]["number"], coverage["mars:metadata"]["Forecast date"])
 
-        all_coords = unique_domains
+        # Within each point, one dataset per distinct time axis (as before), filled from
+        # that point's coverages only.
+        datasets = []
+        for group in self._point_groups(point_key, slot_key):
+            seen_t = set()
+            for coverage in group:
+                t = tuple(coverage["domain"]["axes"]["t"]["values"])
+                if t not in seen_t:
+                    seen_t.add(t)
+                    datasets.append((group, coverage["domain"]))
 
         num = []
         datetime = []
@@ -148,8 +149,8 @@ def to_xarray(self):
         nums = list(set(num))
         datetime = list(set(datetime))
 
-        # Process each coordinate domain
-        for coords in all_coords:
+        # Process each requested point and time axis
+        for group, coords in datasets:
             dataarraydict = {}
             x = coords["axes"][self.x_name]["values"]
             y = coords["axes"][self.y_name]["values"]
@@ -158,7 +159,7 @@ def to_xarray(self):
             steps = [step.replace("Z", "") for step in steps]
             steps = pd.to_datetime(steps)
 
-            cov_idx_list = self._find_coverages(nums, datetime, x, y, z)
+            cov_idx_list = self._find_coverages(nums, datetime, x, y, z, group)
 
             coords = {
                 "latitude": x,
@@ -192,26 +193,22 @@ def to_xarray(self):
                     attrs,
                 )
 
-            ds.append(xr.Dataset(data_vars=dataarraydict, coords=coords))
-
-        # Combine all DataArrays into a Dataset
-        for mars_metadata in self.mars_metadata[0]:
-            if mars_metadata != "date" and mars_metadata != "step":
-                for dss in ds:
-                    dss.attrs[mars_metadata] = self.mars_metadata[0][mars_metadata]
+            dss = xr.Dataset(data_vars=dataarraydict, coords=coords)
+            dss.attrs.update(self._point_dataset_attrs(group))
+            ds.append(dss)
 
         if len(ds) == 1:
             return ds[0]
 
         return ds
 
-    def _find_coverages(self, nums, datetime, x, y, z):
+    def _find_coverages(self, nums, datetime, x, y, z, coverages=None):
         """Find coverages that match the given domain parameters (num, date, x, y, z)
         and return them along with domain parameter indices."""
         result = []
         for i, num in enumerate(nums):
             for j, date in enumerate(datetime):
-                for coverage in self.covjson["coverages"]:
+                for coverage in coverages if coverages is not None else self.covjson["coverages"]:
                     if self._covers_domain(coverage, num, date, x, y, z):
                         result.append((i, j, coverage))
         return result

covjsonkit/decoder/VerticalProfile.py

@@ -112,34 +112,29 @@ def to_xarray(self):
         ]
         ds = []
 
-        # Get coordinates for all domains
-        all_coords = self.get_domains()
-
-        unique_coords = set()  # To track unique coordinate tuples
-        unique_domains = []  # To store unique domains
-
-        for domain in self.domains:
-            # Extract coordinate values
-            x = domain["axes"][self.x_name]["values"][0]
-            y = domain["axes"][self.y_name]["values"][0]
-            z = domain["axes"][self.z_name]["values"][0]
-
-            # Create a unique identifier for the domain
-            coord_tuple = (x, y, z)
-
-            # Check if this coordinate combination is already seen
-            if coord_tuple not in unique_coords:
-                unique_coords.add(coord_tuple)  # Mark as seen
-                unique_domains.append(domain)  # Add to unique domains
-
-        all_coords = unique_domains
+        # One dataset per requested point: coverages are grouped by domain and label,
+        # and points that snapped to the same grid point are kept apart.
+        def domain_key(coverage):
+            axes = coverage["domain"]["axes"]
+            return (
+                axes[self.x_name]["values"][0],
+                axes[self.y_name]["values"][0],
+                axes[self.z_name]["values"][0],
+            )
+
+        def slot_key(coverage):
+            meta = coverage["mars:metadata"]
+            return (meta["number"], meta["Forecast date"], meta["step"])
+
+        groups = self._point_groups(domain_key, slot_key)
         param_values = {}
 
         # Initialize parameter values for all parameters
         for parameter in self.parameters:
             param_values[parameter] = []
 
-        for domain_idx, coords in enumerate(all_coords):
+        for domain_idx, group in enumerate(groups):
+            coords = group[0]["domain"]
             dataarraydict = {}
 
             # Get coordinates
@@ -175,7 +170,7 @@ def to_xarray(self):
                         for k, step in enumerate(steps):
                             if len(param_values[parameter][domain_idx][i][j]) <= k:
                                 param_values[parameter][domain_idx][i][j].append([])
-                            for coverage in self.covjson["coverages"]:
+                            for coverage in group:
                                 new_step = (
                                     dt.fromisoformat(date.replace("Z", "")) + timedelta(hours=parse_step_string(step))
                                 ).isoformat() + "Z"
@@ -211,12 +206,9 @@ def to_xarray(self):
                 dataarray.attrs["long_name"] = self.get_parameter_metadata(parameter)["observedProperty"]["id"]
                 dataarraydict[dataarray.attrs["long_name"]] = dataarray
 
-            ds.append(xr.Dataset(dataarraydict))
-
-        for mars_metadata in self.mars_metadata[0]:
-            for dss in ds:
-                if mars_metadata != "date" and mars_metadata != "step":
-                    dss.attrs[mars_metadata] = self.mars_metadata[0][mars_metadata]
+            dss = xr.Dataset(dataarraydict)
+            dss.attrs.update(self._point_dataset_attrs(group))
+            ds.append(dss)
 
         if len(ds) == 1:
             return ds[0]

covjsonkit/decoder/decoder.py

@@ -91,6 +91,39 @@ def get_mars_metadata(self):
                 mars_metadata.append(coverage["mars:metadata"])
         return mars_metadata
 
+    def _point_groups(self, domain_key, slot_key):
+        """Group coverages into one group per requested point, in first-seen order.
+
+        Coverages with the same ``domain_key`` (spatial/time domain) and ``label``
+        belong to the same point. When several requested points snapped to the same
+        grid point, their coverages repeat for each ``slot_key`` (e.g. number and
+        forecast date), so the n-th coverage for a given domain, label and slot
+        belongs to the n-th such point.
+        """
+        groups = {}
+        order = []
+        seen = {}
+        for coverage in self.covjson["coverages"]:
+            label = coverage.get("mars:metadata", {}).get("label")
+            domain = domain_key(coverage)
+            slot = (domain, label, slot_key(coverage))
+            occurrence = seen.get(slot, 0)
+            seen[slot] = occurrence + 1
+            key = (domain, label, occurrence)
+            if key not in groups:
+                groups[key] = []
+                order.append(key)
+            groups[key].append(coverage)
+        return [groups[key] for key in order]
+
+    def _point_dataset_attrs(self, group):
+        """Dataset attributes for one point: shared MARS metadata plus the point's own label."""
+        attrs = {key: val for key, val in self.mars_metadata[0].items() if key not in ("date", "step", "label")}
+        label = group[0].get("mars:metadata", {}).get("label")
+        if label is not None:
+            attrs["label"] = label
+        return attrs
+
     @abstractmethod
     def get_ranges(self):
         pass