ecoli.library.parquet_emitter

ecoli.library.parquet_emitter.METADATA_PREFIX = 'output_metadata__'

In the config dataset, user-defined metadata for each store (see get_output_metadata()) will be contained in columns with this prefix.

class ecoli.library.parquet_emitter.ParquetEmitter(config)

Bases: Emitter

Emit data to a Parquet dataset.

Configure emitter.

Parameters:

config (dict[str, Any]) –

Should be a dictionary as follows:

{
    'type': 'parquet',
    'emits_to_batch': Number of emits per Parquet row
        group (optional, default: 400),
    # One of the following is REQUIRED
    'out_dir': local output directory (absolute/relative),
    'out_uri': Google Cloud storage bucket URI
}

_finalize()

Convert remaining batched emits to Parquet at sim shutdown and mark sim as successful if success flag was set. In vEcoli, this is done by EcoliSim upon reaching division.

emit(data)

Flattens emit dictionary by concatenating nested key names with double underscores (flatten_dict()), serializes flattened emit with orjson, and writes newline-delimited JSONs in a temporary file to be batched for some number of timesteps before conversion to Parquet by json_to_parquet().

The output directory (config["out_dir"] or config["out_uri"]) will have the following structure:

{experiment_id}
|-- history
|   |-- experiment_id={experiment_id}
|   |   |-- variant={variant}
|   |   |   |-- lineage_seed={seed}
|   |   |   |   |-- generation={generation}
|   |   |   |   |   |-- agent_id={agent_id}
|   |   |   |   |   |   |-- 400.pq (batched emits)
|   |   |   |   |   |   |-- 800.pq
|   |   |   |   |   |   |-- ..
|-- configuration
|   |-- experiment_id={experiment_id}
|   |   |-- variant={variant}
|   |   |   |-- lineage_seed={seed}
|   |   |   |   |-- generation={generation}
|   |   |   |   |   |-- agent_id={agent_id}
|   |   |   |   |   |   |-- config.pq (sim config data)

This Hive-partioned directory structure can be efficiently filtered and queried using DuckDB (see get_dataset_sql()).

Parameters:

data (dict[str, Any])

ecoli.library.parquet_emitter.USE_UINT16 = {'listeners__complexation_listener__complexation_events', 'listeners__ribosome_data__mRNA_TU_index', 'listeners__ribosome_data__n_ribosomes_on_each_mRNA', 'listeners__rna_degradation_listener__count_RNA_degraded_per_cistron', 'listeners__rna_degradation_listener__count_rna_degraded', 'listeners__rna_synth_prob__bound_TF_domains', 'listeners__rna_synth_prob__bound_TF_indexes', 'listeners__rna_synth_prob__expected_rna_init_per_cistron', 'listeners__rna_synth_prob__gene_copy_number', 'listeners__rna_synth_prob__n_bound_TF_per_TU', 'listeners__rna_synth_prob__n_bound_TF_per_cistron', 'listeners__rna_synth_prob__promoter_copy_number', 'listeners__rnap_data__active_rnap_domain_indexes', 'listeners__rnap_data__active_rnap_n_bound_ribosomes', 'listeners__rnap_data__rna_init_event', 'listeners__rnap_data__rna_init_event_per_cistron', 'listeners__transcript_elongation_listener__count_rna_synthesized'}

uint16 is 4x smaller than int64 for values between 0 - 65,535.

ecoli.library.parquet_emitter.USE_UINT32 = {'listeners__fba_results__catalyst_counts', 'listeners__monomer_counts', 'listeners__ribosome_data__n_ribosomes_on_partial_mRNA_per_transcript', 'listeners__ribosome_data__n_ribosomes_per_transcript', 'listeners__ribosome_data__ribosome_init_event_per_monomer', 'listeners__rna_counts__full_mRNA_cistron_counts', 'listeners__rna_counts__full_mRNA_counts', 'listeners__rna_counts__mRNA_cistron_counts', 'listeners__rna_counts__mRNA_counts', 'listeners__rna_counts__partial_mRNA_cistron_counts', 'listeners__rna_counts__partial_mRNA_counts'}

uint32 is 2x smaller than int64 for values between 0 - 4,294,967,295.

ecoli.library.parquet_emitter.flatten_dict(d)

Flatten nested dictionary down to key-value pairs where each key concatenates all the keys needed to reach the corresponding value in the input. Allows each leaf field in a nested emit to be turned into a column in a Parquet file for efficient storage and retrieval.

Parameters:

d (dict)

ecoli.library.parquet_emitter.get_config_value(conn, config_subquery, field)

Gets the saved configuration option (anything in config JSON, with double underscore concatenation for nested fields due to flatten_dict()).

Parameters:

• conn (DuckDBPyConnection) – DuckDB connection

• config_subquery (str) – DuckDB query containing sim config data

• field (str) – Name of configuration option to get value of

Return type:

Any

ecoli.library.parquet_emitter.get_dataset_sql(out_dir, experiment_ids)

Creates DuckDB SQL strings for sim outputs, configs, and metadata on which sims were successful.

Parameters:

• out_dir (str) – Path to output directory for workflows to retrieve data for (relative or absolute local path OR URI beginning with gcs:// or gs:// for Google Cloud Storage bucket)

• experiment_ids (list[str]) – List of experiment IDs to include in query. To read data from more than one experiment ID, the listeners in the output of the first experiment ID in the list must be a strict subset of the listeners in the output of the subsequent experiment ID(s).

Returns:

3-element tuple containing

• history_sql: SQL query for sim output (see read_stacked_columns()),

• config_sql: SQL query for sim configs (see get_field_metadata() and get_config_value())

• success_sql: SQL query for metadata marking successful sims (see read_stacked_columns())

Return type:

tuple[str, str, str]

ecoli.library.parquet_emitter.get_encoding(val, field_name, use_uint16=False, use_uint32=False)

Get optimal PyArrow type and Parquet encoding for input value.

Parameters:

• val (Any)

• field_name (str)

• use_uint16 (bool)

• use_uint32 (bool)

Return type:

tuple[Any, str, str, bool]

ecoli.library.parquet_emitter.get_field_metadata(conn, config_subquery, field)

Gets the saved metadata (see get_output_metadata()) for a given field as a list.

Parameters:

• conn (DuckDBPyConnection) – DuckDB connection

• config_subquery (str) – DuckDB query containing sim config data

• field (str) – Name of field to get metadata for

Return type:

list

ecoli.library.parquet_emitter.get_plot_metadata(conn, config_subquery, variant_name)

Gets dictionary that can be used as metadata kwarg to wholecell.utils.plotting_tools.export_figure().

Parameters:

• conn (DuckDBPyConnection) – DuckDB connection

• config_subquery (str) – DuckDB query containing sim config data

• variant_name (str) – Name of variant

Return type:

dict[str, Any]

ecoli.library.parquet_emitter.json_to_parquet(ndjson, encodings, schema, outfile, filesystem)

Reads newline-delimited JSON file and converts to Parquet file.

Parameters:

• ndjson (str) – Path to newline-delimited JSON file.

• encodings (dict[str, str]) – Mapping of column names to Parquet encodings

• schema (Schema) – PyArrow schema of Parquet file to write

• outfile (str) – Filepath of output Parquet file

• filesystem (FileSystem) – PyArrow filesystem for Parquet output (local if None)

ecoli.library.parquet_emitter.named_idx(col, names, idx)

Create DuckDB expressions for given indices from a list column. Can be used in projection kwarg of read_stacked_columns(). Since each index gets pulled out into its own column, this greatly simplifies aggregations like averages, etc. Only use this if the number of indices is relatively small (<100) and the list column is 1-dimensional. For 2+ dimensions or >100 indices, see ndidx_to_duckdb_expr().

Warning

DuckDB arrays are 1-indexed so this function adds 1 to every supplied index!

Parameters:

• col (str) – Name of list column.

• names (list[str]) – New column names, one for each index.

• idx (list[int]) – Indices to retrieve from col

Returns:

DuckDB SQL expression for a set of named columns corresponding to the values at given indices of a list column

Return type:

str

ecoli.library.parquet_emitter.ndarray_to_ndlist(arr)

Convert a Numpy ndarray into a PyArrow FixedSizeListArray. This is useful for writing user-defined functions (see DuckDB documentation on Python Function API and func kwarg for read_stacked_columns()) that expect a PyArrow return type.

Note that the number of rows in the returned PyArrow array is equal to the size of the first dimension of the input array. This means for a 3 x 4 x 5 Numpy array, the return PyArrow array will have 3 rows where each row is a nested list with 4 lists of length 5.

Parameters:

arr (ndarray)

Return type:

FixedSizeListArray

ecoli.library.parquet_emitter.ndidx_to_duckdb_expr(name, idx)

Returns a DuckDB expression for a column equivalent to converting each row of name into an ndarray name_arr (ndlist_to_ndarray()) and getting name_arr[idx]. idx can contain 1D lists of integers, boolean masks, or ":" (no 2D+ indices like x[[[1,2]]]). See also named_idx() if pulling out a relatively small set of indices.

Warning

DuckDB arrays are 1-indexed so this function adds 1 to every supplied integer index!

Parameters:

• name (str) – Name of column to recursively index

• idx (list[int | list[int] | list[bool] | str]) –

  To get all elements for a dimension, supply the string ":". Otherwise, only single integers or 1D integer lists of indices are allowed for each dimension. Some examples:

  [0, 1] # First row, second column
  [[0, 1], 1] # First and second row, second column
  [0, 1, ":"] # First element of axis 1, second of 2, all of 3
  # Final example differs between this function and Numpy
  # This func: 1st and 2nd of axis 1, all of 2, 1st and 2nd of 3
  # Numpy: Complicated, see Numpy docs on advanced indexing
  [[0, 1], ":", [0, 1]]
  

Return type:

str

ecoli.library.parquet_emitter.ndlist_to_ndarray(s)

Convert a PyArrow series of nested lists with fixed dimensions into a Numpy ndarray. This should really only be necessary if you are trying to perform linear algebra (e.g. matrix multiplication, dot products) inside a user-defined function (see DuckDB documentation on Python Function API and func kwarg for read_stacked_columns()).

For elementwise arithmetic with two nested list columns, this can be used in combination with ndarray_to_ndlist() to define a custom DuckDB function as follows:

import duckdb
from ecoli.library.parquet_emitter import (
    ndarray_to_ndlist, ndlist_to_ndarray)
def sum_arrays(col_0, col_1):
    return ndarray_to_ndlist(
        ndlist_to_ndarray(col_0) +
        ndlist_to_ndarray(col_1)
    )
conn = duckdb.connect()
conn.create_function(
    "sum_2d_int_arrays", # Function name for use in SQL (must be unique)
    sum_arrays, # Python function that takes and returns PyArrow arrays
    [list[list[int]], list[list[int]]], # Input types (2D lists here)
    list[list[int]], # Return type (2D list here)
    type = "arrow" # Tell DuckDB function operates on Arrow arrays
)
conn.sql("SELECT sum_2d_int_arrays(int_col_0, int_col_1) from input_table")
# Note that function must be registered under different name for each
# set of unique input/output types
conn.create_function(
    "sum_2d_int_and_float",
    sum_arrays,
    [list[list[int]], list[list[float]]], # Second input is 2D float array
    list[list[float]], # Adding int to float array gives float in Numpy
    type = "arrow"
)
conn.sql("SELECT sum_2d_int_arrays(int_col_0, float_col_0) from input_table")

Parameters:

s (Array)

Return type:

ndarray

ecoli.library.parquet_emitter.num_cells(conn, subquery)

Return cell count in DuckDB subquery containing experiment_id, variant, lineage_seed, generation, and agent_id columns.

Parameters:

• conn (DuckDBPyConnection)

• subquery (str)

Return type:

int

ecoli.library.parquet_emitter.open_arbitrary_sim_data(sim_data_dict)

Given a mapping from experiment ID(s) to mappings from variant ID(s) to sim_data path(s), pick an arbitrary sim_data to read.

Parameters:

sim_data_dict (dict[str, dict[int, Any]]) – Generated by runscripts.analysis and passed to each analysis script as an argument.

Returns:

PyArrow file object for arbitrarily chosen sim_data to be loaded with pickle.load

Return type:

NativeFile

ecoli.library.parquet_emitter.open_output_file(outfile)

Open a file by its path, whether that be a path on local storage or Google Cloud Storage.

Parameters:

outfile (str) – Path to file. Must have gs:// or gcs:// prefix if on Google Cloud Storage. Can be relative or absolute path if on local storage.

Returns:

PyArrow file object that supports reading, seeking, etc. in bytes

Return type:

NativeFile

ecoli.library.parquet_emitter.read_stacked_columns(history_sql, columns, remove_first=False, func=None, conn=None, order_results=True, success_sql=None)

Loads columns for many cells. If you would like to perform more advanced computatations (aggregations, window functions, etc.) using the optimized DuckDB API, you can omit conn, in which case this function will return an SQL string that can be used as a subquery. For computations that cannot be easily performed using the DuckDB API, you can define a custom function func that will be called on the data for each cell. By default, the return value (whether it be the actual data or an SQL subquery) will also include the experiment_id, variant, lineage_seed, generation, agent_id, and time columns.

For example, to get the average total concentration of three bulk molecules with indices 100, 1000, and 10000 per cell:

import duckdb
from ecoli.library.parquet_emitter import (
    get_dataset_sql, read_stacked_columns)
history_sql, config_sql, _ = get_dataset_sql('out/', 'exp_id')
subquery = read_stacked_columns(
    history_sql,
    # Note DuckDB arrays are 1-indexed
    ["bulk[100 + 1] + bulk[1000 + 1] + bulk[10000 + 1] AS bulk_sum",
    "listeners__enzyme_kinetics__counts_to_molar AS counts_to_molar"],
    order_results=False,
)
query = '''
    SELECT avg(bulk_sum * counts_to_molar) AS avg_total_conc
    FROM ({subquery})
    GROUP BY experiment_id, variant, lineage_seed, generation, agent_id
    '''
conn = duckdb.connect()
data = conn.sql(query).arrow()

Here is a more complicated example that defines a custom function to get the per-cell average RNA synthesis probability per cistron:

import duckdb
import pickle
import pyarrow as pa
from ecoli.library.parquet_emitter import (
    get_dataset_sql, ndlist_to_ndarray, read_stacked_columns)
history_sql, config_sql, _ = get_dataset_sql('out/', 'exp_id')
# Load sim data
with open("reconstruction/sim_data/kb/simData.cPickle", "rb") as f:
    sim_data = pickle.load(f)
# Get mapping from RNAs (TUs) to cistrons
cistron_tu_mat = sim_data.process.transcription.cistron_tu_mapping_matrix
# Custom aggregation function with Numpy dot product and mean
def avg_rna_synth_prob_per_cistron(rna_synth_prob):
    # Convert rna_synth_prob into 2-D Numpy array (time x TU)
    rna_synth_prob = ndlist_to_ndarray(rna_synth_prob[
        "listeners__rna_synth_prob__actual_rna_synth_prob"])
    rna_synth_prob_per_cistron = cistron_tu_mat.dot(rna_synth_prob.T).T
    # Return value must be a PyArrow table
    return pa.table({'avg_rna_synth_prob_per_cistron': [
        rna_synth_prob_per_cistron.mean(axis=0)]})
conn = duckdb.connect()
result = read_stacked_columns(
    history_sql,
    ["listeners__rna_synth_prob__actual_rna_synth_prob"],
    func=avg_rna_synth_prob_per_cistron,
    conn=conn,
)

Parameters:

• history_sql (str) – DuckDB SQL string from get_dataset_sql(), potentially with filters appended in WHERE clause

• columns (list[str]) – Names of columns to read data for. Alternatively, DuckDB expressions of columns (e.g. avg(listeners__mass__cell_mass) AS avg_mass or the output of named_idx() or ndidx_to_duckdb_expr()).

• remove_first (bool) – Remove data for first timestep of each cell

• func (Callable[[Table], Table] | None) – Function to call on data for each cell, should take and return a PyArrow table with columns equal to columns

• conn (DuckDBPyConnection | None) – DuckDB connection instance with which to run query. Typically provided by runscripts.analysis.main() to the plot method of analysis scripts (tweaked some DuckDB settings). Can be omitted to return SQL query string to be used as subquery instead of running query immediately and returning result.

• order_results (bool) – Whether to sort returned table by experiment_id, variant, lineage_seed, generation, agent_id, and time. If no conn is provided, this can usually be disabled and any sorting can be deferred until the last step in the query with a manual ORDER BY. Doing this can greatly reduce RAM usage.

• success_sql (str | None) – Final DuckDB SQL string from get_dataset_sql(). If provided, will be used to filter out unsuccessful sims.

Return type:

Table | str

ecoli.library.parquet_emitter.skip_n_gens(subquery, n)

Modifies a DuckDB SQL query to skip the first n generations of data.

Parameters:

• subquery (str)

• n (int)

Return type:

str

ecoli.library.parquet_emitter.union_by_name(query_sql)

Modifies SQL query string from get_dataset_sql() to include union_by_name = true in the DuckDB read_parquet function. This allows data to be read from simulations that have different columns by filling in nulls and casting as necessary. This comes with a performance penalty and should be avoided if possible.

Parameters:

query_sql (str) – SQL query string from get_dataset_sql()

Return type:

str