diff --git a/config/streams/era5_1deg_forecasting_conditioning/day_of_year_conditioning.yml b/config/streams/era5_1deg_forecasting_conditioning/day_of_year_conditioning.yml new file mode 100644 index 0000000000..452e63dae1 --- /dev/null +++ b/config/streams/era5_1deg_forecasting_conditioning/day_of_year_conditioning.yml @@ -0,0 +1,22 @@ +# (C) Copyright 2025 WeatherGenerator contributors. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. + +day_of_year_conditioning: + type: time_conditioning + conditioning: True + filenames: null + timestep_conditioning: scalar + conditioning_type: time_based + value_type: day + embed: + net: linear + dim_embed: 256 + conditioning_FE: + type: add + dim_embed: 256 diff --git a/config/streams/era5_1deg_forecasting_conditioning/era5.yml b/config/streams/era5_1deg_forecasting_conditioning/era5.yml new file mode 100644 index 0000000000..6d6078cae4 --- /dev/null +++ b/config/streams/era5_1deg_forecasting_conditioning/era5.yml @@ -0,0 +1,110 @@ +# (C) Copyright 2024 WeatherGenerator contributors. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. + +ERA5 : + type : anemoi + filenames : ['aifs-ea-an-oper-0001-mars-o96-1979-2024-1h-v3-with-era51.zarr'] + stream_id : 0 + source_exclude : ['z', 'w_10', 'w_50', 'w_100', 'w_150', 'w_200', 'w_250', 'w_300', 'w_400', 'w_500', 'w_600', 'w_700', 'w_850', 'w_925', 'w_1000', 'skt', 'tcw', 'cp', 'tp', 'q_50', 'q_100'] + target_exclude : ['z', 'w_10', 'w_50', 'w_100', 'w_150', 'w_200', 'w_250', 'w_300', 'w_400', 'w_500', 'w_600', 'w_700', 'w_850', 'w_925', 'w_1000', 'slor', 'sdor', 'tcw', 'cp', 'tp', 'q_50', 'q_100'] + geoinfo_channels : ['z', 'lsm', 'slor', 'sdor', 'insolation', 'cos_local_time', 'sin_local_time', 'cos_julian_day', 'sin_julian_day'] + loss_weight : 1. + location_weight : cosine_latitude + token_size : 8 + tokenize_spacetime : True + max_num_targets: 20000 + frequency : 06:00:00 + embed : + net : transformer + num_tokens : 1 + num_heads : 8 + dim_embed : 512 + num_blocks : 2 + embed_target_coords : + net : linear + dim_embed : 512 + target_readout : + num_layers : 2 + num_heads : 4 + # sampling_rate : 0.2 + pred_head : + ens_size : 1 + num_layers : 1 + channel_weights : + q_10: 0.2 + q_50: 0.2 + q_100: 0.23 + q_150: 0.26 + q_200: 0.29 + q_250: 0.33 + q_300: 0.36 + q_400: 0.42 + q_500: 0.48 + q_600: 0.55 + q_700: 0.61 + q_850: 0.71 + q_925: 0.75 + q_1000: 0.8 + t_10: 0.2 + t_50: 0.2 + t_100: 0.23 + t_150: 0.26 + t_200: 0.29 + t_250: 0.33 + t_300: 0.36 + t_400: 0.42 + t_500: 0.48 + t_600: 0.55 + t_700: 0.61 + t_850: 0.71 + t_925: 0.75 + t_1000: 0.8 + u_10: 0.2 + u_50: 0.2 + u_100: 0.23 + u_150: 0.26 + u_200: 0.29 + u_250: 0.33 + u_300: 0.36 + u_400: 0.42 + u_500: 0.48 + u_600: 0.55 + u_700: 0.61 + u_850: 0.71 + u_925: 0.75 + u_1000: 0.8 + v_10: 0.2 + v_50: 0.2 + v_100: 0.23 + v_150: 0.26 + v_200: 0.29 + v_250: 0.33 + v_300: 0.36 + v_400: 0.42 + v_500: 0.48 + v_600: 0.55 + v_700: 0.61 + v_850: 0.71 + v_925: 0.75 + v_1000: 0.8 + z_10: 0.2 + z_50: 0.2 + z_100: 0.23 + z_150: 0.26 + z_200: 0.29 + z_250: 0.33 + z_300: 0.36 + z_400: 0.42 + z_500: 0.48 + z_600: 0.55 + z_700: 0.61 + z_850: 0.71 + z_925: 0.75 + z_1000: 0.8 + \ No newline at end of file diff --git a/config/streams/era5_1deg_forecasting_conditioning/era5_sst_conditioning.yml b/config/streams/era5_1deg_forecasting_conditioning/era5_sst_conditioning.yml new file mode 100644 index 0000000000..31ba987128 --- /dev/null +++ b/config/streams/era5_1deg_forecasting_conditioning/era5_sst_conditioning.yml @@ -0,0 +1,49 @@ +# (C) Copyright 2025 WeatherGenerator contributors. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. + +# Example field conditioning stream: SST forcing for the atmosphere. +# The conditioning encoder reads SST data at each forecast step and encodes it +# into a latent representation that is injected into the ForecastingEngine. +ERA5_SST: + type: anemoi + conditioning: True + filenames: ['ERA5-1deg-6h-mean-forcing-1978-2024.zarr'] + timestep_conditioning: field + stream_id: 99 + source: ['sea_surface_temperature'] + token_size: 8 + tokenize_spacetime: False + geoinfo_channels: [] + + # Conditioning encoder configuration. + # Option A — load a pre-trained encoder from a model checkpoint: + # encoder: + # model_id: XXX + # model_epoch: XXX + # + # Option B — define an inline encoder (same parameters as the main encoder): + encoder: + ae_local_dim_embed: 512 + ae_local_num_blocks: 2 + ae_local_num_heads: 8 + ae_local_num_queries: 1 + ae_global_dim_embed: 512 + ae_global_num_blocks: 2 + ae_global_num_heads: 8 + + embed: + net: transformer + num_tokens: 1 + num_heads: 8 + dim_embed: 512 + num_blocks: 2 + + conditioning_FE: + type: cross_attention + dim_embed: 256 diff --git a/config/streams/era5_1deg_forecasting_conditioning/time_of_day_conditioning.yml b/config/streams/era5_1deg_forecasting_conditioning/time_of_day_conditioning.yml new file mode 100644 index 0000000000..8b355c67ba --- /dev/null +++ b/config/streams/era5_1deg_forecasting_conditioning/time_of_day_conditioning.yml @@ -0,0 +1,22 @@ +# (C) Copyright 2025 WeatherGenerator contributors. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. + +time_of_day_conditioning: + type: time_conditioning + conditioning: True + filenames: null + timestep_conditioning: scalar + conditioning_type: time_based + value_type: hour + embed: + net: linear + dim_embed: 256 + conditioning_FE: + type: add + dim_embed: 256 diff --git a/packages/readers_extra/src/weathergen/readers_extra/registry.py b/packages/readers_extra/src/weathergen/readers_extra/registry.py index b13ee46550..4d8f00ae8b 100644 --- a/packages/readers_extra/src/weathergen/readers_extra/registry.py +++ b/packages/readers_extra/src/weathergen/readers_extra/registry.py @@ -4,6 +4,10 @@ def get_extra_reader(stream_type: str) -> object | None: # There is no sanity check on them, so they may fail at runtime during imports match stream_type: + case "time_conditioning": + from weathergen.datasets.data_reader_time_conditioning import DataReaderTimeConditioning + + return DataReaderTimeConditioning case "iconart": from weathergen.readers_extra.data_reader_iconart import DataReaderIconArt diff --git a/src/weathergen/datasets/batch.py b/src/weathergen/datasets/batch.py index ea6a0b26ab..c28dc6485f 100644 --- a/src/weathergen/datasets/batch.py +++ b/src/weathergen/datasets/batch.py @@ -25,6 +25,10 @@ class SampleMetaData: global_params: dict | None = None + # Per-step scalar conditioning values, shape (num_output_steps, scalar_dim). + # Populated by MultiStreamDataSampler for streams with timestep_conditioning: scalar. + conditioning: np.typing.NDArray | None = None + class Sample: # keys: stream name, values: SampleMetaData @@ -67,6 +71,10 @@ def to_device(self, device) -> None: if self.meta_info[key].mask is not None else None ) + if self.meta_info[key].conditioning is not None: + self.meta_info[key].conditioning = self.meta_info[key].conditioning.to( + device, non_blocking=True + ) for key, val in self.streams_data.items(): if val is not None: @@ -301,6 +309,7 @@ class ModelBatch: def __init__( self, stream_names: list[str], + conditioning_stream_names: list[str], num_source_samples: int, num_target_samples: int, output_offset, @@ -320,6 +329,13 @@ def __init__( stream_names, num_target_samples, output_steps, self.output_idxs ) + self.conditioning_samples = BatchSamples( + stream_names=conditioning_stream_names, + num_samples=1, + output_steps=output_steps, + output_idxs=self.output_idxs, + ) + self.source2target_matching_idxs = np.full(num_source_samples, -1, dtype=np.int32) self.target2source_matching_idxs = [[] for _ in range(num_target_samples)] @@ -332,6 +348,9 @@ def pin_memory(self): # pin target samples self.target_samples.pin_memory() + # pin conditioning samples + self.conditioning_samples.pin_memory() + return self def to_device(self, device): # -> ModelBatch @@ -341,6 +360,7 @@ def to_device(self, device): # -> ModelBatch self.source_samples.to_device(device) self.target_samples.to_device(device) + self.conditioning_samples.to_device(device) self.device = device @@ -391,6 +411,40 @@ def add_target_stream( ) self.target2source_matching_idxs[target_sample_idx] = source_sample_idx + def add_scalar_conditioning_stream(self, stream_name, conditioning_values): + """ + Add scalar conditioning values for all samples in the batch for a specific stream. + """ + for sample in self.source_samples.samples: + if stream_name not in sample.meta_info: + sample.add_meta_info(stream_name, SampleMetaData(params={})) + sample.meta_info[stream_name].conditioning = conditioning_values + + def add_field_conditioning_stream(self, stream_name, stream_data: StreamData): + """ + Add field conditioning values for all samples in the batch for a specific stream. + The StreamData contains one source step per forecast step. + """ + for sample in self.conditioning_samples.samples: + sample.streams_data[stream_name] = stream_data + + def get_scalar_conditioning_values( + self, stream_name: str, step: int + ) -> np.typing.NDArray | None: + """ + Get scalar conditioning values for all samples at a specific forecast step. + + Returns np.ndarray of shape (num_samples, scalar_dim), or None if not available + for any sample. + """ + values = [] + for sample in self.source_samples.samples: + meta = sample.meta_info.get(stream_name) + if meta is None or meta.conditioning is None or step >= len(meta.conditioning): + return None + values.append(meta.conditioning[step]) + return torch.stack(values, dim=0) if values else None + def is_empty(self): """ Check if batch is empty diff --git a/src/weathergen/datasets/data_reader_time_conditioning.py b/src/weathergen/datasets/data_reader_time_conditioning.py new file mode 100644 index 0000000000..0e52b8421f --- /dev/null +++ b/src/weathergen/datasets/data_reader_time_conditioning.py @@ -0,0 +1,115 @@ +# (C) Copyright 2025 WeatherGenerator contributors. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. + +""" +Synthetic data reader for scalar conditioning. + +Generates synthetic scalar values like time_of_day, day_of_year, and noise_level +that can be used for model conditioning. +""" + +import logging + +import numpy as np + +from weathergen.datasets.data_reader_base import ( + DataReaderBase, + ReaderData, + TimeWindowHandler, + check_reader_data, +) +from weathergen.train.utils import Stage + +_logger = logging.getLogger(__name__) + + +class DataReaderTimeConditioning(DataReaderBase): + """ + DataReader that generates synthetic scalar conditioning values. + + Supports three types: + - time_based: extracts hour/day from timestamp + - constant: fixed value + - random: random uniform value (per timestep) + """ + + def __init__( + self, + tw_handler: TimeWindowHandler, + filename, + stream_info: dict, + stage: Stage, + ) -> None: + super().__init__(tw_handler, stream_info) + + self.source_idx = [] + self.target_idx = [] + self.geoinfo_idx = [] + + self.conditioning_type = stream_info.get("conditioning_type", "time_based") + value_type = stream_info.get("value_type") + if isinstance(value_type, list): + self.value_types = value_type + elif value_type is not None: + self.value_types = [value_type] + else: + raise ValueError("value_type in time_conditioning must be specified in stream_info") + + self.source_channels = [f"time_conditioning_{vt}" for vt in self.value_types] + self.target_channels = [] + self.geoinfo_channels = [] + self.target_channel_weights = [] + + self.mean = np.zeros(len(self.source_channels), dtype=np.float32) + self.stdev = np.ones(len(self.source_channels), dtype=np.float32) + self.mean_geoinfo = np.zeros(0, dtype=np.float32) + self.stdev_geoinfo = np.ones(0, dtype=np.float32) + + def length(self) -> int: + return self.time_window_handler.get_index_range().end + + def _get(self, idx, channels_idx) -> ReaderData: + dt_range = self.time_window_handler.window(idx) + dt = dt_range.start + + values = [] + for vt in self.value_types: + # The values for time conditioning below denote fractional values of the + # day or year, normalized to [0, 1]. + if self.conditioning_type == "time_based": + if vt == "hour": + hours = dt.astype("datetime64[h]").astype(int) % 24 + minutes = dt.astype("datetime64[m]").astype(int) % 60 + total_hours = hours + minutes / 60.0 + value = total_hours / 24.0 + elif vt == "day": + days = dt.astype("datetime64[D]").astype(int) % 365 + value = days / 365.0 + else: + raise ValueError(f"Unknown value_type: {vt}") + values.append(value) + else: + raise ValueError(f"Unknown conditioning_type: {self.conditioning_type}") + + coords = np.zeros((1, 2), dtype=np.float32) + geoinfos = np.zeros((1, 0), dtype=np.float32) + datetimes = np.array([dt], dtype=np.datetime64) + + rdata = ReaderData( + coords=coords, + geoinfos=geoinfos, + data=np.array(values, dtype=np.float32).reshape(1, -1), + datetimes=datetimes, + ) + check_reader_data(rdata, dt_range) + + return rdata + + def get_geoinfo_size(self) -> int: + return 0 diff --git a/src/weathergen/datasets/multi_stream_data_sampler.py b/src/weathergen/datasets/multi_stream_data_sampler.py index bd4a8ad87e..6a0c55cfe3 100644 --- a/src/weathergen/datasets/multi_stream_data_sampler.py +++ b/src/weathergen/datasets/multi_stream_data_sampler.py @@ -28,7 +28,7 @@ from weathergen.datasets.data_reader_obs import DataReaderObs from weathergen.datasets.masking import Masker from weathergen.datasets.stream_data import StreamData, spoof -from weathergen.datasets.tokenizer_masking import TokenizerMasking +from weathergen.datasets.tokenizer_masking import TokenizerMasking, readerdata_to_torch from weathergen.datasets.utils import ( get_tokens_lens, ) @@ -144,7 +144,13 @@ def __init__(self, cf: Config, mode_cfg: dict, stage: Stage): # check samples per mini epoch self.samples_per_mini_epoch = mode_cfg.samples_per_mini_epoch self.check_samples(self._get_fsm()) - self.streams_datasets = self._init_stream_datasets(cf) + ( + self.streams_datasets, + self.scalar_conditioning_datasets, + self.field_conditioning_datasets, + ) = self._init_stream_datasets(cf) + self.scalar_conditioning_stream_names = list(self.scalar_conditioning_datasets.keys()) + self.field_conditioning_stream_names = list(self.field_conditioning_datasets.keys()) # RNG seed setup rs = cf.data_loading.rng_seed @@ -217,11 +223,14 @@ def _calc_baseperms(self, fsm: int) -> np.typing.NDArray: return np.arange(self.max_input_steps, perms_len) - def _init_stream_datasets(self, cf) -> dict[StreamName, _Stream]: + def _init_stream_datasets( + self, cf + ) -> tuple[dict[StreamName, _Stream], dict[StreamName, _Stream], dict[StreamName, _Stream]]: """Load dataset readers for all streams from config.""" streams_datasets: dict[StreamName, _Stream] = {} for stream_name, stream_info in cf.streams.items(): stream_info["data_paths"] = cf.get("data_paths", []) + ds_type = stream_info["type"] # list of sources for current stream streams_datasets[stream_name] = _Stream(stream_info, []) kwargs = { @@ -230,7 +239,7 @@ def _init_stream_datasets(self, cf) -> dict[StreamName, _Stream]: "stage": self._stage, } dataset: type[AnyDataReader] | None = None - match stream_info["type"]: + match ds_type: case "obs": dataset = DataReaderObs case "anemoi": @@ -242,8 +251,29 @@ def _init_stream_datasets(self, cf) -> dict[StreamName, _Stream]: f"for stream name '{stream_name}'." raise ValueError(msg) - for fname in stream_info.get("filenames", [pathlib.Path()]): + filenames_cfg = stream_info.get("filenames", [pathlib.Path()]) + + if filenames_cfg is None: + filenames_cfg = [pathlib.Path()] + else: + pass + + for fname in filenames_cfg: fname = pathlib.Path(fname) + # skip if explicitly pointing to current directory + if fname == pathlib.Path(): + if stream_info.get("conditioning", False): + if is_root(): + logger.info( + f"Opening conditioning dataset with type: {ds_type}" + + f" from stream config {stream_name}.", + ) + ds = dataset(filename=None, **kwargs) + streams_datasets[stream_name].readers += [ds] + continue + else: + msg = f"Did not find input data for {ds_type} stream '{stream_name}'." + raise FileNotFoundError(msg) # dont check if file exists since zarr stores might be directories if fname.exists(): # check if fname is a valid path to allow for simple overwriting @@ -254,12 +284,11 @@ def _init_stream_datasets(self, cf) -> dict[StreamName, _Stream]: filename = next((f for f in filenames if f.exists()), None) if filename is None: msg = ( - f"Did not find input data for {stream_info['type']} " + f"Did not find input data for {ds_type} " f"stream '{stream_name}': {filenames}." ) raise FileNotFoundError(msg) - ds_type = stream_info["type"] if is_root(): logger.info( f"Opening dataset with type: {ds_type}" @@ -277,7 +306,20 @@ def _init_stream_datasets(self, cf) -> dict[StreamName, _Stream]: else [1.0 for _ in ds.target_channels] ) - return streams_datasets + # Separate streams by timestep_conditioning type + regular: dict[StreamName, _Stream] = {} + scalar_conditioning: dict[StreamName, _Stream] = {} + field_conditioning: dict[StreamName, _Stream] = {} + for name, stream in streams_datasets.items(): + tc = stream.info.get("timestep_conditioning") + if tc == "scalar": + scalar_conditioning[name] = stream + elif tc == "field": + field_conditioning[name] = stream + else: + regular[name] = stream + + return regular, scalar_conditioning, field_conditioning def reset(self) -> tuple[Sequence[int], Sequence[int]]: """ @@ -640,12 +682,17 @@ def _preprocess_model_batch( Perform necessary pre-processing of model batch """ stream_names = list(self.streams_datasets.keys()) + conditioning_stream_names = list(self.field_conditioning_datasets.keys()) batch.source_samples.tokens_lens = get_tokens_lens( stream_names, batch.source_samples, source_input_steps ) batch.target_samples.tokens_lens = get_tokens_lens( stream_names, batch.target_samples, target_input_steps ) + if conditioning_stream_names: + batch.conditioning_samples.tokens_lens = get_tokens_lens( + conditioning_stream_names, batch.conditioning_samples, input_steps=1 + ) return batch @@ -676,6 +723,7 @@ def _get_batch(self, idx: int, num_forecast_steps: int): num_output_steps = self._get_output_length(num_forecast_steps) batch = ModelBatch( list(self.streams_datasets.keys()), + list(self.field_conditioning_datasets.keys()), num_source_samples, num_target_samples, self.output_offset, @@ -754,10 +802,77 @@ def _get_batch(self, idx: int, num_forecast_steps: int): source_in_steps = input_steps.max().item() target_in_steps = np.array([tc.get("num_steps_input", 1) for _, tc in target_cfgs.items()]) target_in_steps = 1 if len(target_in_steps) == 0 else target_in_steps.max().item() + + if self.scalar_conditioning_stream_names: + batch = self._build_scalar_conditioning_data(batch, idx, num_forecast_steps) + + if self.field_conditioning_stream_names: + batch = self._build_field_conditioning_data(batch, idx, num_forecast_steps) + batch = self._preprocess_model_batch(batch, source_in_steps, target_in_steps) return batch + def _build_scalar_conditioning_data( + self, batch: ModelBatch, idx: int, num_forecast_steps: int + ) -> ModelBatch: + """Collect per-step scalar conditioning values and store in sample meta_info.""" + num_output_steps = self._get_output_length(num_forecast_steps) + for stream_name, stream_ds in self.scalar_conditioning_datasets.items(): + step_values = [] + for timestep_idx in range(self.output_offset, num_output_steps): + step_dt = idx + (self.time_step * timestep_idx) // self.step_timedelta + rdata = stream_ds.readers[0].get_source(step_dt) + step_values.append(readerdata_to_torch(rdata).data.flatten()) + conditioning_values = ( + torch.stack(step_values, dim=0) + if step_values + else torch.zeros((0, 1), dtype=torch.float32) + ) + batch.add_scalar_conditioning_stream(stream_name, conditioning_values) + return batch + + def _build_field_conditioning_data( + self, batch: ModelBatch, idx: int, num_forecast_steps: int + ) -> ModelBatch: + """Collect per-step field conditioning data with sliding window and + store in conditioning_samples.streams_data.""" + num_output_steps = self._get_output_length(num_forecast_steps) + + for stream_name, stream_ds in self.field_conditioning_datasets.items(): + stream_info = stream_ds.info + + # Create ONE StreamData with input_steps=num_output_steps + # Each source step corresponds to one forecast step's conditioning window + stream_data = StreamData(idx, num_forecast_steps, 1, self.num_healpix_cells) + + # Collect data for each forecast step + for step, _ in enumerate(range(self.output_offset, num_output_steps)): + # Conditioning window for forecast step N is at idx + (N - output_offset) + # For output_offset=1: forecast step 1 → idx+0, step 2 → idx+1, etc. + rdata = collect_datasources(stream_ds.readers, idx + step, "source", self.rng) + + if rdata.is_empty(): + stream_data = None + else: + token_data_list = self.tokenizer.get_tokens_windows(stream_info, [rdata], True) + token_data = token_data_list[0] + if token_data[0] is None: + stream_data = None + else: + time_win = self.time_window_handler.window(idx + step) + src_cells, src_lens = self.tokenizer.get_source( + stream_info, rdata, token_data, (time_win.start, time_win.end), None + ) + + # Add to the correct step index (step = 0, 1, 2, ...) + stream_data.add_source(self._stage, step, rdata, src_lens, src_cells, False) + + # Add the complete StreamData once (not in the loop) + batch.add_field_conditioning_stream(stream_name, stream_data) + + return batch + def __iter__(self) -> ModelBatch: """ Return one batch of data diff --git a/src/weathergen/datasets/tokenizer_masking.py b/src/weathergen/datasets/tokenizer_masking.py index 7c033e398f..f59d15916a 100644 --- a/src/weathergen/datasets/tokenizer_masking.py +++ b/src/weathergen/datasets/tokenizer_masking.py @@ -99,6 +99,9 @@ def cell_to_token_mask(self, idxs_cells, idxs_cells_lens, mask): if num_tokens == 0: return (mask_tokens, mask_channels) + if mask is None: + mask = [True] * len(idxs_cells_lens) + # TODO, TODO, TODO: use np.repeat # https://stackoverflow.com/questions/26038778/repeat-each-values-of-an-array-different-times # build token level mask: for each cell replicate the keep flag across its tokens diff --git a/src/weathergen/model/engines.py b/src/weathergen/model/engines.py index fde31213b6..0cde93d5fc 100644 --- a/src/weathergen/model/engines.py +++ b/src/weathergen/model/engines.py @@ -48,7 +48,7 @@ def __init__(self, cf: Config, sources_size) -> None: self.dtype = get_dtype(self.cf.mixed_precision_dtype) self.sources_size = sources_size # KCT:iss130, what is this? self.embeds = torch.nn.ModuleDict() - self.streams = cf.streams + self.streams = {k: v for k, v in cf.streams.items() if not v.get("timestep_conditioning")} for i, (stream_name, si) in enumerate(self.streams.items()): if si.get("diagnostic", False) or self.sources_size[i] == 0: diff --git a/src/weathergen/model/model.py b/src/weathergen/model/model.py index f4035ea467..0511bba91d 100644 --- a/src/weathergen/model/model.py +++ b/src/weathergen/model/model.py @@ -327,7 +327,10 @@ def __init__(self, cf: Config, sources_size, targets_num_channels, targets_coord self.forecast_engine: ForecastingEngine | IdentityEngine | None = None self.pred_heads = None self.q_cells: torch.Tensor | None = None - self.streams: dict[str, typing.Any] = cf.streams + self.streams: dict[str, typing.Any] = { + k: v for k, v in cf.streams.items() if not v.get("timestep_conditioning") + } + self.target_token_engines = None assert cf.get("forecast", {}).get("att_dense_rate", 1.0) == 1.0, (