Merge pull request #55 from ecmwf/bug/fix_traj_coords

Path now allows inflation on levelist
ecmwf · Dec 12, 2024 · b6a2d6f · b6a2d6f
2 parents a82622c + 46270c2
commit b6a2d6f
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 31 deletions.
diff --git a/covjsonkit/decoder/Path.py b/covjsonkit/decoder/Path.py
@@ -42,10 +42,12 @@ def to_xarray(self):
         # Get coordinates
         x = []
         y = []
+        z = []
         t = []
         for coord in self.get_coordinates()["composite"]["values"]:
             x.append(float(coord[1]))
             y.append(float(coord[2]))
+            z.append(float(coord[3]))
             t.append(coord[0])
 
         values = {}
@@ -108,13 +110,11 @@ def to_xarray(self):
                 points=(["points"], list(range(0, len(x)))),
                 x=(["points"], x),
                 y=(["points"], y),
+                z=(["points"], z),
                 t=(["points"], t),
             ),
         )
         for mars_metadata in self.mars_metadata[0]:
             ds.attrs[mars_metadata] = self.mars_metadata[0][mars_metadata]
 
-        # Add date attribute
-        # ds.attrs["date"] = self.get_coordinates()["t"]["values"][0]
-
         return ds
diff --git a/covjsonkit/encoder/Path.py b/covjsonkit/encoder/Path.py
@@ -109,6 +109,20 @@ def from_polytope(self, result):
             }
         )
 
+        for date in coords.keys():
+            coord = coords[date]["composite"]
+            coords[date]["composite"] = []
+            start = 0
+            for level in set(fields["l"]):
+                for s in set(fields["s"]):
+                    if (date, level, fields["number"][0], fields["param"][0], s) in range_dict:
+                        cor_len = len(range_dict[(date, level, fields["number"][0], fields["param"][0], s)])
+                        end = start + cor_len
+                        for cor in coord[int(start) : int(end)]:
+                            coords[date]["composite"].append([s, cor[0], cor[1], level])
+                        start = end
+        logging.debug("The coordinates returned from walking tree: %s", coords)  # noqa: E501
+
         combined_dict = {}
 
         for date in fields["dates"]:
@@ -122,7 +136,7 @@ def from_polytope(self, result):
                         if para not in combined_dict[date][num]:
                             combined_dict[date][num][para] = {}
                         # for s, value in range_dict[date][level][num][para].items():
-                        for s in fields["s"]:
+                        for s in set(fields["s"]):
                             key = (date, level, num, para, s)
                             # for k, v in range_dict.items():
                             # if k == key:
@@ -138,50 +152,26 @@ def from_polytope(self, result):
 
         logging.debug("The values returned from combined dicts: %s", combined_dict)  # noqa: E501
 
-        levels = fields["levels"]
         if fields["param"] == 0:
             raise ValueError("No parameters were returned, date requested may be out of range")
         for para in fields["param"]:
             self.add_parameter(para)
 
         logging.debug("The parameters added were: %s", self.parameters)  # noqa: E501
 
-        for date in coords.keys():
-            coord = coords[date]["composite"]
-            coords[date]["composite"] = []
-            for level in levels:
-                start = 0
-                for i, s in enumerate(fields["s"]):
-                    end = start + len(coord) / len(fields["s"])
-                    for cor in coord[int(start) : int(end)]:
-                        if len(fields["l"]) == 1:
-                            coords[date]["composite"].append([s, cor[0], cor[1], fields["l"][0]])
-                        elif len(fields["l"]) == 0:
-                            coords[date]["composite"].append([s, cor[0], cor[1], level])
-                        else:
-                            coords[date]["composite"].append([s, cor[0], cor[1], fields["l"][i]])
-                    start = end
-        logging.debug("The coordinates returned from walking tree: %s", coords)  # noqa: E501
-
         for date in combined_dict.keys():
             for num in combined_dict[date].keys():
                 val_dict = {}
-                # for step in combined_dict[date][num][self.parameters[0]].keys():
-                #    val_dict[step] = {}
                 for para in combined_dict[date][num].keys():
                     if para not in val_dict:
                         val_dict[para] = []
                     for step in combined_dict[date][num][para].keys():
                         val_dict[para].extend(combined_dict[date][num][para][step])
-                # for step in val_dict.keys():
                 mm = mars_metadata.copy()
                 mm["number"] = num
-                # mm["step"] = step
-                # temp = []
-                # for coord in coords[date]["composite"]:
-                #    temp.append([step] + coord)
-                # coords[date]["composite"] = temp
                 mm["Forecast date"] = date
+                if "levelist" in mm:
+                    del mm["levelist"]
                 self.add_coverage(mm, coords[date], val_dict)
 
         return self.covjson
diff --git a/covjsonkit/version.py b/covjsonkit/version.py
@@ -1 +1 @@
-__version__ = "0.0.28"
+__version__ = "0.1.0"