Experiments with Decorators

python/pyproject.toml

@@ -0,0 +1,26 @@
+[tool.poetry]
+name = "umbridge"
+version = "0.1.0"
+description =  "UM-Bridge (the UQ and Model Bridge) provides a unified interface for numerical models that is accessible from virtually any programming language or framework. It is primarily intended for coupling advanced models (e.g. simulations of complex physical processes) to advanced statistical or optimization methods."
+authors = ["Your Name <you@example.com>"]
+readme = "README.md"
+packages = [{include = "umbridge"}]
+
+[tool.poetry.dependencies]
+python = "^3.10"
+litestar = "^2.0.0a5"
+requests = "^2.30.0"
+numpy = "^1.24.3"
+pytensor = "^2.11.2"
+aiohttp = "^3.8.4"
+jax = "^0.4.8"
+jaxlib = "^0.4.7"
+
+
+[tool.poetry.group.dev.dependencies]
+black = "^23.3.0"
+pytest = "^7.3.1"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

python/tests/autoumbridge_test.py

@@ -0,0 +1,17 @@
+from umbridge.autoumbridge import autoUm
+from umbridge import serve_models
+
+@autoUm(use_jax=True)
+def f(x):
+    return x**2
+
+def test_wrapper():
+    assert f(2) == 4
+
+    # assert f.gradient(1) == 2
+
+
+
+
+if __name__ == "__main__":
+    serve_models([f])

python/umbridge/__init__.py

@@ -0,0 +1,4 @@
+# init file for umbridge
+from .model import Model
+from .server import serve_models
+from .client import HTTPModel, supported_models

python/umbridge/autoumbridge.py

@@ -0,0 +1,49 @@
+from .model import Model
+from functools import update_wrapper
+from jax import grad
+
+
+def autoUm(use_jax=True):
+    class UmbridgeModel(Model):
+        def __init__(self, wrapped_model):
+            super().__init__(wrapped_model.__name__)
+            self.wrapped_model = wrapped_model
+            update_wrapper(self, wrapped_model)
+
+        def __call__(self, *args, **kwargs):
+            return self.wrapped_model(*args, **kwargs)
+
+        def get_input_sizes(self):
+            return super().get_input_sizes()
+
+        def get_output_sizes(self):
+            return super().get_output_sizes()
+
+        if use_jax:
+            def gradient(self, out_wrt, in_wrt, parameters, sens, config={}):
+                return grad(self.wrapped_model)
+
+            def apply_jacobian(self, out_wrt, in_wrt, parameters, vec, config=...):
+                raise NotImplementedError(f"Method called but not implemented in {self.__class__.__name__}")
+
+            def apply_hessian(self, out_wrt, in_wrt1, in_wrt2, parameters, sens, vec, config={}):
+                raise NotImplementedError(f'Method called but not implemented in {self.__class__.__name__}.')
+
+            def supports_evaluate(self):
+                return True
+            def supports_gradient(self):
+                return True
+            def supports_apply_jacobian(self):
+                return True
+            def supports_apply_hessian(self):
+                return True
+
+    return UmbridgeModel
+
+if __name__ == "__main__":
+
+    @autoUm()
+    def f(x):
+        x**2
+
+    f(2)

python/umbridge/client.py

@@ -0,0 +1,128 @@
+import requests
+from .model import Model
+
+
+def supported_models(url):
+  response = requests.get(f"{url}/Info").json()
+  if (response["protocolVersion"] != 1.0):
+      raise RuntimeWarning("Model has unsupported protocol version!")
+  return response["models"]
+
+class HTTPModel(Model):
+    def __init__(self, url, name):
+        super().__init__(name)
+        self.url = url
+
+        if (name not in supported_models(url)):
+            raise Exception(f'Model {name} not supported by server! Supported models are: {supported_models(url)}')
+
+        input = {}
+        input["name"] = name
+        response = requests.post(f"{self.url}/ModelInfo", json=input).json()
+        self.__supports_evaluate = response["support"].get("Evaluate", False)
+        self.__supports_gradient = response["support"].get("Gradient", False)
+        self.__supports_apply_jacobian = response["support"].get("ApplyJacobian", False)
+        self.__supports_apply_hessian = response["support"].get("ApplyHessian", False)
+
+    def get_input_sizes(self, config={}):
+        input = {}
+        input["name"] = self.name
+        input["config"] = config
+        response = requests.post(f"{self.url}/InputSizes", json=input).json()
+        return response["inputSizes"]
+
+    def get_output_sizes(self, config={}):
+        input = {}
+        input["name"] = self.name
+        input["config"] = config
+        response = requests.post(f"{self.url}/OutputSizes", json=input).json()
+        return response["outputSizes"]
+
+    def supports_evaluate(self):
+        return self.__supports_evaluate
+
+    def supports_gradient(self):
+        return self.__supports_gradient
+
+    def supports_apply_jacobian(self):
+        return self.__supports_apply_jacobian
+
+    def supports_apply_hessian(self):
+        return self.__supports_apply_hessian
+
+    def __check_input_is_list_of_lists(self,parameters):
+        if not isinstance(parameters, list):
+            raise Exception("Parameters must be a list of lists!")
+        if not all(isinstance(x, list) for x in parameters):
+            raise Exception("Parameters must be a list of lists!")
+
+    def __call__(self, parameters, config={}):
+        if not self.supports_evaluate():
+            raise Exception('Evaluation not supported by model!')
+        self.__check_input_is_list_of_lists(parameters)
+
+        inputParams = {}
+        inputParams["name"] = self.name
+        inputParams["input"] = parameters
+        inputParams["config"] = config
+        response = requests.post(f"{self.url}/Evaluate", json=inputParams).json()
+
+        if "error" in response:
+            raise Exception(f'Model returned error of type {response["error"]["type"]}: {response["error"]["message"]}')
+        return response["output"]
+
+    def gradient(self, out_wrt, in_wrt, parameters, sens, config={}):
+        if not self.supports_gradient():
+            raise Exception('Gradient not supported by model!')
+        self.__check_input_is_list_of_lists(parameters)
+
+        inputParams = {}
+        inputParams["name"] = self.name
+        inputParams["outWrt"] = out_wrt
+        inputParams["inWrt"] = in_wrt
+        inputParams["input"] = parameters
+        inputParams["sens"] = sens
+        inputParams["config"] = config
+        response = requests.post(f"{self.url}/Gradient", json=inputParams).json()
+
+        if "error" in response:
+            raise Exception(f'Model returned error of type {response["error"]["type"]}: {response["error"]["message"]}')
+        return response["output"]
+
+    def apply_jacobian(self, out_wrt, in_wrt, parameters, vec, config={}):
+        if not self.supports_apply_jacobian():
+            raise Exception('ApplyJacobian not supported by model!')
+        self.__check_input_is_list_of_lists(parameters)
+
+        inputParams = {}
+        inputParams["name"] = self.name
+        inputParams["outWrt"] = out_wrt
+        inputParams["inWrt"] = in_wrt
+        inputParams["input"] = parameters
+        inputParams["vec"] = vec
+        inputParams["config"] = config
+        response = requests.post(f"{self.url}/ApplyJacobian", json=inputParams).json()
+
+        if "error" in response:
+            raise Exception(f'Model returned error of type {response["error"]["type"]}: {response["error"]["message"]}')
+        return response["output"]
+
+    def apply_hessian(self, out_wrt, in_wrt1, in_wrt2, parameters, sens, vec, config={}):
+        if not self.supports_apply_hessian():
+            raise Exception('ApplyHessian not supported by model!')
+        self.__check_input_is_list_of_lists(parameters)
+
+        inputParams = {}
+        inputParams["name"] = self.name
+        inputParams["outWrt"] = out_wrt
+        inputParams["inWrt1"] = in_wrt1
+        inputParams["inWrt2"] = in_wrt2
+        inputParams["input"] = parameters
+        inputParams["sens"] = sens
+        inputParams["vec"] = vec
+        inputParams["config"] = config
+        response = requests.post(f"{self.url}/ApplyHessian", json=inputParams).json()
+
+        if "error" in response:
+            raise Exception(f'Model returned error of type {response["error"]["type"]}: {response["error"]["message"]}')
+        return response["output"]

python/umbridge/model.py

@@ -0,0 +1,36 @@
+from abc import ABC, abstractmethod
+
+class Model(ABC):
+
+    def __init__(self, name):
+        self.name = name
+
+    @abstractmethod
+    def get_input_sizes(self):
+        return
+
+    @abstractmethod
+    def get_output_sizes(self):
+        return
+
+    @abstractmethod
+    def __call__(self, parameters, config={}):
+        return
+
+    def gradient(self, out_wrt, in_wrt, parameters, sens, config={}):
+        raise NotImplementedError(f'Method called but not implemented in {self.__class__.__name__}.')
+
+    def apply_jacobian(self, out_wrt, in_wrt, parameters, vec, config={}):
+        raise NotImplementedError(f'Method called but not implemented in {self.__class__.__name__}.')
+
+    def apply_hessian(self, out_wrt, in_wrt1, in_wrt2, parameters, sens, vec, config={}):
+        raise NotImplementedError(f'Method called but not implemented in {self.__class__.__name__}.')
+
+    def supports_evaluate(self):
+        return False
+    def supports_gradient(self):
+        return False
+    def supports_apply_jacobian(self):
+        return False
+    def supports_apply_hessian(self):
+        return False