surrogates_interface 4.0.3__tar.gz → 4.1.1__tar.gz

{surrogates_interface-4.0.3 → surrogates_interface-4.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: surrogates_interface
-Version: 4.0.3
+Version: 4.1.1
 Summary: Interface for surrogate models.
 Home-page: https://gitlab.windenergy.dtu.dk/surrogate-models/surrogate-models-interface
 License: MIT
@@ -19,6 +19,7 @@ Provides-Extra: om
 Provides-Extra: smt
 Provides-Extra: test
 Provides-Extra: tf
+Provides-Extra: torch
 Requires-Dist: black ; extra == "code-style"
 Requires-Dist: codespell ; extra == "code-style"
 Requires-Dist: flake8 ; extra == "code-style"
@@ -43,6 +44,7 @@ Requires-Dist: tensorflow-intel (>=2.10.0,<3.0.0) ; (sys_platform == "win32") an
 Requires-Dist: tensorflow-io-gcs-filesystem (<0.32.0) ; (python_version >= "3.9" and python_version < "3.12" and platform_system == "Windows") and (extra == "test" or extra == "tf")
 Requires-Dist: tensorflow-io-gcs-filesystem (>=0.23.1) ; (platform_machine != "arm64" and python_version >= "3.9" and python_version < "3.12" or platform_system != "Darwin" and python_version >= "3.9" and python_version < "3.12") and (extra == "test" or extra == "tf")
 Requires-Dist: tomli (>=2.0.1,<3.0.0) ; (python_version < "3.11") and (extra == "test")
+Requires-Dist: torch (>=2.0.0,<3.0.0) ; extra == "test" or extra == "torch"
 Project-URL: Repository, https://gitlab.windenergy.dtu.dk/surrogate-models/surrogate-models-interface
 Description-Content-Type: text/markdown
 

{surrogates_interface-4.0.3 → surrogates_interface-4.1.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "surrogates_interface"
-version = "4.0.3"
+version = "4.1.1"
 description = "Interface for surrogate models."
 authors = ["DTU Wind and Energy Systems"]
 maintainers = [
@@ -45,6 +45,7 @@ tensorflow-io-gcs-filesystem = [
 tensorflow-intel = {version="^2.10.0", platform = "win32", optional=true}
 smt = {version="^2.0.1", optional=true}
 openmdao = {version="^3.25", optional=true}
+torch = { version = "^2.0.0", optional=true }
 
 [tool.poetry.extras]
 
@@ -74,6 +75,7 @@ test = ["matplotlib",  # Needed by SMT.
 	    "tensorflow-io-gcs-filesystem",
 	    "tensorflow-intel",
 	    "openmdao",
+		"torch",
 	    ]
 
 smt = ["matplotlib",  # Needed by SMT.
@@ -89,6 +91,9 @@ tf = ["scikit-learn",
 
 om = ["openmdao"]
 
+torch = ["torch"]
+
+
 
 [build-system]
 build-backend = "poetry.core.masonry.api"

{surrogates_interface-4.0.3 → surrogates_interface-4.1.1}/surrogates_interface/surrogates.py

@@ -452,6 +452,8 @@ class SurrogateModel:
         x_transformed = x.copy()
         # The input transformation Jacobian starts as an identity matrix,
         # repeated over all points.
+        if x.ndim == 1:
+            x = x.reshape(1, -1)
         input_transformation_jacobian = np.tile(np.eye(x.shape[1]), (x.shape[0], 1, 1))
         for transformer in self.input_transformers:
             # Get Jacobian of current input transformation.
@@ -736,9 +738,9 @@ class TensorFlowModel(SurrogateModel):
         IOError
             In case of an invalid savefile.
         """
-        from tensorflow.keras.models import load_model
+        import tensorflow as tf
 
-        return load_model(filepath, custom_objects, compile, options)
+        return tf.keras.models.load_model(filepath, custom_objects, compile, options)
 
     def _save_model(self, filepath, overwrite=True, save_format="h5"):
         """
@@ -1025,3 +1027,187 @@ class SMTModel(SurrogateModel):
 
 class SurrogateModelFamily:
     pass
+
+
+class PyTorchModel(SurrogateModel):
+    """
+    Class that stores and evaluates an existing PyTorch model.
+    """
+
+    def __init__(self, *args, dtype=None, max_batch=50_000, **kwargs):
+        import torch
+        import torch._dynamo
+
+        torch._dynamo.config.suppress_errors = True
+        super().__init__(*args, **kwargs)
+        self.model.eval()
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device)
+        self.dtype = torch.float16 if dtype is None else dtype
+        self.model.to(self.dtype)
+        self.max_batch = max_batch
+
+    def __to_numpy(self, _x):
+        import torch
+
+        try:
+            import cupy as cp
+        except ImportError:  # pragma: no cover
+            cp = np
+        lalib = cp if self.device == torch.device("cuda") else np  # pragma: no cover
+        return lalib.asarray(_x).astype(lalib.float32)
+
+    def __to_tensor(self, _x):
+        import torch
+
+        if len(_x.shape) < 2:
+            _x = _x.reshape(1, -1)
+        return torch.as_tensor(_x, device=self.device, dtype=self.dtype)
+
+    def __record_in_dtype(self, _x):  # pragma: no cover
+        # TODO: could be used to retain the input/output dtype the same
+        self.in_dtype = _x.dtype
+
+    def _predict_model_output(self, x):
+        import torch
+
+        x = self.__to_tensor(x)
+
+        with torch.no_grad():
+            output = torch.zeros(
+                (x.shape[0], self.n_outputs),
+                device=self.device,
+                dtype=self.dtype,
+            )
+            for i in range(0, x.shape[0], self.max_batch):
+                inter_out = self.model(x[i : i + self.max_batch])
+                output[i : i + self.max_batch].copy_(
+                    inter_out[: x[i : i + self.max_batch].shape[0]], non_blocking=True
+                )
+
+        return self.__to_numpy(output)
+
+    def _predict_model_jacobian(self, x, ad_mode=None):
+        if ad_mode is None:
+            ad_mode = ADMode.REV if (self.n_outputs <= self.n_inputs) else ADMode.FWD
+
+        if ad_mode is ADMode.REV:
+            y, jacobian = self._predict_model_output_and_jacobian_rev(x)
+        elif ad_mode is ADMode.FWD:
+            y, jacobian = self._predict_model_output_and_jacobian_fwd(x)
+
+        return y, jacobian
+
+    def predict_output(self, x):
+        y = super().predict_output(x)
+        return self.__to_numpy(y)
+
+    def _save_model(self, filepath):
+        import torch
+
+        torch.save(self.model.state_dict(), filepath)
+
+    @staticmethod
+    def _load_model(filepath, model_cls):
+        import torch
+
+        model_cls.load_state_dict(torch.load(filepath, weights_only=True))
+        return model_cls
+
+    def _predict_model_jvp(self, x, dx):
+        import torch.func as torch_func
+
+        if not hasattr(self, "batched_jvp_func"):
+            def g(_x):  # fmt: skip
+                return self.model(_x)
+
+            def single_jvp(_x, _v):
+                return torch_func.jvp(g, (_x,), (_v,))
+
+            self.batched_jvp_func = torch_func.vmap(single_jvp)
+
+        x = self.__to_tensor(x)
+        dx = self.__to_tensor(dx)
+        output, jvp_res = self.batched_jvp_func(x, dx)
+        return self.__to_numpy(output.detach()), self.__to_numpy(jvp_res.detach())
+
+    def _predict_model_vjp(self, x, dy):
+        import torch.func as torch_func
+
+        if not hasattr(self, "batched_vjp_func"):
+            def g(_x):  # fmt: skip
+                return self.model(_x)
+
+            def single_vjp(_x, _v):
+                out, f = torch_func.vjp(g, _x)
+                return out, f(_v)
+
+            self.batched_vjp_func = torch_func.vmap(single_vjp)
+
+        x = self.__to_tensor(x)
+        dy = self.__to_tensor(dy)
+        output, vjp_res = self.batched_vjp_func(x, dy)
+        return self.__to_numpy(output.detach()), self.__to_numpy(vjp_res[0].detach())
+
+    def _predict_model_output_and_jacobian_fwd(self, x):
+        import torch.func as torch_func
+
+        if not hasattr(self, "y_jac_fwd_func"):
+            def g(_x):  # fmt: skip
+                result = self.model(_x)
+                return (result, result)
+            self.y_jac_fwd_func = torch_func.vmap(torch_func.jacfwd(g, has_aux=True))
+
+        x = self.__to_tensor(x)
+        jac, output = self.y_jac_fwd_func(x)
+        return self.__to_numpy(output.detach()), self.__to_numpy(jac.detach())
+
+    def _predict_model_output_and_jacobian_rev(self, x):
+        import torch.func as torch_func
+
+        if not hasattr(self, "y_jac_rev_func"):
+            def g(_x):  # fmt: skip
+                result = self.model(_x)
+                return (result, result)
+            self.y_jac_rev_func = torch_func.vmap(torch_func.jacrev(g, has_aux=True))
+
+        x = self.__to_tensor(x)
+        jac, output = self.y_jac_rev_func(x)
+        return self.__to_numpy(output.detach()), self.__to_numpy(jac.detach())
+
+    def _predict_model_jacobian_sbatch(self, x):  # pragma: no cover
+        # Not used for now, but left as a reference for future if GPU memory issues arise.
+        import torch
+        import torch.func as torch_func
+
+        if not hasattr(self, "y_jac_func"):
+            def g(_x):  # fmt: skip
+                result = self.model(_x)
+                return (result, result)
+            self.y_jac_func = torch_func.vmap(torch_func.jacrev(g, has_aux=True))
+
+        x = self.__to_tensor(x)
+        batch_size = x.shape[0]
+        y = torch.zeros(
+            (batch_size, self.n_outputs),
+            device=self.device,
+            dtype=self.dtype,
+        )
+        jac = torch.zeros(
+            (batch_size, self.n_outputs, self.n_inputs),
+            device=self.device,
+            dtype=self.dtype,
+        )
+        for i in range(0, batch_size, self.max_batch):
+            x_mb = x[i : i + self.max_batch]
+            inter_jac, inter_y = self.y_jac_func(x_mb)
+            y[i : i + self.max_batch].copy_(
+                inter_y[: x_mb.shape[0]],
+                non_blocking=True,
+            )
+            jac[i : i + self.max_batch].copy_(
+                inter_jac[: x_mb.shape[0]],
+                non_blocking=True,
+            )
+
+        return self.__to_numpy(y.detach()), self.__to_numpy(jac.detach())

{surrogates_interface-4.0.3 → surrogates_interface-4.1.1}/surrogates_interface/transformers.py

@@ -814,3 +814,42 @@ class PowerTransformerYeoJohnson(Transformer):
         if isinstance(other, self.__class__):
             return np.array_equal(self.lambdas_, other.lambdas_)
         return False
+
+
+class StandardScalerTorch(StandardScaler):
+    import torch
+
+    def __init__(self, scale, mean):
+        super().__init__(scale, mean)
+        device = self.torch.device("cuda" if self.torch.cuda.is_available() else "cpu")
+        self.device = device
+        def __np2torch(v):  # fmt: skip
+            return self.torch.as_tensor(v).to(device)
+        self.gain_ = __np2torch(self.gain_)
+        self.inv_gain_ = __np2torch(self.inv_gain_)
+        self.offset_ = __np2torch(self.offset_)
+
+    def transform(self, x, inplace=True):
+        if inplace:
+            raise ValueError("No inplace...")
+        x = self.torch.as_tensor(x, device=self.device, dtype=self.torch.float32)
+        return super().transform(x, inplace)
+
+    def inverse_transform(self, x, inplace=True):
+        if inplace:
+            raise ValueError("No inplace...")
+        x = self.torch.as_tensor(x, device=self.device, dtype=self.torch.float32)
+        return super().inverse_transform(x, inplace)
+
+    def transform_gradient(self, x):
+        return self.torch.broadcast_to(self.gain_, x.shape).numpy()
+
+    def inverse_transform_gradient(self, x):
+        return self.torch.broadcast_to(self.inv_gain_, x.shape).numpy()
+
+    def __eq__(self, other):
+        if isinstance(other, self.__class__):
+            test_gain = self.torch.equal(self.gain_, other.gain_)
+            test_offset = self.torch.equal(self.offset_, other.offset_)
+            return test_gain and test_offset
+        return False