Add torch default dtype constant

Author: sgreenbury

autoemulate/experimental/data/utils.py

@@ -6,6 +6,7 @@
     InputLike,
     OutputLike,
     TensorLike,
+    TorchDefaultDType,
     TorchScalarDType,
 )
 from sklearn.utils.validation import check_X_y
@@ -28,9 +29,9 @@ def _convert_to_dataset(
         """
         # Convert input to Dataset if not already
         if isinstance(x, np.ndarray):
-            x = torch.tensor(x, dtype=torch.float32)
+            x = torch.tensor(x, dtype=TorchDefaultDType)
         if isinstance(y, np.ndarray):
-            y = torch.tensor(y, dtype=torch.float32)
+            y = torch.tensor(y, dtype=TorchDefaultDType)
 
         if isinstance(x, torch.Tensor | np.ndarray) and isinstance(
             y, torch.Tensor | np.ndarray
@@ -77,7 +78,7 @@ def _convert_to_tensors(
         cls,
         x: InputLike,
         y: InputLike | None = None,
-        dtype: torch.dtype = torch.float32,
+        dtype: torch.dtype = TorchDefaultDType,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         """
         Convert InputLike x, y to Tensor or tuple of Tensors.

autoemulate/experimental/emulators/polynomials.py

@@ -56,7 +56,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         # Transform input using the fitted PolynomialFeatures
         x_np, _ = self._convert_to_numpy(x)
         x_poly = self.poly.transform(x_np)
-        x_poly_tensor = torch.tensor(x_poly, dtype=torch.float32, device=self.device)
+        x_poly_tensor = torch.tensor(x_poly, dtype=x.dtype, device=self.device)
         return self.linear(x_poly_tensor)
 
     @staticmethod

autoemulate/experimental/learners/base.py

@@ -11,7 +11,7 @@
 from autoemulate.experimental.emulators.base import Emulator
 from autoemulate.experimental.simulations.base import Simulator
 
-from ..types import GaussianLike, TensorLike
+from ..types import GaussianLike, TensorLike, TorchDefaultDType
 
 
 @dataclass(kw_only=True)
@@ -40,6 +40,7 @@ class Learner(ValidationMixin, ABC):
     y_train: TensorLike
     in_dim: int = field(init=False)
     out_dim: int = field(init=False)
+    dtype: torch.dtype = field(default=TorchDefaultDType)
 
     def __post_init__(self):
         """
@@ -199,8 +200,8 @@ def summary(self):
                   cumulative number of queries.
         """
         # pull histories into float tensors
-        mse = torch.tensor(self.metrics["mse"], dtype=torch.float32)
-        q = torch.tensor(self.metrics["n_queries"], dtype=torch.float32)
+        mse = torch.tensor(self.metrics["mse"], dtype=self.dtype)
+        q = torch.tensor(self.metrics["n_queries"], dtype=self.dtype)
 
         # build per-query ratios safely (avoid zero division)
         d = {}

autoemulate/experimental/simulations/base.py

@@ -6,7 +6,7 @@
 
 from autoemulate.experimental.data.utils import ValidationMixin, set_random_seed
 from autoemulate.experimental.logging_config import configure_logging
-from autoemulate.experimental.types import TensorLike
+from autoemulate.experimental.types import TensorLike, TorchDefaultDType
 from autoemulate.experimental_design import LatinHypercube
 
 logger = logging.getLogger("autoemulate")
@@ -101,7 +101,10 @@ def out_dim(self) -> int:
         return self._out_dim
 
     def sample_inputs(
-        self, n_samples: int, random_seed: int | None = None
+        self,
+        n_samples: int,
+        random_seed: int | None = None,
+        dtype: torch.dtype = TorchDefaultDType,
     ) -> TensorLike:
         """
         Generate random samples using Latin Hypercube Sampling.
@@ -123,10 +126,8 @@ def sample_inputs(
             set_random_seed(random_seed)  # type: ignore PGH003
         lhd = LatinHypercube(self.param_bounds)
         sample_array = lhd.sample(n_samples)
-        # TODO: have option to set dtype and ensure consistency throughout codebase?
-        # added here as method was returning float64 and elsewhere had tensors of
-        # float32 and this caused issues
-        return torch.tensor(sample_array, dtype=torch.float32)
+
+        return torch.tensor(sample_array, dtype=dtype)
 
     @abstractmethod
     def _forward(self, x: TensorLike) -> TensorLike:

autoemulate/experimental/simulations/epidemic.py

@@ -1,7 +1,7 @@
 import torch
 
 from autoemulate.experimental.simulations.base import Simulator
-from autoemulate.experimental.types import TensorLike
+from autoemulate.experimental.types import TensorLike, TorchDefaultDType
 from autoemulate.simulations.epidemic import simulate_epidemic
 
 
@@ -10,11 +10,7 @@ class Epidemic(Simulator):
     Simulator of infectious disease spread (SIR).
     """
 
-    def __init__(
-        self,
-        param_ranges=None,
-        output_names=None,
-    ):
+    def __init__(self, param_ranges=None, output_names=None):
         if param_ranges is None:
             param_ranges = {"beta": (0.1, 0.5), "gamma": (0.01, 0.2)}
         if output_names is None:
@@ -36,5 +32,4 @@ def _forward(self, x: TensorLike) -> TensorLike:
             Peak infection rate.
         """
         y = simulate_epidemic(x.cpu().numpy()[0])
-        # TODO (#537): update with default dtype
-        return torch.tensor([y], dtype=torch.float32).view(-1, 1)
+        return torch.tensor([y], dtype=TorchDefaultDType).view(-1, 1)

autoemulate/experimental/types.py

@@ -20,4 +20,12 @@
 DeviceLike: TypeAlias = str | torch.device
 
 # Torch dtype's
-TorchScalarDType = (torch.float32, torch.float64, torch.int32, torch.int64)
+TorchScalarDType: tuple[torch.dtype, ...] = (
+    torch.float32,
+    torch.float64,
+    torch.int32,
+    torch.int64,
+)
+
+# Default torch dtype (float32)
+TorchDefaultDType: torch.dtype = torch.float32

tests/experimental/conftest.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 import torch
+from autoemulate.experimental.types import TorchDefaultDType
 from sklearn.datasets import make_regression
 
 N_S = 20
@@ -115,7 +116,7 @@ def noisy_data():
     Generate a highly noisy dataset to test stochasticity effects.
     """
     rng = np.random.RandomState(0)
-    x = torch.tensor(rng.normal(size=(100, 2)), dtype=torch.float32)
-    y = torch.tensor(rng.normal(size=(100,)), dtype=torch.float32)
+    x = torch.tensor(rng.normal(size=(100, 2)), dtype=TorchDefaultDType)
+    y = torch.tensor(rng.normal(size=(100,)), dtype=TorchDefaultDType)
     x2 = x[:4].clone()
     return x, y, x2

tests/experimental/test_experimental_base_simulator.py

@@ -1,17 +1,15 @@
 import pytest
 import torch
 from autoemulate.experimental.simulations.base import Simulator
-from autoemulate.experimental.types import TensorLike
+from autoemulate.experimental.types import TensorLike, TorchDefaultDType
 from torch import Tensor
 
 
 class MockSimulator(Simulator):
     """Mock implementation of Simulator for testing purposes"""
 
     def __init__(
-        self,
-        parameters_range: dict[str, tuple[float, float]],
-        output_names: list[str],
+        self, parameters_range: dict[str, tuple[float, float]], output_names: list[str]
     ):
         # Call parent constructor
         super().__init__(parameters_range, output_names)
@@ -90,7 +88,7 @@ def test_sample_inputs(mock_simulator):
 def test_forward(mock_simulator):
     """Test that forward method works correctly"""
     # Create test input
-    test_input = torch.tensor([[0.5, 0.0, 7.5]], dtype=torch.float32)
+    test_input = torch.tensor([[0.5, 0.0, 7.5]], dtype=TorchDefaultDType)
 
     # Get output
     output = mock_simulator.forward(test_input)
@@ -109,7 +107,7 @@ def test_forward_batch(mock_simulator):
     # Create test batch
     n_samples = 3
     batch = torch.tensor(
-        [[0.5, 0.0, 7.5], [0.2, 0.3, 6.0], [0.8, -0.5, 9.0]], dtype=torch.float32
+        [[0.5, 0.0, 7.5], [0.2, 0.3, 6.0], [0.8, -0.5, 9.0]], dtype=TorchDefaultDType
     )
 
     # Process batch
@@ -188,7 +186,7 @@ def _forward(self, x: TensorLike) -> TensorLike:
             [0.1, 0.5, 0.5],  # Below threshold
             [0.7, 0.5, 0.5],  # Above threshold
         ],
-        dtype=torch.float32,
+        dtype=TorchDefaultDType,
     )
 
     # This should process all samples without errors

tests/experimental/test_experimental_preprocessors.py

@@ -1,6 +1,7 @@
 import pytest
 import torch
 from autoemulate.experimental.data.preprocessors import Standardizer
+from autoemulate.experimental.types import TorchDefaultDType
 
 
 class TestStandardizer:
@@ -37,7 +38,7 @@ def test_small_std_replaced(self):
         Test that small std values are replaced with 1.0.
         """
         mean = torch.zeros((1, 2))
-        std = torch.tensor([[1e-40, 2.0]], dtype=torch.float32)
+        std = torch.tensor([[1e-40, 2.0]], dtype=TorchDefaultDType)
         s = Standardizer(mean, std)
         assert s.std[0, 0] == 1.0
         assert s.std[0, 1] == 2.0