Add butterfly, iron condor, and covered strategies

README.md

@@ -12,9 +12,26 @@ Use cases for Optopsy:
 * Discover performance statistics on **percentage change** for various options strategies on a given stock
 
 ## Supported Option Strategies
-* Calls/Puts
-* Straddles/Strangles
-* Vertical Call/Put Spreads
+
+### Single Leg
+* Long/Short Calls
+* Long/Short Puts
+
+### Two Leg
+* Long/Short Straddles
+* Long/Short Strangles
+* Vertical Call Spreads (Bull/Bear)
+* Vertical Put Spreads (Bull/Bear)
+* Covered Call (synthetic)
+* Protective Put (synthetic)
+
+### Three Leg (Butterflies)
+* Long/Short Call Butterfly
+* Long/Short Put Butterfly
+
+### Four Leg
+* Iron Condor / Reverse Iron Condor
+* Iron Butterfly / Reverse Iron Butterfly
 
 ## Documentation
 Please see the [wiki](https://github.com/michaelchu/optopsy/wiki) for API reference.
@@ -29,7 +46,7 @@ You will need Python 3.8 or newer and Pandas 2.0.0 or newer and Numpy 1.26.0 or
 
 ### Installation
 ```
-pip install optopsy==2.0.2
+pip install optopsy==2.0.3
 ```
 
 ### Example

optopsy/__init__.py

@@ -13,6 +13,19 @@
     short_call_spread,
     long_put_spread,
     short_put_spread,
+    # Butterfly strategies
+    long_call_butterfly,
+    short_call_butterfly,
+    long_put_butterfly,
+    short_put_butterfly,
+    # Iron condor and iron butterfly strategies
+    iron_condor,
+    reverse_iron_condor,
+    iron_butterfly,
+    reverse_iron_butterfly,
+    # Covered strategies
+    covered_call,
+    protective_put,
 )
 from .datafeeds import csv_data
 
@@ -30,5 +43,18 @@
     "short_call_spread",
     "long_put_spread",
     "short_put_spread",
+    # Butterfly strategies
+    "long_call_butterfly",
+    "short_call_butterfly",
+    "long_put_butterfly",
+    "short_put_butterfly",
+    # Iron condor and iron butterfly strategies
+    "iron_condor",
+    "reverse_iron_condor",
+    "iron_butterfly",
+    "reverse_iron_butterfly",
+    # Covered strategies
+    "covered_call",
+    "protective_put",
     "csv_data",
 ]

optopsy/core.py

@@ -1,7 +1,6 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple
 import pandas as pd
 import numpy as np
-from functools import reduce
 from .definitions import evaluated_cols
 from .checks import _run_checks
 
@@ -170,13 +169,25 @@ def _calculate_otm_pct(data: pd.DataFrame) -> pd.DataFrame:
     )
 
 
+def _get_leg_quantity(leg: Tuple) -> int:
+    """Get quantity for a leg, defaulting to 1 if not specified."""
+    return leg[2] if len(leg) > 2 else 1
+
+
 def _apply_ratios(data: pd.DataFrame, leg_def: List[Tuple]) -> pd.DataFrame:
-    """Apply position ratios (long/short multipliers) to entry and exit prices."""
+    """Apply position ratios (long/short multipliers) and quantities to entry and exit prices."""
     for idx in range(1, len(leg_def) + 1):
         entry_col = f"entry_leg{idx}"
         exit_col = f"exit_leg{idx}"
-        entry_kwargs = {entry_col: lambda r: r[entry_col] * leg_def[idx - 1][0].value}
-        exit_kwargs = {exit_col: lambda r: r[exit_col] * leg_def[idx - 1][0].value}
+        leg = leg_def[idx - 1]
+        multiplier = leg[0].value * _get_leg_quantity(leg)
+        # Use default arguments to capture values at each iteration (avoid late binding)
+        entry_kwargs = {
+            entry_col: lambda r, col=entry_col, m=multiplier: r[col] * m
+        }
+        exit_kwargs = {
+            exit_col: lambda r, col=exit_col, m=multiplier: r[col] * m
+        }
         data = data.assign(**entry_kwargs).assign(**exit_kwargs)
 
     return data
@@ -206,6 +217,16 @@ def _assign_profit(
     return data
 
 
+def _rename_leg_columns(
+    data: pd.DataFrame, leg_idx: int, join_on: List[str]
+) -> pd.DataFrame:
+    """Rename columns with leg suffix, excluding join columns."""
+    rename_map = {
+        col: f"{col}_leg{leg_idx}" for col in data.columns if col not in join_on
+    }
+    return data.rename(columns=rename_map)
+
+
 def _strategy_engine(
     data: pd.DataFrame,
     leg_def: List[Tuple],
@@ -237,19 +258,19 @@ def _rule_func(
     ) -> pd.DataFrame:
         return d if r is None else r(d, ld)
 
-    partials = [leg[1](data) for leg in leg_def]
+    # Pre-rename columns for each leg to avoid suffix issues with 3+ legs
+    partials = [
+        _rename_leg_columns(leg[1](data).copy(), idx, join_on or [])
+        for idx, leg in enumerate(leg_def, start=1)
+    ]
     suffixes = [f"_leg{idx}" for idx in range(1, len(leg_def) + 1)]
 
-    return (
-        reduce(
-            lambda left, right: pd.merge(
-                left, right, on=join_on, how="inner", suffixes=suffixes
-            ),
-            partials,
-        )
-        .pipe(_rule_func, rules, leg_def)
-        .pipe(_assign_profit, leg_def, suffixes)
-    )
+    # Merge all legs sequentially
+    result = partials[0]
+    for partial in partials[1:]:
+        result = pd.merge(result, partial, on=join_on, how="inner")
+
+    return result.pipe(_rule_func, rules, leg_def).pipe(_assign_profit, leg_def, suffixes)
 
 
 def _process_strategy(data: pd.DataFrame, **context: Any) -> pd.DataFrame:

optopsy/definitions.py

@@ -56,26 +56,24 @@
 
 triple_strike_internal_cols = [
     "underlying_symbol",
-    "underlying_price_entry",
+    "underlying_price_entry_leg1",
     "expiration",
     "dte_entry",
+    "dte_range",
     "option_type_leg1",
     "strike_leg1",
     "option_type_leg2",
     "strike_leg2",
     "option_type_leg3",
     "strike_leg3",
-    "entry",
-    "exit",
-    "long_profit",
-    "short_profit",
-    "long_pct_change",
-    "short_pct_change",
+    "total_entry_cost",
+    "total_exit_proceeds",
+    "pct_change",
 ]
 
 quadruple_strike_internal_cols = [
     "underlying_symbol",
-    "underlying_price_entry",
+    "underlying_price_entry_leg1",
     "expiration",
     "dte_entry",
     "dte_range",
@@ -87,12 +85,9 @@
     "strike_leg3",
     "option_type_leg4",
     "strike_leg4",
-    "entry",
-    "exit",
-    "long_profit",
-    "short_profit",
-    "long_pct_change",
-    "short_pct_change",
+    "total_entry_cost",
+    "total_exit_proceeds",
+    "pct_change",
 ]
 
 # base columns of dataframe after aggregation(minus the calculated columns)

optopsy/rules.py

@@ -26,3 +26,78 @@ def _rule_non_overlapping_strike(
     )
 
     return data.query(query)
+
+
+def _rule_butterfly_strikes(
+    data: pd.DataFrame, leg_def: List[Tuple]
+) -> pd.DataFrame:
+    """
+    Filter butterfly strategies to ensure proper strike ordering and equal width.
+
+    A butterfly requires:
+    - strike_leg1 < strike_leg2 < strike_leg3
+    - Equal wing widths: (strike_leg2 - strike_leg1) == (strike_leg3 - strike_leg2)
+
+    Args:
+        data: DataFrame containing butterfly strategy data
+        leg_def: List of tuples defining strategy legs
+
+    Returns:
+        Filtered DataFrame with valid butterfly strike configurations
+    """
+    if len(leg_def) != 3:
+        return data
+
+    return data.query(
+        "strike_leg1 < strike_leg2 < strike_leg3 & "
+        "(strike_leg2 - strike_leg1) == (strike_leg3 - strike_leg2)"
+    )
+
+
+def _rule_iron_condor_strikes(
+    data: pd.DataFrame, leg_def: List[Tuple]
+) -> pd.DataFrame:
+    """
+    Filter iron condor strategies to ensure proper strike ordering.
+
+    An iron condor requires 4 strikes in ascending order:
+    - strike_leg1 (long put) < strike_leg2 (short put) < strike_leg3 (short call) < strike_leg4 (long call)
+
+    Args:
+        data: DataFrame containing iron condor strategy data
+        leg_def: List of tuples defining strategy legs
+
+    Returns:
+        Filtered DataFrame with valid iron condor strike configurations
+    """
+    if len(leg_def) != 4:
+        return data
+
+    return data.query(
+        "strike_leg1 < strike_leg2 < strike_leg3 < strike_leg4"
+    )
+
+
+def _rule_iron_butterfly_strikes(
+    data: pd.DataFrame, leg_def: List[Tuple]
+) -> pd.DataFrame:
+    """
+    Filter iron butterfly strategies to ensure proper strike ordering.
+
+    An iron butterfly requires:
+    - strike_leg1 (long put) < strike_leg2 (short put) = strike_leg3 (short call) < strike_leg4 (long call)
+    - The short put and short call share the same strike (ATM)
+
+    Args:
+        data: DataFrame containing iron butterfly strategy data
+        leg_def: List of tuples defining strategy legs
+
+    Returns:
+        Filtered DataFrame with valid iron butterfly strike configurations
+    """
+    if len(leg_def) != 4:
+        return data
+
+    return data.query(
+        "strike_leg1 < strike_leg2 & strike_leg2 == strike_leg3 & strike_leg3 < strike_leg4"
+    )