add calendars, allow arbitrary complex positions

Author: Graeme22

README.md

@@ -1,11 +1,8 @@
 # margin-estimator
-Calculate estimated margin requirements for equities, options, futures, and futures options. Based on CBOE and CME margining.
+Calculate estimated margin requirements for equities options, based on CBOE margining.
 
 > [!NOTE]
-> Not all features are available yet, pending further development.
-> Currently, equity/ETF/index options are supported, for any trade
-> type other than ratio spreads, box spreads, and jaze lizards.
-> Contributions welcome!
+> Box spreads and some complex multi-legged positions may be calculated too conservatively.
 
 ## Installation
 
@@ -15,7 +12,7 @@ $ pip install margin_estimator
 
 ## Usage
 
-Simply pass a list of legs to the `calculate_margin` function along with an `Underlying` object containing information on the underlying, and you'll get back margin requirement estimates for cash and margin accounts!
+Simply pass an arbitrary list of legs to the `calculate_margin` function along with an `Underlying` object containing information on the underlying, and you'll get back margin requirement estimates for cash and margin accounts!
 
 ```python
 from datetime import date
@@ -116,4 +113,4 @@ print(margin)
 >>> cash_requirement=Decimal('7555.00') margin_requirement=Decimal('3661.00')
 ```
 
-Please note that all numbers are baseline estimates based on CBOE/CME guidelines and individual broker margins will likely vary significantly.
+Please note that all numbers are baseline minimums from CBOE guidelines and individual broker margins will likely vary significantly.

pyproject.toml

@@ -4,8 +4,8 @@ build-backend = "hatchling.build"
 
 [project]
 name = "margin-estimator"
-version = "0.1"
-description = "Calculate estimated margin requirements for equities, options, futures, and futures options. Based on CBOE and CME margining."
+version = "0.2"
+description = "Calculate estimated margin requirements for equities options, based on CBOE margining."
 readme = "README.md"
 requires-python = ">=3.11"
 license = {file = "LICENSE"}

src/margin_estimator/margin.py

@@ -1,3 +1,4 @@
+from collections import deque
 from datetime import date, timedelta
 from decimal import Decimal
 
@@ -13,31 +14,76 @@
 
 
 def calculate_margin(legs: list[Option], underlying: Underlying) -> MarginRequirements:
-    """
-    Calculate margin for an arbitrary order according to CBOE's Margin Manual.
-    """
-    if len(legs) == 1:
-        if legs[0].quantity > 0:
-            return _calculate_margin_long_option(legs[0])
-        return _calculate_margin_short_option(legs[0], underlying)
-    if len(legs) == 2 and legs[0].quantity < 0 and legs[1].quantity < 0:
-        return _calculate_margin_short_strangle(legs, underlying)
-    if len(legs) == 2 and legs[0].expiration != legs[1].expiration:
-        short = legs[0] if legs[0].quantity < 0 else legs[1]
-        long = legs[1] if legs[1] != short else legs[0]
-        if short.expiration > long.expiration:
-            margin = _calculate_margin_short_option(short, underlying)
-            return margin + _calculate_margin_long_option(long)
-    calls = [leg for leg in legs if leg.type == OptionType.CALL]
-    puts = [leg for leg in legs if leg.type == OptionType.PUT]
-    extra_puts = sum(c.quantity for c in calls)
-    extra_calls = sum(p.quantity for p in puts)
-    if extra_puts or extra_calls:
-        raise Exception(
-            "Ratio spreads/complex orders not supported! Try splitting your order into "
-            "smaller components and summing the results."
+    # sort by expiry to cover near-term risk first
+    shorts = sorted(leg for leg in legs if leg.quantity < 0)
+    longs = {
+        leg: leg.quantity for leg in sorted([leg for leg in legs if leg.quantity > 0])
+    }
+    covered: list[Option] = []
+    naked_shorts: list[Option] = []
+
+    # step 1: match spreads
+    for short in shorts:
+        unmatched = abs(short.quantity)
+        # iterate our long inventory to find a cover
+        for long, available in longs.items():
+            if unmatched <= 0:
+                break
+            if available <= 0:
+                continue
+            # constraint: same type, long expiry >= short expiry
+            if long.type == short.type and long.expiration >= short.expiration:
+                paired = min(unmatched, available)
+                covered.append(short.model_copy(update={"quantity": -paired}))
+                covered.append(long.model_copy(update={"quantity": paired}))
+                unmatched -= paired
+                longs[long] -= paired
+
+        # remaining short quantity is naked
+        if unmatched > 0:
+            naked_shorts.append(short.model_copy(update={"quantity": -unmatched}))
+
+    # step 2: match strangles
+    strangles: list[tuple[Option, Option]] = []
+    naked_calls = deque(leg for leg in naked_shorts if leg.type == OptionType.CALL)
+    naked_puts = deque(leg for leg in naked_shorts if leg.type == OptionType.PUT)
+
+    while naked_calls and naked_puts:
+        call = naked_calls.popleft()
+        put = naked_puts.popleft()
+        q = min(abs(call.quantity), abs(put.quantity))
+        strangles.append(
+            (
+                call.model_copy(update={"quantity": -q}),
+                put.model_copy(update={"quantity": -q}),
+            )
         )
-    return _calculate_margin_spread(legs)
+        # handle remaining quantity
+        if rem := abs(call.quantity) - q:
+            naked_calls.appendleft(call.model_copy(update={"quantity": -rem}))
+        if rem := abs(put.quantity) - q:
+            naked_puts.appendleft(put.model_copy(update={"quantity": -rem}))
+
+    # all unmatched options at this point go here
+    naked = list(naked_calls)
+    naked.extend(naked_puts)
+    naked.extend(
+        leg.model_copy(update={"quantity": q}) for leg, q in longs.items() if q
+    )
+
+    # step 3: calculate totals
+    total = MarginRequirements(cash_requirement=ZERO, margin_requirement=ZERO)
+    if covered:
+        total += _calculate_margin_spread(covered)
+    for call, put in strangles:
+        total += _calculate_margin_short_strangle([call, put], underlying)
+    for leg in naked:
+        if leg.quantity > 0:
+            total += _calculate_margin_long_option(leg)
+        else:
+            total += _calculate_margin_short_option(leg, underlying)
+
+    return total
 
 
 def _calculate_margin_long_option(option: Option) -> MarginRequirements:

src/margin_estimator/models.py

@@ -1,8 +1,9 @@
-from datetime import date
+import re
+from datetime import date, datetime
 from decimal import Decimal
 from enum import StrEnum
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 
 class ETFType(StrEnum):
@@ -17,12 +18,51 @@ class OptionType(StrEnum):
 
 
 class Option(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
     expiration: date
     price: Decimal
     quantity: int
     strike: Decimal
     type: OptionType
 
+    def __lt__(self, other: "Option"):
+        return (self.expiration, self.strike, self.type.value) < (
+            other.expiration,
+            other.strike,
+            other.type.value,
+        )
+
+    @classmethod
+    def from_occ(cls, symbol: str, price: Decimal, quantity: int) -> "Option":
+        match = re.match(r"(\d{6})([CP])(\d{5})(\d{3})", symbol[6:])
+        assert match
+        exp = datetime.strptime(match.group(1), "%y%m%d").date()
+        option_type = match.group(2)
+        strike = int(match.group(3)) + Decimal(match.group(4)) / 1000
+        return cls(
+            expiration=exp,
+            price=price,
+            quantity=quantity,
+            strike=strike,
+            type=OptionType(option_type),
+        )
+
+    @classmethod
+    def from_dxfeed(cls, symbol: str, price: Decimal, quantity: int) -> "Option":
+        match = re.match(r"\.([A-Z]+)(\d{6})([CP])(\d+(\.\d+)?)", symbol)
+        assert match
+        exp = datetime.strptime(match.group(2), "%y%m%d").date()
+        option_type = match.group(3)
+        strike = Decimal(match.group(4))
+        return cls(
+            expiration=exp,
+            price=price,
+            quantity=quantity,
+            strike=strike,
+            type=OptionType(option_type),
+        )
+
 
 class MarginRequirements(BaseModel):
     cash_requirement: Decimal
@@ -34,6 +74,12 @@ def __add__(self, other: "MarginRequirements"):
             margin_requirement=self.margin_requirement + other.margin_requirement,
         )
 
+    def __eq__(self, other):
+        return (
+            self.cash_requirement == other.cash_requirement
+            and self.margin_requirement == other.margin_requirement
+        )
+
 
 class Underlying(BaseModel):
     etf_type: ETFType | None = None