dcegm.pre_processing.batches.single_segment

Functions

create_single_segment_of_batches(...)	Create a single segment of evenly sized batches.
correct_for_uneven_last_batch(batches_list, ...)	Check if the last batch has the same length as the others.
prepare_and_align_batch_arrays(batches_list, ...)	Prepare the lists we get out of the algorithm (and after correction) for the jax
extend_child_state_choices_to_aggregate_choices(...)	In case of uneven batches, we need to extend the child state objects to cover the

Module Contents

dcegm.pre_processing.batches.single_segment.create_single_segment_of_batches(bool_state_choices_to_batch, model_structure)

Create a single segment of evenly sized batches.

If the last batch is not evenly we correct it.

dcegm.pre_processing.batches.single_segment.correct_for_uneven_last_batch(batches_list, child_states_to_integrate_stochastic_list, child_state_choices_to_aggr_choice_list, child_state_choice_idxs_to_interp_list, state_choice_space_dict, map_state_choice_to_parent_state)

Check if the last batch has the same length as the others.

If not, we need to save the information separately.

dcegm.pre_processing.batches.single_segment.prepare_and_align_batch_arrays(batches_list, child_states_to_integrate_stochastic_list, child_state_choices_to_aggr_choice_list, child_state_choice_idxs_to_interp_list, state_choice_space_dict, map_state_choice_to_parent_state, discrete_states_names)

Prepare the lists we get out of the algorithm (and after correction) for the jax calculations.

They all need to have the same length of leading axis

dcegm.pre_processing.batches.single_segment.extend_child_state_choices_to_aggregate_choices(idx_to_aggregate_choice, max_child_state_index_batch, idx_to_interpolate, out_of_bounds_state_choice_idx)

In case of uneven batches, we need to extend the child state objects to cover the same number of state choices in each batch.

As this object has in each batch the shape of n_state_choices x n_