treble_tsdk.collections.simulation_collection
Classes
- class treble_tsdk.collections.simulation_collection.SimulationCollection
- __init__(initial: list[Simulation], client: TSDKClient)
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_simulations(sims: Simulation | list[Simulation])
Add simulations to the collection.
- 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"))
- copy_custom_columns(other_collection: CollectionBase, key_columns: list[str] | None = None, columns_to_copy: list[str] | None = None)
- 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_ir_collection(simulation_subset: polars.DataFrame | None = None, inherit_columns: list[dict[str, str]] | None = None) IRCollection
Convert the collection to an IR collection.
- Parameters:
simulation_subset – A subset of the simulation collection to convert to an IR collection. If not provided, all simulations will be converted.
inherit_columns – A list of column mappings to inherit from the simulation collection. Each dict maps source column name (key) to target column name (value) in the IR collection. For example:
[{"tags": "simulation_tags"}]maps"tags"to"simulation_tags".
- 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.
- 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.
- 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.
- 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 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.