surrogates_interface 4.0.2__tar.gz → 4.1.0__tar.gz

{surrogates_interface-4.0.2 → surrogates_interface-4.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: surrogates_interface
-Version: 4.0.2
+Version: 4.1.0
 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.2 → surrogates_interface-4.1.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "surrogates_interface"
-version = "4.0.2"
+version = "4.1.0"
 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.2 → surrogates_interface-4.1.0}/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,172 @@ class SMTModel(SurrogateModel):
 
 class SurrogateModelFamily:
     pass
+
+
+class PyTorchModel(SurrogateModel):
+    import torch
+    import torch._dynamo  # fmt: skip
+    torch._dynamo.config.suppress_errors = True
+    import torch.func as torch_func  # fmt: skip
+
+    # from torch.func import jacfwd, jacrev, jvp, vjp, vmap  # fmt: skip
+
+    def __init__(self, *args, dtype=torch.float16, max_batch=50_000, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.model.eval()
+        self.device = self.torch.device(
+            "cuda" if self.torch.cuda.is_available() else "cpu"
+        )
+        self.model.to(self.device)
+        self.dtype = dtype
+        self.model.to(dtype)
+        self.max_batch = max_batch
+
+    def __to_numpy(self, _x):
+        try:
+            import cupy as cp
+        except ImportError:  # pragma: no cover
+            cp = np
+        lalib = (
+            cp if self.device == self.torch.device("cuda") else np
+        )  # pragma: no cover
+        return lalib.asarray(_x).astype(lalib.float32)
+
+    def __to_tensor(self, _x):
+        if len(_x.shape) < 2:
+            _x = _x.reshape(1, -1)
+        return self.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):
+        x = self.__to_tensor(x)
+
+        with self.torch.no_grad():
+            output = self.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):
+        self.torch.save(self.model.state_dict(), filepath)
+
+    @staticmethod
+    def _load_model(filepath, model_cls):
+        import torch  # fmt: skip
+        model_cls.load_state_dict(torch.load(filepath, weights_only=True))
+        return model_cls
+
+    def _predict_model_jvp(self, x, dx):
+        if not hasattr(self, "batched_jvp_func"):
+            def g(_x):  # fmt: skip
+                return self.model(_x)
+            def single_jvp(_x, _v):  # fmt: skip
+                return self.torch_func.jvp(g, (_x,), (_v,))
+            self.batched_jvp_func = self.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):
+        if not hasattr(self, "batched_vjp_func"):
+            def g(_x):  # fmt: skip
+                return self.model(_x)
+            def single_vjp(_x, _v):  # fmt: skip
+                out, f = self.torch_func.vjp(g, _x)
+                return out, f(_v)
+            self.batched_vjp_func = self.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):
+        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 = self.torch_func.vmap(
+                self.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):
+        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 = self.torch_func.vmap(
+                self.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."""
+        if not hasattr(self, "y_jac_func"):
+            def g(_x):  # fmt: skip
+                result = self.model(_x)
+                return (result, result)
+            self.y_jac_func = self.torch_func.vmap(
+                self.torch_func.jacrev(g, has_aux=True)
+            )
+
+        x = self.__to_tensor(x)
+        batch_size = x.shape[0]
+        y = self.torch.zeros(
+            (batch_size, self.n_outputs),
+            device=self.device,
+            dtype=self.dtype,
+        )
+        jac = self.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.2 → surrogates_interface-4.1.0}/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

surrogates_interface-4.0.2/surrogates_interface/test/test_domains.py

@@ -1,44 +0,0 @@
-# -*- coding: utf-8 -*-
-import numpy as np
-import numpy.testing as npt
-import pytest
-
-from surrogates_interface.domains import BoxDomain, GenericDomain
-
-
-class TestBoxDomain:
-    @staticmethod
-    def test_in_domain():
-        min = np.array([-1.0, +2.0, -15.0])
-        max = np.array([+1.0, +7.0, +8.0])
-        x = np.array(
-            [[0.0, 0.0, 0.0], [0.0, 4.0, 9.0], [0.0, 4.0, 0.0], [3.0, 4.0, 6.0]]
-        )
-        box = BoxDomain(min, max)
-        out = box.in_domain(x)
-        npt.assert_array_equal(out, np.array([False, False, True, False]))
-
-    @staticmethod
-    def test_bad_box():
-        min = np.array([-1.0, +7.0])
-        max = np.array([+1.0, -15.0])
-        with pytest.raises(AssertionError):
-            BoxDomain(min, max)
-
-    @staticmethod
-    def test_eq():
-        box1 = BoxDomain([-1.0, +2.0, -15.0], [+1.0, +7.0, +8.0])
-        box2 = BoxDomain([-1.0, -5.0, -15.0], [+1.0, +7.0, +8.0])
-        assert not box1 == box2
-
-    @staticmethod
-    def test_eq_different_type():
-        class DummyDomain(GenericDomain):
-            fields = []
-
-            def in_domain():
-                pass
-
-        box1 = BoxDomain([-1.0, +2.0, -15.0], [+1.0, +7.0, +8.0])
-        box2 = DummyDomain()
-        assert not box1 == box2

surrogates_interface-4.0.2/surrogates_interface/test/test_init.py

@@ -1,12 +0,0 @@
-try:
-    import tomllib  # From Python 3.11.
-except ModuleNotFoundError:
-    import tomli as tomllib  # Before Python 3.11.
-
-import surrogates_interface
-
-
-def test_version():
-    with open("pyproject.toml", "rb") as fid:
-        data = tomllib.load(fid)
-    assert surrogates_interface.__version__ == data["tool"]["poetry"]["version"]

surrogates_interface-4.0.2/surrogates_interface/test/test_pca_jacobian.py

@@ -1,137 +0,0 @@
-# -*- coding: utf-8 -*-
-
-import numpy as np
-import numpy.testing as npt
-import openmdao.api as om
-import pytest
-from sklearn.decomposition import PCA
-from sklearn.preprocessing import MinMaxScaler, PowerTransformer
-
-from surrogates_interface.openmdao import InputOutputType, SurrogateModelComp
-from surrogates_interface.surrogates import SurrogateModel
-from surrogates_interface.transformers import Transformer
-
-
-class Paraboloid(SurrogateModel):
-    """Define a paraboloid."""
-
-    coef_1 = np.array([0.2, 0.5, 0.3])
-    coef_2 = np.array([1.2, 1.5, 1.3])
-
-    def _predict_model_output(self, x):
-        y = np.column_stack(
-            (np.sum(self.coef_1 * x**2, axis=1), np.sum(self.coef_2 * x**2, axis=1))
-        )
-        return y
-
-    def _predict_model_jacobian(self, x):
-        y = self._predict_model_output(x)
-        jac = np.stack((2.0 * self.coef_1 * x, 2.0 * self.coef_2 * x), axis=1)
-        return y, jac
-
-
-@pytest.fixture
-def make_data():
-    rng = np.random.default_rng(seed=123)
-    n_points = 100
-    n_inputs_original = 5
-    n_outputs_original = 4
-    input = rng.uniform(-1.0, +1.0, size=(n_points, n_inputs_original))
-    # Introduce linear correlation between some input features.
-    input[:, 1] = 2.0 * input[:, 0]
-    input[:, 2] = 3.0 * input[:, 0] - 5.0
-    # Build input transformers.
-    sklearn_input_transformers = [
-        PowerTransformer(standardize=False),
-        PCA(n_components=3, whiten=False),
-    ]
-    input_transformed = input.copy()
-    for transformer in sklearn_input_transformers:
-        input_transformed = transformer.fit_transform(input_transformed)
-    si_input_transformers = [
-        Transformer.from_sklearn(t) for t in sklearn_input_transformers
-    ]
-    # Make a dummy output. We only need it to create the transformers.
-    output = np.zeros((n_points, n_outputs_original))
-    output[:, 0] = np.sin(input[:, 0]) * np.cos(input[:, 4])
-    output[:, 1] = np.sum(input**2, axis=1)
-    # Introduce linear correlation between some output features.
-    output[:, 2] = 1.1 * output[:, 0]
-    output[:, 3] = 0.5 * output[:, 1]
-    # Build output transformers.
-    sklearn_output_transformers = [
-        MinMaxScaler(feature_range=(-1.0, +1.0), clip=False),
-        PCA(n_components=2, whiten=False),
-    ]
-    output_transformed = output.copy()
-    for transformer in sklearn_output_transformers:
-        output_transformed = transformer.fit_transform(output_transformed)
-    si_output_transformers = [
-        Transformer.from_sklearn(t) for t in sklearn_output_transformers
-    ]
-    # Build the surrogate model.
-    si_model = Paraboloid(
-        [],  # model is dummy.
-        input_transformers=si_input_transformers,
-        output_transformers=si_output_transformers,
-        n_inputs=n_inputs_original,
-        n_outputs=n_outputs_original,
-    )
-    # Evaluate it on the input, and overwrite the transformed output.
-    output_transformed = si_model._predict_model_output(input_transformed)
-    # Apply inverse output transformations and overwrite the original output.
-    output = output_transformed.copy()
-    for transformer in reversed(si_output_transformers):
-        output = transformer.inverse_transform(output, inplace=False)
-    return (
-        si_model,
-        input,
-        output,
-    )
-
-
-def test_jacobian(make_data):
-    si_model, input, _ = make_data
-    # Compute exact Jacobian.
-    jacobian_si = si_model.predict_jacobian(input)
-    # Compute Jacobian using central finite differences.
-    jacobian_fd = np.zeros_like(jacobian_si)
-    h = 1e-6
-    for i in range(input.shape[1]):
-        xm = input.copy()
-        xp = input.copy()
-        xm[:, i] -= h
-        xp[:, i] += h
-        jacobian_fd[:, :, i] = si_model.predict_output(xp) - si_model.predict_output(xm)
-    jacobian_fd /= 2.0 * h
-    # Test.
-    npt.assert_allclose(jacobian_si, jacobian_fd, rtol=2e-7)
-
-
-def test_om_compute_partials_joined_joined(make_data):
-    # Get the data.
-    si_model, input, _ = make_data
-    # Make the OpenMDAO problem.
-    problem = om.Problem()
-    component = SurrogateModelComp(
-        model=si_model,
-        input_type=InputOutputType.JOINED,
-        output_type=InputOutputType.JOINED,
-        n_points=input.shape[0],
-    )
-    problem.model.add_subsystem(
-        "surrogate", component, promotes_inputs=["x"], promotes_outputs=["y"]
-    )
-    problem.setup()
-    # Set the input.
-    problem.set_val("x", input)
-    # Check partial derivatives.
-    om_partials = problem.check_partials(
-        out_stream=None, step=1e-6, step_calc="abs", form="central"
-    )
-    # Jacobian computed by finite differences.
-    jacobian_fd = om_partials["surrogate"][("y", "x")]["J_fd"]
-    # Exact Jacobian.
-    jacobian_exact = om_partials["surrogate"][("y", "x")]["J_fwd"]
-    # Test.
-    npt.assert_allclose(jacobian_fd, jacobian_exact, rtol=2e-7)