Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
156 changes: 78 additions & 78 deletions covjsonkit/decoder/TimeSeries.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,18 +9,10 @@ def __init__(self, covjson):
super().__init__(covjson)
self.domains = self.get_domains()
self.ranges = self.get_ranges()
if "x" in self.covjson["coverages"][0]["domain"]["axes"]:
self.x_name = "x"
else:
self.x_name = "latitude"
if "y" in self.covjson["coverages"][0]["domain"]["axes"]:
self.y_name = "y"
else:
self.y_name = "longitude"
if "z" in self.covjson["coverages"][0]["domain"]["axes"]:
self.z_name = "z"
else:
self.z_name = "levelist"
first_axes = self.covjson["coverages"][0]["domain"]["axes"]
self.x_name = "x"
self.y_name = "y"
self.z_name = "z" if "z" in first_axes else None

def get_domains(self):
domains = []
Expand Down Expand Up @@ -48,9 +40,9 @@ def get_coordinates(self):
coord_dict[param] = []
# Get x,y,z,t coords and unpack t coords and match to x,y,z coords
for ind, domain in enumerate(self.domains):
x = domain["axes"][self.x_name]["values"][0]
y = domain["axes"][self.y_name]["values"][0]
z = domain["axes"][self.z_name]["values"][0]
longitude = domain["axes"][self.x_name]["values"][0]
latitude = domain["axes"][self.y_name]["values"][0]
level = domain["axes"][self.z_name]["values"][0] if self.z_name else None
fct = domain["axes"]["t"]["values"][0]
ts = domain["axes"]["t"]["values"]
if "number" in self.mars_metadata[ind]:
Expand All @@ -62,7 +54,7 @@ def get_coordinates(self):
for t in ts:
# Have to replicate these coords for each parameter
# coordinates.append([x, y, z, t])
coords.append([x, y, z, fct, t, num])
coords.append([latitude, longitude, level, fct, t, num])
coord_dict[param].append(coords)
return coord_dict

Expand All @@ -75,13 +67,16 @@ def to_geotiff(self):
def to_geojson(self):
features = []
for coverage in self.covjson["coverages"]:
lat = coverage["domain"]["axes"][self.x_name]["values"][0]
lon = coverage["domain"]["axes"][self.y_name]["values"][0]
z = coverage["domain"]["axes"][self.z_name]["values"][0]
longitude = coverage["domain"]["axes"][self.x_name]["values"][0]
latitude = coverage["domain"]["axes"][self.y_name]["values"][0]
datetimes = coverage["domain"]["axes"]["t"]["values"]
if "mars:metadata" in coverage:
mars_metadata = coverage["mars:metadata"]

geom_coords = [longitude, latitude]
if self.z_name:
geom_coords.append(coverage["domain"]["axes"][self.z_name]["values"][0])

values = {}
for key in coverage["ranges"]:
values[key] = coverage["ranges"][key]["values"]
Expand All @@ -96,7 +91,7 @@ def to_geojson(self):
features.append(
{
"type": "Feature",
"geometry": {"type": "Point", "coordinates": [lon, lat, z]},
"geometry": {"type": "Point", "coordinates": geom_coords},
"properties": param_vals,
}
)
Expand All @@ -114,29 +109,31 @@ def to_xarray(self):
if not has_forecast_date:
return self._to_xarray_no_forecast_date()

dims = ["latitude", "longitude", "levelist", "number", "datetime", "t"]
if self.z_name:
dims = ["latitude", "longitude", "levelist", "number", "datetime", "t"]
else:
dims = ["latitude", "longitude", "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
unique_coords = set()
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
longitude = domain["axes"][self.x_name]["values"][0]
latitude = domain["axes"][self.y_name]["values"][0]
t = tuple(domain["axes"]["t"]["values"])

# Create a unique identifier for the domain
coord_tuple = (x, y, z, t)
if self.z_name:
z = domain["axes"][self.z_name]["values"][0]
coord_tuple = (longitude, latitude, z, t)
else:
coord_tuple = (longitude, latitude, 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
unique_coords.add(coord_tuple)
unique_domains.append(domain)

all_coords = unique_domains

Expand All @@ -151,28 +148,36 @@ def to_xarray(self):
# Process each coordinate domain
for coords in all_coords:
dataarraydict = {}
x = coords["axes"][self.x_name]["values"]
y = coords["axes"][self.y_name]["values"]
z = coords["axes"][self.z_name]["values"]
longitude = coords["axes"][self.x_name]["values"]
latitude = coords["axes"][self.y_name]["values"]
steps = coords["axes"]["t"]["values"]
steps = [step.replace("Z", "") for step in steps]
steps = pd.to_datetime(steps)

cov_idx_list = self._find_coverages(nums, datetime, x, y, z)

coords = {
"latitude": x,
"longitude": y,
"levelist": z,
"number": nums,
"datetime": datetime,
"t": steps,
}
if self.z_name:
z = coords["axes"][self.z_name]["values"]
cov_idx_list = self._find_coverages(nums, datetime, longitude, latitude, z)
coord_dict = {
"latitude": latitude,
"longitude": longitude,
"levelist": z,
"number": nums,
"datetime": datetime,
"t": steps,
}
else:
cov_idx_list = self._find_coverages(nums, datetime, longitude, latitude, None)
coord_dict = {
"latitude": latitude,
"longitude": longitude,
"number": nums,
"datetime": datetime,
"t": steps,
}

for parameter in self.parameters:
param_values = [[[] for _ in range(len(datetime))] for _ in range(len(nums))]

# Extract parameter values for the current domain
for i, j, cov in cov_idx_list:
param_values[i][j] = cov["ranges"][parameter]["values"]

Expand All @@ -186,15 +191,13 @@ def to_xarray(self):
"units": self.get_parameter_metadata(parameter)["unit"]["symbol"],
"long_name": long_name,
}
dataarraydict[long_name] = (
dims,
[[[param_values]]],
attrs,
)
if self.z_name:
dataarraydict[long_name] = (dims, [[[param_values]]], attrs)
else:
dataarraydict[long_name] = (dims, [[param_values]], attrs)

ds.append(xr.Dataset(data_vars=dataarraydict, coords=coords))
ds.append(xr.Dataset(data_vars=dataarraydict, coords=coord_dict))

# 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:
Expand All @@ -205,43 +208,37 @@ def to_xarray(self):

return ds

def _find_coverages(self, nums, datetime, x, y, z):
"""Find coverages that match the given domain parameters (num, date, x, y, z)
and return them along with domain parameter indices."""
def _find_coverages(self, nums, datetime, longitude, latitude, z):
result = []
for i, num in enumerate(nums):
for j, date in enumerate(datetime):
for coverage in self.covjson["coverages"]:
if self._covers_domain(coverage, num, date, x, y, z):
if self._covers_domain(coverage, num, date, longitude, latitude, z):
result.append((i, j, coverage))
return result

def _covers_domain(self, coverage, num, date, x, y, z):
"""check if coverage matches the given domain parameters (num, date, x, y, z)"""
return (
def _covers_domain(self, coverage, num, date, longitude, latitude, z):
axes = coverage["domain"]["axes"]
match = (
coverage["mars:metadata"]["number"] == num
and coverage["mars:metadata"]["Forecast date"] == date
and coverage["domain"]["axes"][self.x_name]["values"] == x
and coverage["domain"]["axes"][self.y_name]["values"] == y
and coverage["domain"]["axes"][self.z_name]["values"] == z
and axes[self.x_name]["values"] == longitude
and axes[self.y_name]["values"] == latitude
)
if match and self.z_name and z is not None:
match = axes[self.z_name]["values"] == z
return match

def _to_xarray_no_forecast_date(self):
"""Convert monthly-means CovJSON (no 'Forecast date' in metadata) to xarray.

In this layout every coverage contains all of its time steps packed into
the domain's ``t`` axis (produced by ``from_polytope_month``). Each
coverage corresponds to one spatial point / ensemble member combination.
The resulting Dataset has a ``t`` dimension whose values come directly
from the domain axis.
"""Monthly-means path: all time steps are packed into a single coverage's
t-axis and there is no "Forecast date" in mars:metadata.
"""
ds_list = []

for coverage in self.covjson["coverages"]:
domain = coverage["domain"]["axes"]
x = domain[self.x_name]["values"]
y = domain[self.y_name]["values"]
z = domain[self.z_name]["values"]
longitude = domain[self.x_name]["values"]
latitude = domain[self.y_name]["values"]

steps = domain["t"]["values"]
steps = [s.replace("Z", "") for s in steps]
Expand All @@ -266,10 +263,13 @@ def _to_xarray_no_forecast_date(self):
dataarraydict[long_name] = dataarray

coord_dict = dict(
latitude=(["latitude"], x),
longitude=(["longitude"], y),
levelist=(["levelist"], z),
latitude=(["latitude"], latitude),
longitude=(["longitude"], longitude),
)
if self.z_name:
z = domain[self.z_name]["values"]
coord_dict["levelist"] = (["levelist"], z)

dss = xr.Dataset(dataarraydict, coords=coord_dict)

# Attach MARS metadata (skip keys that vary per time step)
Expand Down
Loading
Loading