Implementing a Custom QuantAMM Strategy

When designing a QuantAMM strategy, an update rule, you’re essentially creating a function that maps market observations to weight changes. Key considerations:

Weight Calculation Paths

There are two approaches to implementing weight calculation in a custom strategy:

  1. Vectorized Path - Compute all weight changes at once using convolution operations (faster, good for simulation)

  2. Scan Path - Compute weight changes sequentially, one step at a time (matches on-chain execution)

You can implement one or both paths. See Weight Calculation Paths for detailed information.

Vectorized Implementation

The vectorized approach computes all weight outputs at once. This is typically faster for simulation and training.

from jax import jit
from jax import numpy as jnp
from typing import Dict, Any, Optional
from functools import partial
from quantammsim.pools.G3M.quantamm.TFMM_base_pool import TFMMBasePool

class MyCustomRule(TFMMBasePool):
    @partial(jit, static_argnums=(2))
    def calculate_rule_outputs(
        self,
        params: Dict[str, Any],
        run_fingerprint: Dict[str, Any],
        prices: jnp.ndarray,
        additional_oracle_input: Optional[jnp.ndarray] = None,
    ) -> jnp.ndarray:
        # Your vectorized weight calculation logic here
        # Returns weight changes for all time steps at once
        ...

Scan-Based Implementation

The scan-based approach processes prices one at a time, maintaining state between steps. This mirrors how weights are computed on-chain and is useful for verifying production behavior.

class MyCustomRule(TFMMBasePool):
    def get_initial_rule_state(
        self,
        initial_price: jnp.ndarray,
        params: Dict[str, Any],
        run_fingerprint: Dict[str, Any],
    ) -> Dict[str, jnp.ndarray]:
        """Initialize the estimator carry state from the first price."""
        # Return initial state dictionary
        return {
            "ewma": initial_price,
            # ... other state variables
        }

    def calculate_rule_output_step(
        self,
        carry: Dict[str, jnp.ndarray],
        price: jnp.ndarray,
        params: Dict[str, Any],
        run_fingerprint: Dict[str, Any],
    ) -> tuple:
        """Single-step weight update.

        Returns:
            (new_carry, rule_output) tuple
        """
        # Update state and compute weight change for this step
        new_carry = {
            "ewma": ...,  # Updated state
        }
        rule_output = ...  # Weight change for this step
        return new_carry, rule_output

Implementing Both Paths

For maximum flexibility, implement both paths. They should produce numerically equivalent results:

class MyCustomRule(TFMMBasePool):
    # Vectorized path
    def calculate_rule_outputs(self, params, run_fingerprint, prices, ...):
        ...

    # Scan path
    def get_initial_rule_state(self, initial_price, params, run_fingerprint):
        ...

    def calculate_rule_output_step(self, carry, price, params, run_fingerprint):
        ...

# Check which paths are supported
pool = MyCustomRule()
print(pool.supports_vectorized_path())  # True
print(pool.supports_scan_path())        # True

Creating a Custom Rule

To create a custom update rule:

  1. Inherit from TFMMBasePool

  2. Implement either the vectorized path (calculate_rule_outputs) or the scan path (get_initial_rule_state + calculate_rule_output_step), or both

  3. (Optional) Provide the logic for any custom parameters in the pool’s helper function init_base_parameters()

  4. Register the pool with JAX as a pytree node using register_pytree_node() (see note in Implementing a new AMM)

Note that the simulator does not enforce causality, so be careful to make sure no look-ahead bias is introduced in the raw weight calculation. If you stick to using provided QuantAMM estimators, e.g. the gradient estimator calc_gradients(), then you can be confident that no look-ahead bias is introduced.

Using Custom Rules

To use your custom rule, add it to the function create_pool() giving it a string name, and then pass this string name to the rule key in the run_fingerprint dictionary to use it.

run_fingerprint = {
    'tokens': ['BTC', 'USDC'],
    'rule': 'custom',  # Must register your rule in the creator.py file
    'initial_pool_value': 1000000.0
}

Design Considerations

When designing update rules, consider:

  1. Responsiveness: How quickly should weights change?

  2. Stability: Avoid oscillations or extreme changes

  3. Robustness: Handle edge cases and unusual market conditions

  4. Computational Efficiency: Rules run frequently, keep them fast

  5. Memory Usage: Consider how much historical data you need

The base TFMM implementation handles many edge cases and constraints, allowing you to focus on the core strategy logic in your update rule. Note of course that if the strategy is novel, for QuantAMM V1 it will have to be implemented as a smart contract. Contact the team if you’d like to discuss.