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): This method has been removed. Please use start_remote_processing instead

as_remote_processing_live_progress(show_all_tasks: bool = False)

cancel_remote_processing(): Cancel remote IR processing for the IR collection. Cancels all remote IR processing tasks for the IR collection.

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') → 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)

Enrich the IR collection with acoustic parameters. Adds the requested acoustic_parameters columns to the dataframe.

Parameters:

ir_subset – A subset of the IRCollection to enrich with acoustic parameters. If None, all IRs in the collection are used.
acoustic_parameters – A list of acoustic parameters to enrich the IR collection with. If None, all available acoustic parameters are used.
work_dir – The work directory to use for the acoustic parameters. If None, a temporary directory is used.

estimate_device_render_cost(work_dir: str, ir_subset: polars.DataFrame | None = None) → Decimal: This method has been removed.

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 | Path | None = None, behaviour: IRDataLoaderBehaviour = IRDataLoaderBehaviour.lazy) → IRDataLoader

get_ir_info(key: str) → IRInfo

get_remote_processing_progress() → OneOffTaskBundleProgress | None

Fetch live progress for all remote IR processing tasks shipped from this collection.

Returns None when the collection has no shipped processing tasks.

get_remote_processing_token_cost(optimize_mono_download: bool = False) → IRCollectionProcessingCost

Get the token cost of the remote IR processing tasks for the IR collection.

Parameters:: optimize_mono_download (bool) – Must match the flag passed to start_remote_processing so the quote covers the same work.

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): This method has been removed. Please use start_remote_processing instead

plot(): Plot the collection using a interactive widget.

print_remote_processing_token_cost(optimize_mono_download: bool = False)

Print the token cost of the remote IR processing tasks for the IR collection.

Parameters:: optimize_mono_download (bool) – Must match the flag passed to start_remote_processing so the quote covers the same work.

remove_column(name: str): Remove a column from the collection.

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_device_render_info(params: DeviceRenderParams | list[DeviceRenderParams] | dict[str, DeviceRenderParams | list[DeviceRenderParams]], ir_subset: pl.DataFrame | None = None)

Attach device render info to a set of IRs in the collection (in-place).

A scalar assigns a single device render to each affected base IR, while a list fans the base IR out into one new row per DeviceRenderParams. The base IR is always consumed; for a list of length N each affected row is replaced by N device-render variants. Acoustic-parameter columns (if present) are nulled on the new rows since they no longer apply to the rendered IR.

Call styles:

Broadcast (scalar) – params is a single DeviceRenderParams and every row in ir_subset (or the whole collection if ir_subset is None) is rewritten with the same device/orientation.
Broadcast (fan-out) – params is a list[DeviceRenderParams] and every affected row is expanded into one variant per list entry.
Per-id mapping – params is a dict[str, DeviceRenderParams | list[DeviceRenderParams]] keyed by IRInfo.id so each affected row gets its own device/orientation (or its own fan-out when the value is a list). ir_subset must be None in this case because the mapping’s keys define the affected set.

Moving IRs are skipped with a warning in broadcast mode and rejected in the per-id mapping mode, since device renders are not supported for them.

Parameters:

params – The device-render parameters to apply (scalar broadcast, list fan-out, or per-id mapping of either).
ir_subset – A subset of the IR collection to update in broadcast mode. Must be None when a per-id mapping is supplied. Defaults to None, which means all IRs.

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

Prescribe filters for a set of IRs within the IR collection (in-place).

A single filter chain assigns that chain to each affected base IR, while a list of chains fans the base IR out into one new row per chain. The base IR is always consumed; for a fan-out of N chains each affected row is replaced by N filtered variants.

Call styles:

Broadcast (single chain) – filters is a list[FilterDefinition] and every row in ir_subset (or the whole collection if ir_subset is None) is rewritten with the same filter chain.
Broadcast (fan-out) – filters is a list[list[FilterDefinition]] and every affected row is expanded into one variant per chain.
Per-id mapping – filters is a dict[str, list[FilterDefinition] | list[list[FilterDefinition]]] keyed by IRInfo.id so each affected row gets its own chain (or its own fan-out when the value is a list of chains). ir_subset must be None in this case because the mapping’s keys define the affected set.

Parameters:

filters – A single filter chain to broadcast, a list of chains to fan out, or a per-id mapping of either.
ir_subset – A subset of the IR collection to set the filters for. Only valid in broadcast mode. Defaults to None, which means all IRs.

start_remote_processing(optimize_mono_download: bool = False)

Perform remote IR processing for the IR collection. Processes all IRs according to their processing recipe (Device renders and filters). Any IRs that have already been processed remotely are skipped.

Parameters:: optimize_mono_download (bool) – When True remote IR processing will include IRs that do not need device renders and/or filters and package their mono channels into a package for a more optimized download.

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 is_mutable: bool

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 remote_processing_status: SolveTaskStatus | None

Aggregate status of the remote IR-processing tasks backing this collection.

Returns:: The aggregate status, or None if this collection has no remote processing tasks.

property required_columns: list[str]: List of required columns for the collection.