peal.data.dataset_utils.process_confounder_data_controlled

peal.data.dataset_utils.process_confounder_data_controlled(raw_data, config, mode, extract_instances_tensor, set_negative_to_zero=True)[source]

Subsample rows so that target and confounder co-occur at a fixed rate.

Rows are consumed in file order and sorted into the four (attribute, confounder) groups until every group reaches its quota, which is derived from config.num_samples and either config.full_confounder_config (four fractions summing to 1) or config.confounder_probability (the aligned groups (0,0) and (1,1) each get p/2, the misaligned ones (1-p)/2). Rows with attribute or confounder values >= 2 are skipped. Each group is then split with config.split so that all partitions keep the same confounder ratio.

Parameters:

  • raw_data (iterable) – Rows (csv lines or json records) passed to extract_instances_tensor.

  • config (DataConfig) – Provides num_samples, split and the confounder settings.

  • mode (str) – "train", "val", "test"; otherwise all selected keys.

  • extract_instances_tensor (callable) – f(idx, line) -> (key, tensor, attribute, confounder).

  • set_negative_to_zero (bool, optional) – Clamp negative tensor values (and the attribute/confounder ints) to 0.

Returns:

  • data (dict) – Mapping key -> tensor for every parsed row (not only selected ones).

  • keys_out (list) – Selected keys of the requested split, shuffled in place with the module-level random state.

Raises:

AssertionError – If the quotas cannot be filled from the available rows or the totals do not add up to config.num_samples.