credit.datasets.gen_2.multi_source
==================================

.. py:module:: credit.datasets.gen_2.multi_source

.. autoapi-nested-parse::

   multi_source.py
   ---------------
   MultiSourceDataset: primary entry point for all data loading.

   Wraps one or more registered source datasets and returns a dict nested by
   source name.  Only sources whose keys appear under ``config["source"]`` are
   instantiated; absent sources are silently skipped.

   Sample structure returned by __getitem__::

       {
           "input":    {<user_provided_name>: {"<user_provided_name>/prognostic/3d/T": tensor, ...}, ...},
           "target":   {<user_provided_name>: {"<user_provided_name>/prognostic/3d/T": tensor, ...}, ...},  # return_target only
           "metadata": {<user_provided_name>: {"input_datetime": int, "target_datetime": int}, ...},
       }

   Usage::

       from credit.datasets.gen_2.multi_source import MultiSourceDataset
       from credit.samplers import DistributedMultiStepBatchSampler
       from torch.utils.data import DataLoader

       dataset = MultiSourceDataset(config["data"], return_target=True)
       sampler = DistributedMultiStepBatchSampler(dataset, batch_size=4,
                     shuffle=True, num_replicas=1, rank=0)
       loader = DataLoader(dataset, batch_sampler=sampler, num_workers=4)

   Extending with a new source::

       # In _SOURCE_REGISTRY, add:
       "new_source": ("credit.datasets.new_source", "NewSourceDataset"),

       # The dataset class must accept (config, return_target) and expose
       # a ``datetimes`` attribute (pd.DatetimeIndex).



Attributes
----------

.. autoapisummary::

   credit.datasets.gen_2.multi_source.logger


Classes
-------

.. autoapisummary::

   credit.datasets.gen_2.multi_source.MultiSourceDataset


Functions
---------

.. autoapisummary::

   credit.datasets.gen_2.multi_source.make_single_source_subconfig
   credit.datasets.gen_2.multi_source.route_to_dataset_class


Module Contents
---------------

.. py:data:: logger

.. py:function:: make_single_source_subconfig(config: dict[str, Any], user_dataset_name: str) -> dict[str, Any]

   Return a modified config dict containing only the specified source.

   This is used internally to instantiate each sub-dataset with a config
   containing just its own source config block, to avoid confusion with
   multisource config fields (e.g. HRRR vs HRRR_NAT vs HRRR_SUBH).

   :param config: Original multisource config dict.
   :param user_dataset_name: Unique dataset name specified by the user in config["source"] (e.g. "Example_ERA5").

   :returns: New config dict containing only the specified source's config block.


.. py:function:: route_to_dataset_class(source_cfg: dict[str, Any]) -> type

   Return the appropriate Dataset class based on the "dataset_type" field in the source config.

   The module containing the class is imported lazily on first call so that
   optional heavy dependencies are not loaded unless this source type is used.

   Built-in source types (e.g. "local", "arco_era5") are matched against
   ``_SOURCE_REGISTRY``. If no built-in matches, falls back to
   ``credit.datasets._DATASET_REGISTRY`` so that datasets registered via
   ``custom_objects`` in the config can be used as a source's ``dataset_type``.
   Both registries are matched case-sensitively.

   :param source_cfg: Config dict for a single source (e.g. config["source"]["Example_ERA5"]).

   :returns: Dataset class corresponding to the "dataset_type" field.

   :raises ValueError: If the "dataset_type" field is missing or does not correspond to a registered dataset.


.. py:class:: MultiSourceDataset(config: dict[str, Any], return_target: bool = False, label: str | None = None)

   Bases: :py:obj:`credit.datasets.gen_2.base_dataset.AbstractBaseDataset`


   CREDIT Dataset that combines multiple source datasets.

   Instantiates one sub-dataset per source key found in ``config["source"]``,
   computes the intersection of their valid timestamps, and delegates each
   ``__getitem__`` call to all active sub-datasets.

   See module docstring for full output structure and usage examples.

   Note that we inherit from AbstractBaseDataset _rather_ than BaseDataset.

   :ivar datasets: Ordered mapping of lowercase source name to its Dataset
                   instance (e.g. ``{"era5": ERA5Dataset, "mrms": MRMSDataset}``).
   :ivar datetimes: DatetimeIndex of timestamps valid for *all* active sources
                    (intersection of each source's own ``datetimes``).
   :ivar static_metadata: Per-source static metadata aggregated from each
                          sub-dataset's ``static_metadata`` attribute.



   .. py:attribute:: datasets
      :type:  dict[str, credit.datasets.gen_2.base_dataset.BaseDataset]


   .. py:attribute:: dt
      :type:  pandas.Timedelta


   .. py:attribute:: datetimes
      :type:  pandas.Index


   .. py:attribute:: static_metadata
      :type:  dict[str, dict[str, Any]]


   .. py:method:: __len__() -> int


   .. py:method:: __getitem__(args: tuple[pandas.Timestamp, int]) -> dict[str, dict[str, Any]]

      Return a dict of per-source sample dicts.

      :param args: ``(t, i)`` where *t* is the current timestamp (nanoseconds
                   or pd.Timestamp) and *i* is the within-sequence step index
                   produced by the sampler.

      :returns: Dict keyed by data type, each value being a dict of source name
                to that source's data::

                    {"input": {"era5": {...}, ...}, "target": {...}, "metadata": {...}}