Allow merge in concat dataset

[elynnwu]

When training multiple label datasets, we have a use case for having merge inside a concat. Ideally we will just have every variables we need in each labeled dataset, but it is often that one dataset is missing some variables. For example, if c96-shield does not have CO2, it will be nice if we can just merge a separate dataset. The resulting config will look something like this:

  dataset:
    concat:
    - data_path: /climate-default/2024-06-20-era5-1deg-8layer-1940-2022-netcdfs
      labels:
      - era5
      subset:
        start_time: '1979-01-01T00:00:00'
        stop_time: '1986-03-31T18:00:00'
    - merge:
       - data_path: /climate-default/2024-07-24-vertically-resolved-c96-1deg-shield-amip-ensemble-dataset/netCDFs/ic_0001
          labels:
          - c96-shield
          subset:
            start_time: '1979-01-01'
            stop_time: '2020-12-31'
       - data_path: /climate-default/shield-co2-data
          labels:
          - c96-shield

We support concat inside merge but not the other way around because concat is supposed to concatenate across time. But in the context of multi-label datasets, each concat is actually a different data source, it'll be convenient if we can have merge within each dataset.

[mcgibbon]

As long as we're in a regular habit of using merge datasets, I think this is a good feature to have.

As a side note, we should make sure we require merged datasets to all have the same labels (I think we already do).

[jpdunc23]

As you mentioned, we can combine merges and concat with concat inside of merge, and if we're careful about the merging then we can use this to combine multiple data sources. Here's an example combining multiple data sources to train on both the CM4 1pct and piControl datasets:

`train_loader` yaml example:

train_loader:
  dataset:
    merge:
    - concat:
        - data_path: /climate-default
          file_pattern: 2025-10-21-cm4-1pctCO2-140yr-no-smoothing-coupled-ocean.zarr
          engine: zarr
          subset:
            start_time: '0256-01-03T12:00:00'
        - data_path: /climate-default
          file_pattern: 2025-10-20-cm4-piControl-200yr-no-smoothing-coupled-ocean.zarr
          engine: zarr
          subset:
            start_time: '0151-01-03T12:00:00'
            stop_time: '0190-12-29T12:00:00'
      strict: false
    - concat:
        - data_path: /climate-default
          file_pattern: 2025-10-16-cm4-1pctCO2-140yr-ocean-no-smoothing.zarr
          engine: zarr
          subset:
            start_time: '0256-01-03T12:00:00'
        - data_path: /climate-default
          file_pattern: 2025-10-15-cm4-piControl-200yr-ocean-no-smoothing.zarr
          engine: zarr
          subset:
            start_time: '0151-01-03T12:00:00'
            stop_time: '0190-12-29T12:00:00'
      strict: false

I had to use the concat config attribute strict: false to get this to work. If you had a combination of datasets requiring merge and others that don't need it, you could probably hack it by using a "fake" merge, i.e. by just repeating the dataset.

train_loader yaml example with fake merge:

train_loader:
  dataset:
    merge:
    - concat:
        - data_path: /climate-default
          file_pattern: real_merge_override.zarr
          engine: zarr
        - data_path: /climate-default
          file_pattern: fake_merge.zarr
          engine: zarr
      strict: false
    - concat:
        - data_path: /climate-default
          file_pattern: real_merge_base.zarr
          engine: zarr
        - data_path: /climate-default
          file_pattern: fake_merge.zarr
          engine: zarr
      strict: false

This is likely to be pretty inefficient since it opens the same dataset twice, so I agree this is probably worth implementing.