treble_tsdk.collections.ir_collection

Classes

IRCollection

class treble_tsdk.collections.ir_collection.IRCollection
__init__(initial: list[IRInfo], client: TSDKClient, simulation_cache: dict[str, Simulation] | None = None, device_cache: dict[str, DeviceObj] | None = None)

Initialize the collection base.

Parameters:

  • key_column (str) – The column to use as the key for the collection (dataframe). This column is used to identify the items in the collection. Rows with the same key are considered to reference the same item.

  • default_schema (dict[str, pl.DataType]) – The default schema for the collection.

  • required_columns (list[str]) – The required columns for the collection.

  • schema_mapper (dict[str, Callable[[T], Any]]) – The schema mapper for the collection. This function will be called for each item in the collection with the item as a parameter to get the value for the column.

add_column(name: str, mapper: Callable[[T], Any], dtype: polars.DataType | None = None)

Add a column to the collection using a mapper function.

Parameters:

  • name (str) – Name of the column.

  • mapper (Callable[[T], Any]) – Mapper function to get a value. Called for each row in the collection with the item (e.g. Simulation or IRInfo) as a parameter.

  • dtype (pl.DataType) – Data type of the column. If not provided the mapper function will be used to determine the data type.

add_irs(irs: IRInfo | list[IRInfo])
add_simulations(sims: Simulation | list[Simulation])
apply(func: Callable[[polars.DataFrame], polars.DataFrame]) → _Self

Apply a DataFrame transformation and return a new collection of the same type.

The returned collection shares cached objects with the original but has an independent DataFrame. The transformation must preserve the required columns.

Parameters:

func (Callable[[pl.DataFrame], pl.DataFrame]) – A function that receives the underlying Polars DataFrame and returns a transformed DataFrame.

Returns:

A new collection of the same type with the transformed DataFrame.

Example:

filtered = coll.apply(lambda df: df.filter(pl.col("x") > 5).sort("y"))
apply_filters(ir_subset: polars.DataFrame | None = None, work_dir: str | None = None, behaviour: IRDataLoaderBehaviour = IRDataLoaderBehaviour.eager)

Apply the prescribed filters to the IRs within the IR collection.

Parameters:

  • ir_subset (pl.DataFrame) – A subset of the IR collection to apply the filters to. Defaults to None, which means all IRs in the collection.

  • work_dir (str) – The work directory to store the processed IRs.

  • behaviour (IRDataLoaderBehaviour) – The behaviour of the IR data loader.

copy_custom_columns(other_collection: CollectionBase, key_columns: list[str] | None = None, columns_to_copy: list[str] | None = None)
create_device_render_collection(devices: list[DeviceObj], device_orientations: list[Rotation], ir_subset: polars.DataFrame | None = None, orientation_application: Literal['cartesian', 'zipped'] = 'cartesian', copy_custom_columns: bool = True) → IRCollection

Create IRInfo objects for device renders without actually performing the renders.

Parameters:

  • devices – List of devices to create IRInfos for.

  • device_orientations – List of device orientations.

  • ir_subset – A subset of the IRCollection to create device render IRInfos for. If None, all IRs in the collection are used.

  • orientation_application – “cartesian” creates IRInfos for all device/orientation combinations, “zipped” pairs orientations and spatial IRs one-to-one.

  • copy_custom_columns – If True, copy custom columns from the source collection to the render collection, it tries to match rows by source_id, simulation_type and receiver_id. Acoustic parameters columns are ignored.

Returns:

IRCollection with the device render IRInfos.

enrich_with_acoustic_parameters(ir_subset: polars.DataFrame | None = None, acoustic_parameters: list[str] | None = None, work_dir: str | None = None)
estimate_device_render_cost(work_dir: str, ir_subset: polars.DataFrame | None = None) → Decimal
filter_collection(*args, **kwargs) → _Self

Filter the collection and return a new collection with matching rows.

Accepts the same arguments as polars.DataFrame.filter().

Returns:

A new collection of the same type containing only the matching rows.

Example:

nearby = coll.filter_collection(pl.col("source_receiver_dist") < 5.0)
get_data_loader(work_dir: str | None = None, behaviour: IRDataLoaderBehaviour = IRDataLoaderBehaviour.lazy) → IRDataLoader
get_ir_info(key: str) → IRInfo
head(n: int = 5) → _Self

Return a new collection with the first n rows.

Parameters:

n (int) – Number of rows, defaults to 5.

Returns:

A new collection of the same type.

perform_device_render_bulk(work_dir: str, ir_subset: polars.DataFrame | None = None, processing_batch_size: int = 100, ignore_max_frequency: bool = False, output_mode: ProgressOutputMode = ProgressOutputMode.TQDM_OUTPUT)
plot()

Plot the collection using a interactive widget.

remove_column(name: str)

Remove a column from the collection.

sample_with_distribution(column: str, n_samples: int, distribution: Gaussian | Uniform | LogNormal | Exponential | Bimodal) → _Self

Sample rows where the specified column follows a target distribution.

Returns a new collection of the same type containing only the sampled rows.

Parameters:

  • column (str) – The column to sample from.

  • n_samples (int) – The number of rows to sample.

  • distribution (Distribution) – Distribution to sample from.

Returns:

A new collection containing the sampled rows.

sample_with_gaussian_distribution(column: str, n_samples: int, target_mean: float, target_std: float, seed: int = 42)

Sample rows where the specified column follows a target Gaussian distribution.

This is a convenience wrapper around sample_with_distribution.

Parameters:

  • column (str) – The column to sample from.

  • n_samples (int) – The number of rows to sample.

  • target_mean (float) – The target mean for the distribution.

  • target_std (float) – The target standard deviation for the distribution.

  • seed (int) – The random seed for reproducibility.

set_filters(filters: list[FilterDefinition], ir_subset: polars.DataFrame | None = None)

Prescribe a chain of filters for a set of IRs within the IR collection.

Parameters:

  • filters (list[FilterDefinition]) – A list of filter definitions to apply to the IRs.

  • ir_subset (pl.DataFrame) – A subset of the IR collection to set the filters for. Defaults to None, which means all IRs.

tail(n: int = 5) → _Self

Return a new collection with the last n rows.

Parameters:

n (int) – Number of rows, defaults to 5.

Returns:

A new collection of the same type.

write_parquet(path: str | Path, **kwargs)

Write the dataframe to a parquet file.

Parameters:

  • path (str | Path) – Path to the parquet file.

  • kwargs – Additional keyword arguments passed to polars.DataFrame.write_parquet.

property columns: list[str]

List of columns in the DataFrame

property dataframe: polars.DataFrame

Get the collection data as a Polars DataFrame.

Returns pl.DataFrame:

DataFrame containing all collection items and computed columns.

property editable_dataframe: polars.DataFrame

DataFrame with only the user-modifiable (non-locked) columns. The key column is always included in the editable dataframe although it is locked.

property lazy_df: polars.LazyFrame

Lazy access to underlying Polars DataFrame

property locked_columns: list[str]

List of locked columns for the collection. Columns in this list are read-only and cannot be modified by the user.

property required_columns: list[str]

List of required columns for the collection.