[SYSTEMDS-3835] Add physiological representations to Scuro

src/main/python/systemds/scuro/dataloader/tabular_loader.py

@@ -0,0 +1,78 @@
+# -------------------------------------------------------------
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+# -------------------------------------------------------------
+from dataclasses import dataclass
+import numpy as np
+from typing import List, Optional, Union
+
+from systemds.scuro.dataloader.base_loader import BaseLoader
+from systemds.scuro.modality.type import ModalityType
+
+
+@dataclass
+class TabularStats:
+    num_instances: int
+    num_features: int
+    output_shape: tuple
+    output_shape_is_known: bool = True
+
+
+class TabularLoader(BaseLoader):
+    def __init__(
+        self,
+        source_path: str,
+        indices: List[str],
+        feature_names: Optional[List[str]] = None,
+        data_type: Union[np.dtype, str] = np.float32,
+        chunk_size: Optional[int] = None,
+        normalize: bool = False,
+        file_format: str = "npy",
+        modality_type: Optional[ModalityType] = ModalityType.EMBEDDING,
+    ):
+        super().__init__(source_path, indices, data_type, chunk_size, modality_type)
+        self.feature_names = feature_names
+        self.normalize = normalize
+        self.file_format = file_format.lower()
+        if self.file_format != "npy":
+            raise ValueError(f"Unsupported file format: {self.file_format}")
+        self.stats = self.get_stats(source_path)
+
+    def extract(self, file: str, index: Optional[Union[str, List[str]]] = None):
+        self.file_sanity_check(file)
+        data = np.load(file).astype(self._data_type, copy=False).reshape(-1)
+
+        if self.normalize:
+            mean = np.mean(data)
+            std = np.std(data)
+            data = (data - mean) / (std + 1e-8)
+
+        self.metadata.append(self.modality_type.create_metadata(data))
+        self.data.append(data)
+
+    def get_stats(self, source_path: str) -> TabularStats:
+        num_instances = 0
+        num_features = 0
+        for file_name in self.indices:
+            file = source_path + file_name + "." + self.file_format
+            self.file_sanity_check(file)
+            data = np.load(file)
+            num_features = max(num_features, int(np.prod(data.shape)))
+            num_instances += 1
+        return TabularStats(num_instances, num_features, (num_features,))

src/main/python/systemds/scuro/representations/aggregate.py

@@ -89,7 +89,7 @@ def __init__(self, aggregation_function="mean", pad_modality=True, params=None):
         self.aggregation_function_name = aggregation_function
 
         self.parameters = {
-            "aggregation_function": self._aggregation_function.keys(),
+            "aggregation_function": list(self._aggregation_function.keys()),
         }
 
     def get_current_parameters(self):

src/main/python/systemds/scuro/representations/average.py

@@ -18,12 +18,12 @@
 # under the License.
 #
 # -------------------------------------------------------------
-import copy
 from typing import List
 
 import numpy as np
 
 from systemds.scuro.modality.modality import Modality
+from systemds.scuro.representations.representation import RepresentationStats
 from systemds.scuro.representations.utils import pad_sequences
 
 from systemds.scuro.representations.fusion import Fusion
@@ -41,11 +41,38 @@ def __init__(self, params=None):
         self.associative = True
         self.commutative = True
 
-    def execute(self, modalities: List[Modality], labels=None):
-        data = np.asarray(copy.deepcopy(modalities[0].data), dtype=float)
+    def execute(self, modalities: List[Modality]):
+        data = np.array(modalities[0].data, dtype=np.float64)
         for i in range(1, len(modalities)):
-            data += np.asarray(modalities[i].data, dtype=float)
+            data += np.asarray(modalities[i].data, dtype=np.float64)
 
         data /= len(modalities)
 
-        return np.array(data)
+        return data
+
+    def get_output_stats(self, input_stats_list) -> RepresentationStats:
+        stats_list = self._fusion_input_stats(input_stats_list)
+        if not stats_list:
+            return RepresentationStats(0, (0,))
+
+        num_instances = max(s.num_instances for s in stats_list)
+        rank = len(stats_list[0].output_shape)
+        if rank > 0 and all(len(s.output_shape) == rank for s in stats_list):
+            output_shape = tuple(
+                max(s.output_shape[d] for s in stats_list) for d in range(rank)
+            )
+        else:
+            output_shape = max(stats_list, key=self._stats_num_elements).output_shape
+        output_shape_is_known = all(s.output_shape_is_known for s in stats_list)
+        return RepresentationStats(num_instances, output_shape, output_shape_is_known)
+
+    def estimate_peak_memory_bytes(self, input_stats) -> dict:
+        stats_list = self._as_stats_list(input_stats)
+        input_bytes = sum(self._stats_bytes(s) for s in stats_list)
+        output_bytes = self._stats_bytes(self.get_output_stats(input_stats))
+
+        raw_bytes = self._raw_input_bytes(input_stats)
+        cpu_peak = (
+            int((raw_bytes + input_bytes + 2 * output_bytes) * 1.15) + 8 * 1024 * 1024
+        )
+        return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": 0}

src/main/python/systemds/scuro/representations/bert.py

@@ -102,6 +102,7 @@ def get_output_stats(self, input_stats) -> RepresentationStats:
                 input_stats.num_instances,
                 (self.max_seq_length, 768),
                 aggregate_dim=(0,),
+                dtype=self.data_type,
             )
         else:
             self.stats = RepresentationStats(
@@ -111,6 +112,7 @@ def get_output_stats(self, input_stats) -> RepresentationStats:
                     0,
                     1,
                 ),
+                dtype=self.data_type,
             )
         if self.params and "_pushdown_aggregation" in self.params:
             output_shape = (768,)

src/main/python/systemds/scuro/representations/bow.py

@@ -18,6 +18,8 @@
 # under the License.
 #
 # -------------------------------------------------------------
+import os
+
 import numpy as np
 from sklearn.feature_extraction.text import CountVectorizer
 
@@ -30,6 +32,8 @@
 from systemds.scuro.drsearch.operator_registry import register_representation
 from systemds.scuro.dataloader.text_loader import TextStats
 
+_MAX_VOCAB_FEATURES = int(os.environ.get("SCURO_BOW_MAX_FEATURES", "100000"))
+
 
 @register_representation(ModalityType.TEXT)
 class BoW(UnimodalRepresentation):
@@ -43,14 +47,17 @@ def __init__(self, ngram_range=2, min_df=2, output_file=None, params=None):
 
     def get_output_stats(self, input_stats: TextStats) -> RepresentationStats:
         vocab_estimate = min(
-            100000,
+            _MAX_VOCAB_FEATURES,
             max(
                 1000,
                 input_stats.num_instances * input_stats.max_length * self.ngram_range,
             ),
         )
         return RepresentationStats(
-            input_stats.num_instances, (vocab_estimate,), output_shape_is_known=False
+            input_stats.num_instances,
+            (vocab_estimate,),
+            output_shape_is_known=False,
+            dtype=self.data_type,
         )
 
     def estimate_output_memory_bytes(self, input_stats: TextStats) -> int:
@@ -73,7 +80,10 @@ def estimate_peak_memory_bytes(self, input_stats: TextStats) -> dict:
     def transform(self, modality, aggregation=None):
         transformed_modality = TransformedModality(modality, self)
         vectorizer = CountVectorizer(
-            ngram_range=(1, self.ngram_range), min_df=self.min_df
+            ngram_range=(1, self.ngram_range),
+            min_df=self.min_df,
+            max_features=_MAX_VOCAB_FEATURES,
+            dtype=np.float32,
         )
 
         X = (

src/main/python/systemds/scuro/representations/clip.py

@@ -109,13 +109,17 @@ def get_output_stats(self, input_stats) -> RepresentationStats:
                     input_stats.max_length,
                     512,
                 ),
+                dtype=self.data_type,
             )
         elif not isinstance(input_stats, RepresentationStats):
-            return RepresentationStats(input_stats.num_instances, (512,))
+            return RepresentationStats(
+                input_stats.num_instances, (512,), dtype=self.data_type
+            )
         else:
             return RepresentationStats(
                 input_stats.num_instances,
                 (input_stats.output_shape[0], 512),
+                dtype=self.data_type,
             )
 
     def estimate_peak_memory_bytes(self, input_stats) -> dict:
@@ -394,7 +398,10 @@ def _get_parameters(self):
     def get_output_stats(self, input_stats) -> RepresentationStats:
         if not isinstance(input_stats, RepresentationStats):
             self.stats = RepresentationStats(
-                input_stats.num_instances, (512,), aggregate_dim=(0,)
+                input_stats.num_instances,
+                (512,),
+                aggregate_dim=(0,),
+                dtype=self.data_type,
             )
         else:
             self.stats = RepresentationStats(
@@ -404,6 +411,7 @@ def get_output_stats(self, input_stats) -> RepresentationStats:
                     0,
                     1,
                 ),
+                dtype=self.data_type,
             )
         if self.params and "_pushdown_aggregation" in self.params:
             output_shape = (512,)

src/main/python/systemds/scuro/representations/concatenation.py

@@ -34,11 +34,70 @@
 
 @register_fusion_operator()
 class Concatenation(Fusion):
-    def __init__(self, params=None):
-        """
-        Combines modalities using concatenation
-        """
+    def __init__(self, params=None, preserve_leading_axis=False):
         super().__init__("Concatenation")
+        if params is not None:
+            preserve_leading_axis = params.get(
+                "preserve_leading_axis", preserve_leading_axis
+            )
+
+        self.preserve_leading_axis = bool(preserve_leading_axis)
+        self.preserves_leading_axis = self.preserve_leading_axis
+
+    def get_current_parameters(self):
+        current_params = super().get_current_parameters()
+        current_params["preserve_leading_axis"] = self.preserve_leading_axis
+        return current_params
+
+    @staticmethod
+    def _as_dense(modality):
+        dtype = modality.metadata[0]["data_layout"]["type"]
+        data = modality.data
+        arr = (
+            np.asarray(data, dtype=dtype) if not isinstance(data, np.ndarray) else data
+        )
+        if arr.dtype == object:
+            instances = [np.asarray(instance, dtype=dtype) for instance in data]
+            rest = tuple(
+                max(i.shape[d] for i in instances) for d in range(instances[0].ndim)
+            )
+            arr = np.zeros((len(instances), *rest), dtype=dtype)
+            for i, instance in enumerate(instances):
+                arr[(i, *(slice(0, s) for s in instance.shape))] = instance
+        return arr
+
+    @staticmethod
+    def _to_window_feature_matrix(arr):
+        if arr.ndim == 1:
+            return arr[:, None, None]
+        if arr.ndim == 2:
+            return arr[:, :, None]
+        return arr.reshape(arr.shape[0], arr.shape[1], -1)
+
+    @staticmethod
+    def _flatten_feature_shape(shape):
+        if len(shape) == 0:
+            return (1, 1)
+        if len(shape) == 1:
+            return (shape[0], 1)
+        return (shape[0], int(np.prod(shape[1:])))
+
+    def _concat_on_leading_axis(self, modalities: List[Modality]):
+        arrays = [
+            self._to_window_feature_matrix(self._as_dense(modality))
+            for modality in modalities
+        ]
+
+        num_windows = max(arr.shape[1] for arr in arrays)
+        aligned = []
+        for arr in arrays:
+            if arr.shape[1] < num_windows:
+                pad_width = [(0, 0)] * arr.ndim
+                pad_width[1] = (0, num_windows - arr.shape[1])
+                arr = np.pad(arr, pad_width=pad_width, mode="constant")
+            aligned.append(arr)
+
+        return np.concatenate(aligned, axis=-1)
 
     def execute(self, modalities: List[Modality]):
         if len(modalities) == 1:
@@ -47,6 +106,9 @@ def execute(self, modalities: List[Modality]):
                 dtype=modalities[0].metadata[0]["data_layout"]["type"],
             )
 
+        if self.preserve_leading_axis:
+            return self._concat_on_leading_axis(modalities)
+
         max_emb_size = self.get_max_embedding_size(modalities)
         size = len(modalities[0].data)
 
@@ -71,35 +133,32 @@ def execute(self, modalities: List[Modality]):
         return np.array(data)
 
     def get_output_stats(self, input_stats_list) -> RepresentationStats:
-        if isinstance(input_stats_list, RepresentationStats):
-            return input_stats_list
-
-        stats_list = list(input_stats_list)
+        stats_list = self._fusion_input_stats(input_stats_list)
         if not stats_list:
             return RepresentationStats(0, (0,))
 
-        num_instances = stats_list[0].num_instances
-        total_dim = sum(s.output_shape[-1] for s in stats_list)
-        output_shape = (total_dim,)
-
-        return RepresentationStats(num_instances, output_shape)
+        num_instances = max(s.num_instances for s in stats_list)
+        shapes = [tuple(s.output_shape) for s in stats_list]
+        if self.preserve_leading_axis:
+            shapes = [self._flatten_feature_shape(shape) for shape in shapes]
+        rank = len(shapes[0])
 
-    def estimate_peak_memory_bytes(self, input_stats_list) -> dict:
-        elem_size = np.dtype(np.float32).itemsize
-
-        def stats_bytes(s: RepresentationStats) -> int:
-            numel = int(np.prod(s.output_shape)) if len(s.output_shape) > 0 else 1
-            return int(s.num_instances * numel * elem_size)
+        if rank >= 1 and all(len(shape) == rank for shape in shapes):
+            leading = tuple(max(shape[d] for shape in shapes) for d in range(rank - 1))
+            output_shape = (*leading, sum(shape[-1] for shape in shapes))
+        else:
+            output_shape = max(stats_list, key=self._stats_num_elements).output_shape
 
-        current_output = 0
-        peak = 0
-        for s in input_stats_list:
-            chunk = stats_bytes(s)
-            new_output = current_output + chunk
+        output_shape_is_known = all(s.output_shape_is_known for s in stats_list)
+        return RepresentationStats(num_instances, output_shape, output_shape_is_known)
 
-            step_peak = current_output + chunk + new_output + chunk
-            peak = max(peak, step_peak)
-            current_output = new_output
+    def estimate_peak_memory_bytes(self, input_stats) -> dict:
+        stats_list = self._as_stats_list(input_stats)
+        input_bytes = sum(self._stats_bytes(s) for s in stats_list)
+        output_bytes = self._stats_bytes(self.get_output_stats(input_stats))
 
-        cpu_peak = int(peak * 1.15 + 16 * 1024 * 1024)
+        raw_bytes = self._raw_input_bytes(input_stats)
+        cpu_peak = (
+            int((raw_bytes + 2 * input_bytes + output_bytes) * 1.1) + 8 * 1024 * 1024
+        )
         return {"cpu_peak_bytes": cpu_peak, "gpu_peak_bytes": 0}