Enhance the capability of Bridgescaler for supporting tensors

bridgescaler/__init__.py

@@ -1,4 +1,51 @@
+from importlib.metadata import version, PackageNotFoundError
+
+from packaging.version import Version
+
+
+# 1. PyTorch Checks
+REQUIRED_TORCH_VERSION = Version("2.0.0")
+
+def get_torch_status() -> tuple[bool, Version | None]:
+    try:
+        return True, Version(version("torch"))
+    except PackageNotFoundError:
+        return False, None
+
+TORCH_AVAILABLE, TORCH_VERSION = get_torch_status()
+
+def require_torch() -> None:
+    if not TORCH_AVAILABLE:
+        raise ImportError("PyTorch is required but not installed")
+    if TORCH_VERSION < REQUIRED_TORCH_VERSION:
+        raise RuntimeError(
+            f"PyTorch >= {REQUIRED_TORCH_VERSION} required; found {TORCH_VERSION}"
+        )
+
+# 2. Base Imports
 from .backend import save_scaler, load_scaler, print_scaler, read_scaler
 from .group import GroupStandardScaler, GroupRobustScaler, GroupMinMaxScaler
 from .deep import DeepStandardScaler, DeepMinMaxScaler, DeepQuantileTransformer
-from .distributed import DStandardScaler, DMinMaxScaler, DQuantileScaler
+from .distributed import (DStandardScaler, DMinMaxScaler, DQuantileScaler)
+
+# 3. Conditional Torch Imports
+if TORCH_AVAILABLE:
+    try: # Ensure that no errors are raised if PyTorch is installed but does not meet the required version.
+        from .distributed_tensor import (
+            DStandardScalerTensor,
+            DMinMaxScalerTensor,
+        )
+        from .backend_tensor import print_scaler_tensor, read_scaler_tensor
+    except:
+        pass
+
+# 4. Define Public API
+__all__ = [
+    # Utilities
+    "save_scaler", "load_scaler", "print_scaler", "read_scaler",
+    "TORCH_AVAILABLE",
+    # Scalers
+    "GroupStandardScaler", "GroupRobustScaler", "GroupMinMaxScaler",
+    "DeepStandardScaler", "DeepMinMaxScaler", "DeepQuantileTransformer",
+    "DStandardScaler", "DMinMaxScaler", "DQuantileScaler",
+]

bridgescaler/backend.py

@@ -156,4 +156,4 @@ def create_synthetic_data():
     for l in range(locs.shape[0]):
         x_data_dict[names[l]] = np.random.normal(loc=locs[l], scale=scales[l], size=num_examples)
     x_data = pd.DataFrame(x_data_dict)
-    return x_data
+    return x_data

bridgescaler/backend_tensor.py

@@ -0,0 +1,43 @@
+from . import require_torch
+require_torch()   # enforce torch availability/version at import time
+import torch
+
+import json
+import numpy as np
+
+from bridgescaler.distributed_tensor import DStandardScalerTensor, DMinMaxScalerTensor
+from .backend import NumpyEncoder, object_hook
+
+
+scaler_objs = {"DStandardScalerTensor": DStandardScalerTensor,
+               "DMinMaxScalerTensor": DMinMaxScalerTensor,
+               }
+
+
+def print_scaler_tensor(scaler):
+    """
+    Modify the print_scaler() in backend.py for tensors.
+    """
+    scaler_params = scaler.__dict__
+    scaler_params["type"] = str(type(scaler))[1:-2].split(".")[-1]
+
+    for keys in scaler_params:
+        if type(scaler_params[keys]) == torch.Tensor:
+            scaler_params[keys] = scaler_params[keys].cpu().numpy().copy()
+
+    return json.dumps(scaler_params, indent=4, sort_keys=True, cls=NumpyEncoder)
+
+
+def read_scaler_tensor(scaler_str):
+    """
+    Modify the read_scaler() in backend.py for tensors.
+    """
+    scaler_params = json.loads(scaler_str, object_hook=object_hook)
+    scaler = scaler_objs[scaler_params["type"]]()
+    del scaler_params["type"]
+    for k, v in scaler_params.items():
+        if k == "x_columns_":
+            setattr(scaler, k, v)
+        else:
+            setattr(scaler, k, torch.tensor(v))
+    return scaler