Add backends and mir interface

src/earthkit/regrid/backends/__init__.py

@@ -0,0 +1,210 @@
+# (C) Copyright 2023 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+import threading
+from abc import ABCMeta
+from abc import abstractmethod
+from collections import namedtuple
+
+from earthkit.regrid.utils.config import CONFIG
+
+BackendKey = namedtuple("BackendKey", ["name", "path"])
+
+
+class Backend(metaclass=ABCMeta):
+    @abstractmethod
+    def interpolate(self, values, in_grid, out_grid, method, **kwargs):
+        pass
+
+
+class BackendOrder:
+    DEFAULT_ORDER = ["local-matrix", "plugins", "remote-matrix", "system-matrix", "mir"]
+    BUILT_IN = {"local-matrix", "remote-matrix", "system-matrix", "mir"}
+    UNIQUE = {"system-matrix", "mir"}
+
+    def select(self, order, names, backends):
+        if isinstance(names, str):
+            if names in self.UNIQUE:
+                return [b for k, b in backends.items() if k.name == names]
+
+        r = []
+        if names:
+            _order = list(names)
+        else:
+            _order = order or self.DEFAULT_ORDER
+
+        print("names", names)
+        print("order", _order)
+
+        for m in _order:
+            print("m", m)
+            for k, b in backends.items():
+                print(" k", k)
+                if k.name == m or (m == "plugins" and k.name not in self.BUILT_IN):
+                    print("  -> b", b)
+                    r.append(b)
+
+        return r
+
+
+BACKEND_ORDER = BackendOrder()
+
+
+class BackendManager:
+    BACKENDS = {}
+
+    def __init__(self, *args, policy="all", **kwargs):
+        self.order = []
+        # self._has_settings = False
+        self.lock = threading.Lock()
+        # self.update(*args, **kwargs)
+
+        # # initialise the backend list
+        # self.BACKENDS.update(self.local())
+        # self.BACKENDS.update(self.remote())
+        # self.BACKENDS.update(self.mir())
+        self.update()
+
+        print("BACKENDS", self.BACKENDS)
+
+    def update(self):
+        with self.lock:
+            self.BACKENDS.clear()
+            self.BACKENDS.update(self._local())
+            self.BACKENDS.update(self._remote())
+            self.BACKENDS.update(self._mir())
+
+            from earthkit.regrid.utils.config import CONFIG
+
+            self.order = CONFIG.get("backend-order", []) or []
+
+    def backends(self, backend=None):
+        with self.lock:
+            return BACKEND_ORDER.select(self.order, backend, self.BACKENDS)
+
+    def _local(self):
+        from earthkit.regrid.utils.config import CONFIG
+
+        from .matrix import LocalMatrixBackend
+
+        dirs = CONFIG.get("local-matrix-directories", [])
+        if dirs is None:
+            dirs = []
+        if isinstance(dirs, str):
+            dirs = [dirs]
+
+        r = {}
+        for d in dirs:
+            r[BackendKey("local-matrix", d)] = LocalMatrixBackend(d)
+        return r
+
+    def _remote(self):
+        from earthkit.regrid.utils.config import CONFIG
+
+        from .matrix import RemoteMatrixBackend
+        from .matrix import SystemRemoteMatrixBackend
+
+        dirs = CONFIG.get("remote-matrix-directories", [])
+        if dirs is None:
+            dirs = []
+        if isinstance(dirs, str):
+            dirs = [dirs]
+
+        r = {}
+        for d in dirs:
+            r[BackendKey("remote-matrix", d)] = RemoteMatrixBackend(d)
+
+        r[BackendKey("system-matrix", None)] = SystemRemoteMatrixBackend()
+
+        return r
+
+    def _mir(self):
+        from .mir import MirBackend
+
+        return {BackendKey("mir", None): MirBackend()}
+
+
+MANAGER = BackendManager()
+
+# def add_matrix_source(path):
+#     global DB_LIST
+#     for item in DB_LIST[1:]:
+#         if item.matrix_source() == path:
+#             return item
+#     db = MatrixDb.from_path(path)
+#     DB_LIST.append(db)s
+#     return db
+
+
+# def find(*args, matrix_source=None, **kwargs):
+#     if matrix_source is None:
+#         return SYS_DB.find(*args, **kwargs)
+#     else:
+#         db = add_matrix_source(matrix_source)
+#         return db.find(*args, **kwargs)
+
+
+# class Backend(metaclass=ABCMeta):
+#     @abstractmethod
+#     def interpolate(self, values, in_grid, out_grid, method, **kwargs):
+#         pass
+
+
+# class BackendLoader:
+#     kind = "backend"
+
+#     def load_module(self, module):
+#         return import_module(module, package=__name__).backend
+
+#     def load_entry(self, entry):
+#         entry = entry.load()
+#         if callable(entry):
+#             return entry
+#         return entry.backend
+
+#     def load_remote(self, name):
+#         return None
+
+
+# class BackendMaker:
+#     BACKENDS = {}
+
+#     def __init__(self):
+#         # Preregister the most important backends
+#         from .mir import MirBackend
+#         from .matrix import LocalMatrixBackend
+#         from .matrix import RemoteMatrixBackend
+
+#         self.BACKENDS["mir"] = MirBackend
+#         self.BACKENDS["local"] = LocalMatrixBackend
+#         self.BACKENDS["remote"] = RemoteMatrixBackend
+
+#     def __call__(self, name, *args, **kwargs):
+#         loader = BackendLoader()
+
+#         if name in self.BACKENDS:
+#             klass = self.BACKENDS[name]
+#         else:
+#             klass = find_plugin(os.path.dirname(__file__), name, loader)
+#             self.BACKENDS[name] = klass
+
+#         backend = klass(*args, **kwargs)
+
+#         if getattr(backend, "name", None) is None:
+#             backend.name = name
+
+#         return backend
+
+#     def __getattr__(self, name):
+#         return self(name.replace("_", "-"))
+
+
+# get_backend = BackendMaker()
+
+CONFIG.on_change(MANAGER.update)

src/earthkit/regrid/db.py → src/earthkit/regrid/backends/db.py

@@ -542,6 +542,10 @@ def matrix_source(self):
     def from_path(path):
         return MatrixDb(LocalAccessor(path))
 
+    @staticmethod
+    def from_url(url):
+        return MatrixDb(UrlAccessor(url))
+
     def __len__(self):
         return len(self.index)
 
@@ -556,22 +560,22 @@ def _reset(self):
 
 
 SYS_DB = MatrixDb(SystemAccessor())
-DB_LIST = [SYS_DB]
+# DB_LIST = [SYS_DB]
 
 
-def add_matrix_source(path):
-    global DB_LIST
-    for item in DB_LIST[1:]:
-        if item.matrix_source() == path:
-            return item
-    db = MatrixDb.from_path(path)
-    DB_LIST.append(db)
-    return db
+# def add_matrix_source(path):
+#     global DB_LIST
+#     for item in DB_LIST[1:]:
+#         if item.matrix_source() == path:
+#             return item
+#     db = MatrixDb.from_path(path)
+#     DB_LIST.append(db)
+#     return db
 
 
-def find(*args, matrix_source=None, **kwargs):
-    if matrix_source is None:
-        return SYS_DB.find(*args, **kwargs)
-    else:
-        db = add_matrix_source(matrix_source)
-        return db.find(*args, **kwargs)
+# def find(*args, matrix_source=None, **kwargs):
+#     if matrix_source is None:
+#         return SYS_DB.find(*args, **kwargs)
+#     else:
+#         db = add_matrix_source(matrix_source)
+#         return db.find(*args, **kwargs)

src/earthkit/regrid/backends/matrix.py

@@ -0,0 +1,65 @@
+# (C) Copyright 2023 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+from abc import abstractmethod
+from functools import cached_property
+
+from . import Backend
+
+
+class MatrixBackend(Backend):
+    def __init__(self, path_or_url=None):
+        self.path_or_url = path_or_url
+
+    def interpolate(self, values, in_grid, out_grid, method, **kwargs):
+        z, shape = self.db.find(in_grid, out_grid, method, **kwargs)
+
+        if z is None:
+            raise ValueError(f"No matrix found! {in_grid=} {out_grid=} {method=}")
+
+        # This should check for 1D (GG) and 2D (LL) matrices
+        values = values.reshape(-1, 1)
+
+        print("values.shape", values.shape)
+        print("z.shape", z.shape)
+
+        values = z @ values
+
+        print("values.shape", values.shape)
+
+        return values.reshape(shape)
+
+    @property
+    @abstractmethod
+    def db(self):
+        pass
+
+
+class LocalMatrixBackend(MatrixBackend):
+    @cached_property
+    def db(self):
+        from .db import MatrixDb
+
+        return MatrixDb.from_path(self.path_or_url)
+
+
+class RemoteMatrixBackend(MatrixBackend):
+    @cached_property
+    def db(self):
+        from .db import MatrixDb
+
+        return MatrixDb.from_url(self.path_or_url)
+
+
+class SystemRemoteMatrixBackend(RemoteMatrixBackend):
+    @cached_property
+    def db(self):
+        from .db import SYS_DB
+
+        return SYS_DB

src/earthkit/regrid/backends/mir.py

@@ -0,0 +1,20 @@
+# (C) Copyright 2023 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+from . import Backend
+
+
+class MirBackend(Backend):
+    def interpolate(self, values, in_grid, out_grid, method, **kwargs):
+        try:
+            import mir
+        except ImportError:
+            raise ImportError("The 'mir' package is required for this operation")
+
+        return mir.interpolate(values, in_grid, out_grid, method, **kwargs)