TFMMBasePool

class TFMMBasePool[source]

Bases: AbstractPool

TFMMBasePool is an abstract base class for implementing TFMM (Temporal Function Market Making) liquidity pools.

This class extends the AbstractPool class and provides a foundation for specific TFMM pool implementations. It defines additional abstract methods that are specific to TFMM pools, such as weight calculation.

Abstract Methods:: calculate_rule_outputs: Calculate the raw weight outputs of assets in the pool based on oracle values and parameters. calculate_fine_weights: Function to handle how raw weights get mapped to per-block/per-minute weights. Two standard methods are provided, for when 1) rules output raw weight _changes_ and 2) when rule output raw _weights_ themselves. See MomentumPool and MinVariancePool as prototypical examples of each respectively.

In addition to the methods from AbstractPool, subclasses of TFMMBasePool must implement these TFMM-specific methods to define the behavior of the pool.

Note

This class is designed to be subclassed, not instantiated directly. Concrete implementations should provide specific logic for weight calculation and slippage estimation. It is recommended to implement the functions used within implementations of these methods as external JAX functions that are jitted and then used within pool methods. This separation of concerns comes from that JAX is a functional programming language and we want to keep the pool methods pure. Finally, note that due to this separation of concerns this class does not hold any state, for example pool parameters.

__init__()[source]: Initialize a new TFMMBasePool instance.

calculate_reserves_with_fees(params, run_fingerprint, prices, start_index=None, additional_oracle_input=None, weights=None, local_prices=None, initial_reserves=None)[source]

Calculate reserves with fees and dynamic weights.

TFMM pools calculate weights dynamically based on price history. This method handles the full complexity of weight adjustments, fees, and arbitrage.

Implementation Steps:

Extract local price window
Calculate dynamic weights based on price history
Apply arbitrage frequency adjustments
Initialize reserves based on pool value
Calculate reserve changes using quantAMM precalcs

param params:: Pool parameters including weight calculation parameters
type params:: Dict[str, Any]
param run_fingerprint:: Simulation settings including: - bout_length: Simulation window length - n_assets: Number of tokens - arb_frequency: Arbitrage check frequency - initial_pool_value: Starting pool value - fees: Trading fees - gas_cost: Arbitrage threshold - arb_fees: Arbitrage fees - do_arb: Enable arbitrage - all_sig_variations: Valid trade combinations
type run_fingerprint:: Dict[str, Any]
param prices:: Historical price data
type prices:: jnp.ndarray
param start_index:: Window start position
type start_index:: jnp.ndarray
param additional_oracle_input:: Extra data for weight calculation
type additional_oracle_input:: Optional[jnp.ndarray]
returns:: Time series of pool reserves
rtype:: jnp.ndarray

Parameters:

params (Dict[str, Any])
run_fingerprint (Dict[str, Any])
prices (Array)
start_index (Array | None)
additional_oracle_input (Array | None)
weights (Array | None)
local_prices (Array | None)
initial_reserves (Array | None)

Return type:

Array

calculate_reserves_zero_fees(params, run_fingerprint, prices, start_index, additional_oracle_input=None)[source]

Parameters:

params (Dict[str, Any])
run_fingerprint (Dict[str, Any])
prices (Array)
start_index (Array)
additional_oracle_input (Array | None)

Return type:

Array

calculate_reserves_with_dynamic_inputs(params, run_fingerprint, prices, start_index, fees_array, arb_thresh_array, arb_fees_array, trade_array, lp_supply_array=None, additional_oracle_input=None)[source]

Parameters:

params (Dict[str, Any])
run_fingerprint (Dict[str, Any])
prices (Array)
start_index (Array)
fees_array (Array)
arb_thresh_array (Array)
arb_fees_array (Array)
trade_array (Array)
lp_supply_array (Array)
additional_oracle_input (Array | None)

Return type:

Array

calculate_rule_outputs(params, run_fingerprint, prices, additional_oracle_input=None)[source]

Calculate raw weight adjustments based on price history (vectorized path).

This is the first step in TFMM’s two-step weight calculation process. Subclasses should implement either this method OR calculate_rule_output_step.

Parameters:

params (Dict[str, Any]) – Pool parameters for weight calculation
run_fingerprint (Dict[str, Any]) – Simulation settings
prices (jnp.ndarray) – Historical price data
additional_oracle_input (Optional[jnp.ndarray]) – Extra data for weight calculation

Returns:

Raw weight adjustment values

Return type:

jnp.ndarray

calculate_rule_output_step(carry, price, params, run_fingerprint)[source]

Calculate a single step of weight update (scan-based path).

This method represents how the strategy would run in production, where we are given current state, receive a new price observation, and output new state along with the weight update for this timestep.

This is the core primitive that enables causality-preserving simulation. The state (carry) contains all information needed to compute the next step without any look-ahead bias.

Subclasses should implement either this method OR calculate_rule_outputs.

Parameters:

carry (Dict[str, jnp.ndarray]) – Current state containing estimator variables. Typical keys include: - ‘ewma’: Exponentially weighted moving average of prices (shape: n_assets,) - ‘running_a’: Running accumulator for gradient estimation (shape: n_assets,) Additional keys may be present depending on the pool implementation.
price (jnp.ndarray) – Current price observation (shape: n_assets,)
params (Dict[str, Any]) – Pool parameters (k, lamb, etc.)
run_fingerprint (Dict[str, Any]) – Simulation settings (chunk_period, max_memory_days, etc.)

Returns:

(new_carry, rule_output) where: - new_carry: Updated state dict with same structure as input carry - rule_output: Weight update/output for this timestep (shape: n_assets,)

Return type:

tuple

get_initial_rule_state(initial_price, params, run_fingerprint)[source]

Initialize the carry state for scanning.

This creates the initial state needed to begin the scan-based weight calculation. The initial state is typically derived from the first price observation.

Required if using scan-based path (calculate_rule_output_step).

Parameters:

initial_price (jnp.ndarray) – First price observation (shape: n_assets,)
params (Dict[str, Any]) – Pool parameters
run_fingerprint (Dict[str, Any]) – Simulation settings

Returns:

Initial carry state with keys appropriate for this pool type.

Return type:

Dict[str, jnp.ndarray]

supports_vectorized_path()[source]

Check if pool implements vectorized weight calculation.

Returns True if this pool class overrides calculate_rule_outputs, indicating it supports the vectorized (convolution-based) path.

Returns:: True if vectorized path is supported.
Return type:: bool

supports_scan_path()[source]

Check if pool implements scan-based weight calculation.

Returns True if this pool class overrides calculate_rule_output_step, indicating it supports the scan-based path.

Returns:: True if scan path is supported.
Return type:: bool

calculate_rule_outputs_scan(params, run_fingerprint, prices, additional_oracle_input=None)[source]

Calculate raw weight outputs using jax.lax.scan over single-step updates.

This method produces the same outputs as calculate_rule_outputs, but uses an explicit scan loop over the single-step update method. This mirrors how the strategy would be executed in production, where we process one price at a time.

Parameters:

params (Dict[str, Any]) – Pool parameters
run_fingerprint (Dict[str, Any]) – Simulation settings
prices (jnp.ndarray) – Historical price data (shape: time_steps, n_assets)
additional_oracle_input (Optional[jnp.ndarray]) – Extra data for weight calculation (not used in scan-based approach)

Returns:

Raw weight outputs with same shape and values as calculate_rule_outputs

Return type:

jnp.ndarray

get_initial_guardrail_state(initial_weights)[source]

Initialize the weight carry state for scanning with guardrails.

Parameters:: initial_weights (jnp.ndarray) – Initial portfolio weights (shape: n_assets,)
Returns:: Initial weight carry state with ‘prev_actual_weight’
Return type:: Dict[str, jnp.ndarray]

calculate_coarse_weight_step(estimator_carry, weight_carry, price, params, run_fingerprint)[source]

Compute raw weight update and apply guardrails for a single step.

This method calls calculate_rule_output_step to get the raw weight output, then applies guardrails (normalization, min/max constraints, max change limits).

Parameters:

estimator_carry (Dict[str, jnp.ndarray]) – Current estimator state (ewma, running_a, etc.)
weight_carry (Dict[str, jnp.ndarray]) – Current weight state with ‘prev_actual_weight’
price (jnp.ndarray) – Current price observation (shape: n_assets,)
params (Dict[str, Any]) – Pool parameters
run_fingerprint (Dict[str, Any]) – Simulation settings

Returns:

(new_estimator_carry, new_weight_carry, step_output) where: - new_estimator_carry: Updated estimator state - new_weight_carry: Updated weight state with ‘prev_actual_weight’ - step_output: Dict with ‘actual_start’, ‘scaled_diff’, ‘target_weight’

Return type:

tuple

calculate_fine_weights_step(estimator_carry, weight_carry, price, params, run_fingerprint)[source]

Compute a single interpolation block of fine weights for one price step.

This method calls calculate_coarse_weight_step to get the guardrailed weight outputs, then generates the full interpolation block.

Parameters:

estimator_carry (Dict[str, jnp.ndarray]) – Current estimator state
weight_carry (Dict[str, jnp.ndarray]) – Current weight state with ‘prev_actual_weight’
price (jnp.ndarray) – Current price observation
params (Dict[str, Any]) – Pool parameters
run_fingerprint (Dict[str, Any]) – Simulation settings

Returns:

(new_estimator_carry, new_weight_carry, interpolation_block) where: - new_estimator_carry: Updated estimator state - new_weight_carry: Updated weight state - interpolation_block: Array of shape (chunk_period, n_assets)

Return type:

tuple

calculate_weights_scan(params, run_fingerprint, prices, start_index, additional_oracle_input=None)[source]

Calculate fine weights using sequential single-step interpolation blocks.

This method produces the same outputs as calculate_weights, but uses a truly sequential approach: 1. Warm up the estimator over the burn-in period (single-step updates) 2. Reset weight state to initial_weights at bout start 3. Scan over bout prices using calculate_fine_weights_step 4. Concatenate interpolation blocks

This mirrors how weights would be computed in a production system processing prices one step at a time.

Parameters:

params (Dict[str, Any]) – Pool parameters
run_fingerprint (Dict[str, Any]) – Simulation settings including chunk_period, bout_length, n_assets
prices (jnp.ndarray) – Full price history including burn-in period
start_index (jnp.ndarray) – Start index for the bout period (after burn-in)
additional_oracle_input (Optional[jnp.ndarray]) – Extra data for weight calculation

Returns:

Fine weights matching calculate_weights output

Return type:

jnp.ndarray

calculate_weights_hybrid(params, run_fingerprint, prices, start_index, additional_oracle_input=None, *args, **kwargs)[source]

Calculate the weights of assets in the pool using scan-based raw weight calculation.

This method produces the same outputs as calculate_weights, but uses calculate_rule_outputs_scan instead of calculate_rule_outputs.

Parameters:

params (Dict[str, Any]) – Pool parameters.
run_fingerprint (Dict[str, Any]) – Simulation settings.
prices (jnp.ndarray) – Current prices of the assets.
start_index (jnp.ndarray) – Start index for slicing
additional_oracle_input (Optional[jnp.ndarray], optional) – Additional input from an oracle. Defaults to None.

Returns:

Calculated weights for each asset in the pool.

Return type:

jnp.ndarray

calculate_readouts(params, run_fingerprint, prices, start_index, additional_oracle_input=None)[source]

Calculate readouts (internal estimator variables, other running variables) for the pool, based on price history.

This method can potentially have some overlap with calculate_rule_outputs, but for most TFMM pools it will simply correspond to the readout values for the gradient estimator (the ewma of prices and running a), sliced in the same way that the raw weight outputs are sliced.

Parameters:

params (Dict[str, Any]) – Pool parameters for weight calculation
run_fingerprint (Dict[str, Any]) – Simulation settings
prices (jnp.ndarray) – Historical price data
start_index (jnp.ndarray) – Start index for slicing
additional_oracle_input (Optional[jnp.ndarray]) – Extra data for weight calculation

Returns:

Dict containing readout values for the pool

Return type:

dict

abstract calculate_fine_weights(rule_output, initial_weights, run_fingerprint, params)[source]

Refine raw weight outputs into final weights.

Second step of TFMM’s weight calculation process. Converts raw weight adjustments into valid pool weights.

Parameters:

rule_output (jnp.ndarray) – Output from calculate_rule_outputs
initial_weights (jnp.ndarray) – Starting weights
run_fingerprint (Dict[str, Any]) – Simulation settings
params (Dict[str, Any]) – Pool parameters

Returns:

Final refined weights

Return type:

jnp.ndarray

calculate_weights(params, run_fingerprint, prices, start_index, additional_oracle_input=None, *args, **kwargs)[source]

Calculate the weights of assets in the pool.

Routes to either vectorized or scan-based weight calculation based on the weight_calculation_method in run_fingerprint: - “auto” (default): Use vectorized if available, else scan - “vectorized”: Force vectorized path (errors if not supported) - “scan”: Force scan path (errors if not supported)

Parameters:

(Dict[str (run_fingerprint)
Any]) (Simulation settings.)
(Dict[str
Any])
(jnp.ndarray) (start_index)
(jnp.ndarray)
(Optional[jnp.ndarray] (additional_oracle_input)
optional) (Additional input from an oracle. Defaults to None.)
params (Dict[str, Any])
run_fingerprint (Dict[str, Any])
prices (Array)
start_index (Array)
additional_oracle_input (Array | None)

Returns:

jnp.ndarray

Return type:

Calculated weights for each asset in the pool.

calculate_weights_vectorized(params, run_fingerprint, prices, start_index, additional_oracle_input=None, *args, **kwargs)[source]

Calculate weights using the vectorized path (calculate_rule_outputs).

Parameters:

(Dict[str (run_fingerprint)
Any]) (Simulation settings.)
(Dict[str
Any])
(jnp.ndarray) (start_index)
(jnp.ndarray)
(Optional[jnp.ndarray] (additional_oracle_input)
optional) (Additional input from an oracle. Defaults to None.)
params (Dict[str, Any])
run_fingerprint (Dict[str, Any])
prices (Array)
start_index (Array)
additional_oracle_input (Array | None)

Returns:

jnp.ndarray

Return type:

Calculated weights for each asset in the pool.

calculate_final_weights(params, run_fingerprint, prices, start_index, additional_oracle_input=None, *args, **kwargs)[source]

Calculate the weights of assets in the pool.

This method should be implemented by subclasses to define how weights are calculated based on current prices, pool parameters, and optional additional oracle input.

Parameters:

(Dict[str (run_fingerprint)
Any]) (Simulation settings.)
(Dict[str
Any])
(jnp.ndarray) (start_index)
(jnp.ndarray)
(Optional[jnp.ndarray] (additional_oracle_input)
optional) (Additional input from an oracle. Defaults to None.)
params (Dict[str, Any])
run_fingerprint (Dict[str, Any])
prices (Array)
start_index (Array)
additional_oracle_input (Array | None)

Returns:

jnp.ndarray

Return type:

Calculated weights for each asset in the pool.

calculate_all_signature_variations(params)[source]

Calculate all valid trading signature combinations.

Abstract method that subclasses may implement to define valid trading patterns. Can be used by reserve calculation methods to determine possible arbitrage opportunities.

Parameters:: params (Dict[str, Any]) – Pool parameters that may influence valid trade combinations
Returns:: Array of valid trading signature combinations
Return type:: jnp.ndarray
Raises:: NotImplementedError – Base class does not implement this method

make_vmap_in_axes(params, n_repeats_of_recurred=0)[source]

Configure JAX vectorization axes for pool parameters.

FMM pools handle subsidiary parameters differently for vectorization due to their potentially more complex parameter structure.

Parameters:

params (Dict[str, Any]) – Pool parameters to vectorize
n_repeats_of_recurred (int, optional) – Number of times to repeat recurrent parameters, by default 0

Returns:

vmap axes configuration with subsidiary_params handled separately

Return type:

Dict[str, Any]

is_trainable()[source]

Indicate if pool weights can be trained.

TFMM pools are trainable by default, as their weights change based on market conditions.

Returns:: Always True for TFMM pools as weights are trainable
Return type:: bool

classmethod process_parameters(update_rule_parameters, run_fingerprint)[source]

Process TFMM pool parameters from web interface input.

Handles common TFMM parameters and delegates pool-specific processing to subclasses. Supports both per-token and global parameters.

Parameters:

update_rule_parameters (Dict[str, Any]) – Raw parameters from web interface, each containing: - name: Parameter identifier - value: Parameter values per token
run_fingerprint (Dict[str, Any]) – Run fingerprint dictionary

Returns:

Processed parameters including: - logit_lamb: Memory parameter - k: Update rate parameter - Additional pool-specific parameters

Return type:

Dict[str, np.ndarray]

Notes

Handles parameter broadcasting for single values
Validates parameter dimensions
Processes memory_days and k_per_day specially
Allows subclasses to add specific parameters

calculate_weights_direct(params, prices, maximum_change=0.0003, minimum_weight=0.03, initial_weights=None, initial_running_a=None, initial_ewma=None, *args, **kwargs)[source]

Calculate the weights of assets in the pool, directly from the prices. This is used to quickly calculate the weights from any given price array, without doing any chunking or fine-weighting.

Parameters:

(Dict[str (params)
Any]) (Pool parameters.)
(jnp.ndarray) (prices)
(jnp.ndarray (initial_ewma)
optional) (Initial ewma value of the gradient estimator)
(jnp.ndarray
optional)
(jnp.ndarray
optional)
params (Dict[str, Any])
prices (Array)
maximum_change (float)
minimum_weight (float)
initial_weights (Array | None)
initial_running_a (Array | None)
initial_ewma (Array | None)

Returns:

jnp.ndarray

Return type:

Calculated weights for each asset in the pool.