Skip to content
Snippets Groups Projects
Commit c5d2f62d authored by Nils-Arne Dreier's avatar Nils-Arne Dreier
Browse files

feat: copy metadata from yac source

parent bf6312b3
No related branches found
No related tags found
1 merge request!42feat: copy metadata from yac source
......@@ -5,8 +5,18 @@ from argparse import ArgumentParser
from itertools import chain, groupby
import numpy as np
import yaml
import zarr
from coyote import Coyote, group_comm_rank, group_comm_size, init, run, start_datetime
from coyote import (
Coyote,
ensure_enddef,
get_field_metadata,
group_comm_rank,
group_comm_size,
init,
run,
start_datetime,
)
from ._data_handler import DataHandler
from ._distribute_work import distribute_work
......@@ -74,11 +84,12 @@ def main():
item.group = group # amend the variable object with the group
yield item
for _name, item in group.groups():
item.parent = group
yield from collect_data_vars(item)
data_vars = list(chain(*[collect_data_vars(z) for z in args.datasets]))
logging.info(f"Found {len(data_vars)} variables")
if len(data_vars) == 0:
all_data_vars = list(chain(*[collect_data_vars(z) for z in args.datasets]))
logging.info(f"Found {len(all_data_vars)} variables")
if len(all_data_vars) == 0:
raise RuntimeError("No variables found by the hiopy worker.")
# group the variables by the crs grid_mapping.
......@@ -86,7 +97,8 @@ def main():
grouped_data_vars = {
gid: list(variables)
for gid, variables in groupby(
sorted(data_vars, key=lambda v: grid_id(v, v.group)), key=lambda v: grid_id(v, v.group)
sorted(all_data_vars, key=lambda v: grid_id(v, v.group)),
key=lambda v: grid_id(v, v.group),
)
}
distributed_data_vars = distribute_work(grouped_data_vars, group_comm_size())
......@@ -184,6 +196,24 @@ def main():
frac_mask=frac_mask,
)
def get_source_triple(v):
if "hiopy::parent" in v.attrs:
return get_source_triple(v.group.parent[v.attrs["hiopy::parent"]][v.basename])
else:
src_comp, src_grid = v.attrs["hiopy::yac_source"]
src_field = v.attrs["hiopy::src_name"] if "src_name" else v.basename
return src_comp, src_grid, src_field
ensure_enddef()
if group_comm_rank() == 0:
for v in all_data_vars:
if "hiopy::copy_metadata" in v.attrs:
comp, grid, field = get_source_triple(v)
metadata = yaml.safe_load(get_field_metadata(comp, grid, field))
for key, value in metadata["cf"]:
v.attrs[key] = value
del v.attrs["hiopy::copy_metadata"] # copy only once
run()
for dh in data_handlers:
......
......@@ -82,9 +82,6 @@ PYBIND11_MODULE(coyote, m) {
"cell_to_vertex"_a,
py::arg("x_cells") = std::nullopt,
py::arg("y_cells") = std::nullopt)
.def("get_field_collection_size", &Coyote::get_field_collection_size)
.def("get_field_timestep", &Coyote::get_field_timestep)
.def("get_field_metadata", &Coyote::get_field_metadata)
.def("get_all_comp_grid_fields", &Coyote::get_all_comp_grid_fields);
py::class_<Event>(m, "Event")
......@@ -101,6 +98,10 @@ PYBIND11_MODULE(coyote, m) {
})
.def("release", &Event::release);
m.def("ensure_enddef", &ensure_enddef);
m.def("get_field_timestep", &get_field_timestep);
m.def("get_field_collection_size", &get_field_collection_size);
m.def("get_field_metadata", &get_field_metadata);
m.def("run", &coyote_run, py::call_guard<py::gil_scoped_release>());
m.def("init", &coyote_init, "group_name"_a, "nthreads"_a = 1);
m.def("start_datetime", &coyote_start_datetime);
......
......@@ -211,49 +211,6 @@ namespace coyote {
}
}
int get_field_collection_size(const std::string& source_comp,
const std::string& source_grid,
const std::string& field_name) const {
return yac_cget_field_collection_size_instance(
CoyoteEnv::getInstance().get_yac_instance_id(),
source_comp.c_str(),
source_grid.c_str(),
field_name.c_str());
}
std::string get_field_timestep(const std::string& source_comp,
const std::string& source_grid,
const std::string& field_name) const {
std::string ret_val = "";
const char* timestep = yac_cget_field_timestep_instance(
CoyoteEnv::getInstance().get_yac_instance_id(),
source_comp.c_str(),
source_grid.c_str(),
field_name.c_str());
if (timestep) {
ret_val = std::string(timestep);
}
return ret_val;
}
std::string get_field_metadata(const std::string& source_comp,
const std::string& source_grid,
const std::string& field_name) const {
std::string ret_val = "";
const char* metadata = yac_cget_field_metadata_instance(
CoyoteEnv::getInstance().get_yac_instance_id(),
source_comp.c_str(),
source_grid.c_str(),
field_name.c_str());
if (metadata) {
ret_val = std::string(metadata);
}
return ret_val;
}
std::string get_start_datetime() const {
return yac_cget_start_datetime_instance(CoyoteEnv::getInstance().get_yac_instance_id());
}
......@@ -387,24 +344,6 @@ namespace coyote {
return pImpl->get_all_comp_grid_fields();
}
int Coyote::get_field_collection_size(const std::string& source_comp,
const std::string& source_grid,
const std::string& field_name) const {
return pImpl->get_field_collection_size(source_comp, source_grid, field_name);
}
std::string Coyote::get_field_timestep(const std::string& source_comp,
const std::string& source_grid,
const std::string& field_name) const {
return pImpl->get_field_timestep(source_comp, source_grid, field_name);
}
std::string Coyote::get_field_metadata(const std::string& source_comp,
const std::string& source_grid,
const std::string& field_name) const {
return pImpl->get_field_metadata(source_comp, source_grid, field_name);
}
std::string Coyote::get_start_datetime() const { return pImpl->get_start_datetime(); }
std::string Coyote::get_end_datetime() const { return pImpl->get_end_datetime(); }
......@@ -418,6 +357,49 @@ namespace coyote {
MPI_Comm coyote_group_comm() { return CoyoteEnv::getInstance().get_group_comm(); }
void ensure_enddef() { CoyoteEnv::getInstance().ensure_enddef(); }
int get_field_collection_size(const std::string& source_comp,
const std::string& source_grid,
const std::string& field_name) {
return yac_cget_field_collection_size_instance(
CoyoteEnv::getInstance().get_yac_instance_id(),
source_comp.c_str(),
source_grid.c_str(),
field_name.c_str());
}
std::string get_field_timestep(const std::string& source_comp,
const std::string& source_grid,
const std::string& field_name) {
std::string ret_val = "";
const char* timestep = yac_cget_field_timestep_instance(
CoyoteEnv::getInstance().get_yac_instance_id(),
source_comp.c_str(),
source_grid.c_str(),
field_name.c_str());
if (timestep) {
ret_val = std::string(timestep);
}
return ret_val;
}
std::string get_field_metadata(const std::string& source_comp,
const std::string& source_grid,
const std::string& field_name) {
const char* metadata = yac_cget_field_metadata_instance(
CoyoteEnv::getInstance().get_yac_instance_id(),
source_comp.c_str(),
source_grid.c_str(),
field_name.c_str());
if (metadata) {
return std::string(metadata);
}
return "";
}
void coyote_run() { CoyoteEnv::getInstance().run(); }
} // namespace coyote
......@@ -55,15 +55,6 @@ namespace coyote {
std::optional<std::span<const int>> global_cell_idx = std::nullopt);
// functions to get the necessary info on the available variables
int get_field_collection_size(const std::string& source_comp,
const std::string& source_grid,
const std::string& field_name) const;
std::string get_field_timestep(const std::string& source_comp,
const std::string& source_grid,
const std::string& field_name) const;
std::string get_field_metadata(const std::string& source_comp,
const std::string& source_grid,
const std::string& field_name) const;
std::string get_start_datetime() const;
std::string get_end_datetime() const;
......@@ -75,6 +66,16 @@ namespace coyote {
MPI_Comm coyote_group_comm();
std::string coyote_start_datetime();
std::string coyote_end_datetime();
void ensure_enddef();
int get_field_collection_size(const std::string& source_comp,
const std::string& source_grid,
const std::string& field_name);
std::string get_field_timestep(const std::string& source_comp,
const std::string& source_grid,
const std::string& field_name);
std::string get_field_metadata(const std::string& source_comp,
const std::string& source_grid,
const std::string& field_name);
void coyote_run();
} // namespace coyote
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment