fix: add utility function

ecmwf · Dec 12, 2024 · 202765c · 202765c
1 parent 3b89b82
commit 202765c
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/src/anemoi/utils/data_structures.py b/src/anemoi/utils/data_structures.py
@@ -337,3 +337,15 @@ def register_buffer(self, *, caller, name, persistent=False):
         assert '__' not in name, name
         for k, array in self.arrays.items():
             caller.register_buffer(f'{name}__{k}', array, persistent=persistent)
+
+
+def define_anemoi_tensor(data: dict, sources: list[str], indices: dict) -> TorchNestedAnemoiTensor:
+    import torch
+
+    tt = {}
+    for name in sources:
+        valid_values_mask = ~torch.isnan(data[name]).any((0, 2))
+        tt[name] = data[name][:, valid_values_mask][..., indices[name]]
+
+    return TorchNestedAnemoiTensor(tt)
+