Introduce DataSource

ml4ht/data/data_source.py

@@ -0,0 +1,200 @@
+from typing import (
+    Any,
+    Dict,
+    List,
+    Tuple,
+    Callable,
+    Optional,
+    Sequence,
+    Iterator,
+    Mapping,
+)
+
+from torch.utils.data import IterableDataset, Dataset, get_worker_info
+import numpy as np
+from ml4ht.data.defines import EXCEPTIONS
+
+
+DataIndex = Mapping[str, Any]
+Data = Dict[str, np.ndarray]
+# Given `DataIndex`, return (input data, output data)
+DataSource = Callable[[DataIndex], Tuple[Data, Data]]
+TrainingExample = Tuple[
+    Dict[str, np.ndarray],
+    Dict[str, np.ndarray],
+]
+
+
+def range_epoch_idx_generator(
+    true_epoch_len: int,
+    shuffle: bool,
+    training_epochs_per_true_epoch: int = 1,
+) -> Iterator[Sequence[int]]:
+    """
+    Yields optionally shuffled lists of indices from a range
+    :param true_epoch_len: how many indices to yield from
+    :param shuffle: Whether to shuffle the sample ids each epoch
+    :param training_epochs_per_true_epoch: How many training epochs to break a true epoch into.
+        Defaults to 1 to yield true epochs.
+    :yield: Epoch worth of sample ids
+    """
+    assert training_epochs_per_true_epoch <= true_epoch_len
+    sample_ids = np.arange(true_epoch_len)
+    while True:
+        if shuffle:
+            np.random.shuffle(sample_ids)
+        epochs = np.array_split(sample_ids, training_epochs_per_true_epoch)
+        yield from epochs
+
+
+def data_index_epoch_generator(
+    data_indices: List[DataIndex],
+    shuffle: bool,
+    training_epochs_per_true_epoch: int = 1,
+) -> Iterator[Sequence[DataIndex]]:
+    """
+    Yields optionally shuffled epochs of `DataIndex`s
+    :param data_indices: sample ids (e.g. MRNs) for each epoch
+    :param shuffle: Whether to shuffle the sample ids each epoch
+    :param training_epochs_per_true_epoch: How many training epochs to break a true epoch into.
+        Defaults to 1 to yield true epochs.
+    :yield: Epoch worth of DataIndices
+    """
+    epoch_idx_generator = range_epoch_idx_generator(
+        true_epoch_len=len(data_indices),
+        shuffle=shuffle,
+        training_epochs_per_true_epoch=training_epochs_per_true_epoch,
+    )
+    for indices in epoch_idx_generator:
+        yield [data_indices[idx] for idx in indices]
+
+
+class TrainingDataset(Dataset):
+    """
+    Indexable of training data from `DataSources` for use in pytorch `DataLoader`.
+    To be able to skip errors during data loading, use `TrainingIterableDataset` below.
+    """
+
+    def __init__(
+        self,
+        data_sources: List[DataSource],
+        data_indices: Sequence[DataIndex],
+    ):
+        """
+        :param data_sources: Sources of data for inputs and outputs to model
+        :param data_indices: Indices for each epoch of data
+        """
+        super(TrainingIterableDataset).__init__()
+        self.data_sources = data_sources
+        self.data_indices = data_indices
+
+    def __len__(self):
+        return len(self.data_indices)
+
+    def training_example(self, data_idx: DataIndex) -> Tuple[Data, Data]:
+        combined_inputs, combined_outputs = {}, {}
+        for data_source in self.data_sources:
+            data_in, data_out = data_source(data_idx)
+            for name, modality in data_in.items():
+                if name in combined_inputs:
+                    raise ValueError(
+                        f"Multiple DataSources share the same input modality called {name}",
+                    )
+                combined_inputs[name] = modality
+            for name, modality in data_out.items():
+                if name in combined_outputs:
+                    raise ValueError(
+                        f"Multiple DataSources share the same output modality called {name}",
+                    )
+                combined_outputs[name] = modality
+        return combined_inputs, combined_outputs
+
+    def __getitem__(self, item: int) -> Tuple[Data, Data]:
+        data_idx = self.data_indices[item]
+        return self.training_example(data_idx)
+
+
+class TrainingIterableDataset(IterableDataset):
+    """
+    Iterable over epochs of training data from `DataSources` for use in pytorch `DataLoader`.
+    More flexible version of `TrainingDataset` that can skip errors,
+    but requires defining an Iterator over epochs of `DataIndex`s.
+    """
+
+    # keys in error list
+    error_key = "error"
+    source_key = "source"
+    idx_key = "data_idx"
+
+    def __init__(
+        self,
+        data_sources: List[DataSource],
+        epoch_indices_iterator: Iterator[Sequence[DataIndex]],
+        raise_errors: bool = True,
+        verbose: bool = True,
+    ):
+        """
+        :param data_sources: Sources of data for inputs and outputs to model
+        :param epoch_indices_iterator: Iterator over epochs of `DataIndex`s
+            For example, see `sample_id_epoch_indices` above.
+        :param raise_errors: Whether to raise or skip errors during calling of
+            `DataFetcher`s
+        """
+        super(TrainingIterableDataset).__init__()
+        self.data_sources = data_sources
+        self.epoch_indices_iterator = epoch_indices_iterator
+        self.raise_errors = raise_errors
+        self.verbose = verbose
+        self.errors = []
+
+    def __iter__(self):
+        """
+        One epoch of training data.
+        """
+        epoch_indices = next(self.epoch_indices_iterator)
+        worker_info = get_worker_info()
+        if worker_info is not None:  # multi-process case
+            split_indices = np.array_split(epoch_indices, worker_info.num_workers)
+            epoch_indices = split_indices[worker_info.id]
+        return filter(
+            lambda x: x is not None,  # skips the errors
+            map(self.training_example, epoch_indices),
+        )
+
+    def format_error(self, error: Dict[str, Any]) -> str:
+        return f"On index `{error[self.idx_key]}` got error `{error[self.error_key]}` from source `{error[self.source_key]}`"
+
+    def training_example(
+        self,
+        data_idx: DataIndex,
+    ) -> Optional[Tuple[Data, Data]]:
+        combined_inputs, combined_outputs = {}, {}
+        for data_source in self.data_sources:
+            try:
+                data_in, data_out = data_source(data_idx)
+            except EXCEPTIONS as e:
+                if self.raise_errors:
+                    raise e
+                self.errors.append(
+                    {
+                        self.error_key: repr(e),
+                        self.source_key: data_source.__name__,
+                        self.idx_key: data_idx,
+                    },
+                )
+                if self.verbose:
+                    print(self.format_error(self.errors[-1]))
+                return
+            for name, modality in data_in.items():
+                if name in combined_inputs:
+                    raise ValueError(
+                        f"Multiple DataSources share the same input modality called {name}",
+                    )
+                combined_inputs[name] = modality
+            for name, modality in data_out.items():
+                if name in combined_outputs:
+                    raise ValueError(
+                        f"Multiple DataSources share the same output modality called {name}",
+                    )
+                combined_outputs[name] = modality
+        return combined_inputs, combined_outputs

ml4ht/data/explore.py

@@ -1,12 +1,14 @@
 from functools import partial
 from multiprocessing import Pool, cpu_count
-from typing import Callable, List
+from typing import Callable, List, TypeVar, Dict
 
 import pandas as pd
+import numpy as np
 
 from ml4ht.data.data_description import DataDescription
 from ml4ht.data.defines import EXCEPTIONS, SampleID, SampleGetter
 from ml4ht.data.sample_getter import DataDescriptionSampleGetter
+from ml4ht.data.data_source import DataSource, DataIndex
 
 ERROR_COL = "error"
 NO_LOADING_OPTIONS_ERROR = ValueError("No loading options")
@@ -15,9 +17,12 @@
 LOADING_OPTION_COL = "state"
 
 
+T = TypeVar("T")
+
+
 def build_df(
-    summarizer: Callable[[SampleID], pd.DataFrame],
-    sample_ids: List[SampleID],
+    summarizer: Callable[[T], pd.DataFrame],
+    sample_ids: List[T],
     multiprocess_workers: int = 0,
 ) -> pd.DataFrame:
     """
@@ -122,13 +127,25 @@ def explore_data_descriptions(
     )
 
 
+def _summarize_tensor(tensor: np.ndarray) -> Dict[str, float]:
+    out = {}
+    out["mean"] = tensor.mean()
+    out["median"] = np.median(tensor)
+    out["std"] = tensor.std()
+    out["min"] = tensor.min()
+    out["max"] = tensor.max()
+    out["argmax"] = tensor.argmax()
+    out["shape"] = tensor.shape
+    return out
+
+
 # sample getter explore
 def _pipeline_sample_getter_summarize_sample_id(
     sample_id: SampleID,
     sample_getter: SampleGetter,
 ) -> pd.DataFrame:
     """
-    Get the summary, dates, states or errors from each TensorMap for a sample id
+    Get the summary, dates, states or errors from each input and output for a sample id
     """
     out = {SAMPLE_ID_COL: sample_id}
     try:
@@ -137,12 +154,9 @@ def _pipeline_sample_getter_summarize_sample_id(
             **data[0],
             **data[1],
         }.items():
-            out[f"{name}_mean"] = [tensor.mean()]
-            out[f"{name}_std"] = [tensor.std()]
-            out[f"{name}_min"] = [tensor.min()]
-            out[f"{name}_max"] = [tensor.max()]
-            out[f"{name}_argmax"] = [tensor.argmax()]
-            out[f"{name}_shape"] = [tensor.shape]
+            tensor_summary = _summarize_tensor(tensor)
+            for field, val in tensor_summary.items():
+                out[f"{name}_{field}"] = [val]
     except EXCEPTIONS as e:
         out[ERROR_COL] = [_format_exception(e)]
     out = pd.DataFrame(out)
@@ -155,7 +169,7 @@ def explore_sample_getter(
     multiprocess_workers: int = 0,
 ) -> pd.DataFrame:
     """
-    Get the datetime selected,
+    Summarize the values and errors of a sample getter
     """
     summarize = partial(
         _pipeline_sample_getter_summarize_sample_id,
@@ -166,3 +180,69 @@ def explore_sample_getter(
         sample_ids,
         multiprocess_workers,
     )
+
+
+def _data_source_auto_summarize_sample_id(
+    data_idx: DataIndex,
+    data_source: DataSource,
+) -> pd.DataFrame:
+    out = {f"idx: {k}": [v] for k, v in data_idx.items()}
+    try:
+        data_in, data_out = data_source(data_idx)
+        for name, tensor in data_in.items():
+            out.update(
+                {
+                    f"input: {name}_{field}": [val]
+                    for field, val in _summarize_tensor(tensor).items()
+                },
+            )
+        for name, tensor in data_out.items():
+            out.update(
+                {
+                    f"output: {name}_{field}": [val]
+                    for field, val in _summarize_tensor(tensor).items()
+                },
+            )
+    except EXCEPTIONS as e:
+        out[ERROR_COL] = [_format_exception(e)]
+    return pd.DataFrame(out)
+
+
+def _data_source_summarize_sample_id(
+    data_idx: DataIndex,
+    data_source: DataSource,
+) -> pd.DataFrame:
+    out = {f"idx: {k}": [v] for k, v in data_idx.items()}
+    try:
+        data_in, data_out = data_source(data_idx)
+        for name, tensor in data_in.items():
+            out[f"input: {name}"] = [tensor]
+        for name, tensor in data_out.items():
+            out[f"output: {name}"] = [tensor]
+    except EXCEPTIONS as e:
+        out[ERROR_COL] = [_format_exception(e)]
+    return pd.DataFrame(out)
+
+
+def explore_data_source(
+    data_source: DataSource,
+    data_indices: List[DataIndex],
+    auto_summarize=True,
+    multiprocess_workers: int = 0,
+) -> pd.DataFrame:
+    """
+    Get DataSource values over all the data indices.
+    `auto_summarize` converts those values into summary statistics
+    in the output `DataFrame`.
+    """
+    summarize = partial(
+        _data_source_auto_summarize_sample_id
+        if auto_summarize
+        else _data_source_summarize_sample_id(),
+        data_source=data_source,
+    )
+    return build_df(
+        summarize,
+        data_indices,
+        multiprocess_workers,
+    )

ml4ht/data/util/data_source_util.py

@@ -0,0 +1,19 @@
+from typing import Sequence
+
+import pandas as pd
+
+from ml4ht.data.data_source import DataIndex, range_epoch_idx_generator
+
+
+def df_epoch_index_generator(
+    df: pd.DataFrame,
+    shuffle: bool,
+    training_epochs_per_true_epoch: int,
+) -> Sequence[DataIndex]:
+    epoch_idx_generator = range_epoch_idx_generator(
+        true_epoch_len=len(df),
+        shuffle=shuffle,
+        training_epochs_per_true_epoch=training_epochs_per_true_epoch,
+    )
+    for epoch_indices in epoch_idx_generator:
+        yield df.iloc[epoch_indices].reindex()