From 5954e5681b2a0cf5a732e047bae37be4895616d6 Mon Sep 17 00:00:00 2001
From: Tom Palmer <remlapmot@hotmail.com>
Date: Mon, 27 Apr 2026 13:46:15 +0100
Subject: [PATCH 1/3] Add expand_only option to SEQopts to return expanded
 dataset without running full analysis

---
 pySEQTarget/SEQopts.py    |  5 +++++
 pySEQTarget/SEQuential.py |  9 +++++++--
 tests/test_expansion.py   | 38 +++++++++++++++++++++++++++++++++++++-
 3 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/pySEQTarget/SEQopts.py b/pySEQTarget/SEQopts.py
index 6dea2ea..5cd32ae 100644
--- a/pySEQTarget/SEQopts.py
+++ b/pySEQTarget/SEQopts.py
@@ -36,6 +36,9 @@ class SEQopts:
     :type excused: bool
     :param excused_colnames: Column names (at the same length of treatment_level) specifying excused conditions, default ``[]``
     :type excused_colnames: List[str]
+    :param expand_only: If True, ``SEQuential.expand()`` returns the expanded dataset and skips weighting,
+        modelling, and survival steps
+    :type expand_only: bool
     :param followup_class: Boolean to force followup values to be treated as classes
     :type followup_class: bool
     :param followup_include: Boolean to force regular followup values into model covariates
@@ -121,6 +124,7 @@ class SEQopts:
     denominator: Optional[str] = None
     excused: bool = False
     excused_colnames: List[str] = field(default_factory=lambda: [])
+    expand_only: bool = False
     followup_class: bool = False
     followup_include: bool = True
     followup_max: int = None
@@ -161,6 +165,7 @@ class SEQopts:
     def _validate_bools(self):
         bools = [
             "excused",
+            "expand_only",
             "followup_class",
             "followup_include",
             "followup_spline",
diff --git a/pySEQTarget/SEQuential.py b/pySEQTarget/SEQuential.py
index a439ac7..7b32a31 100644
--- a/pySEQTarget/SEQuential.py
+++ b/pySEQTarget/SEQuential.py
@@ -107,9 +107,11 @@ def __init__(
         _param_checker(self)
         _data_checker(self)
 
-    def expand(self) -> None:
+    def expand(self):
         """
-        Creates the sequentially nested, emulated target trial structure
+        Creates the sequentially nested, emulated target trial structure.
+        If ``expand_only`` is set in parameters, returns the expanded dataset as a
+        :class:`polars.DataFrame` and skips all subsequent analysis steps.
         """
         start = time.perf_counter()
         kept = [
@@ -169,6 +171,9 @@ def expand(self) -> None:
         end = time.perf_counter()
         self._expansion_time = _format_time(start, end)
 
+        if self.expand_only:
+            return self.DT
+
     def bootstrap(self, **kwargs) -> None:
         """
         Internally sets up bootstrapping - creating a list of IDs to use per iteration
diff --git a/tests/test_expansion.py b/tests/test_expansion.py
index e9e81cb..eac5b39 100644
--- a/tests/test_expansion.py
+++ b/tests/test_expansion.py
@@ -1,6 +1,7 @@
 import polars as pl
+from polars.testing import assert_frame_equal
 
-from pySEQTarget import SEQuential
+from pySEQTarget import SEQopts, SEQuential
 
 
 def _make_model(data):
@@ -90,3 +91,38 @@ def test_expansion_truncates_each_trial_independently():
 
     # Trial 1 starts at time=1, outcome at time=3 → followup 0,1,2
     assert sorted(trial_1["followup"].to_list()) == [0, 1, 2]
+
+
+def test_expand_only_returns_expanded_dataframe():
+    """expand_only=True should return the expanded DataFrame directly and the
+    return value should equal self.DT from a standard expand() call."""
+    data = pl.DataFrame(
+        {
+            "ID": [1, 1, 1, 1, 1],
+            "time": [0, 1, 2, 3, 4],
+            "eligible": [1, 0, 0, 0, 0],
+            "treatment": [0, 1, 0, 1, 0],
+            "outcome": [0, 0, 0, 0, 0],
+        }
+    )
+
+    model_only = SEQuential(
+        data,
+        id_col="ID",
+        time_col="time",
+        eligible_col="eligible",
+        treatment_col="treatment",
+        outcome_col="outcome",
+        time_varying_cols=[],
+        fixed_cols=[],
+        parameters=SEQopts(expand_only=True),
+    )
+    result = model_only.expand()
+
+    assert isinstance(result, pl.DataFrame)
+    assert_frame_equal(result, model_only.DT)
+
+    model_full = _make_model(data)
+    model_full.expand()
+
+    assert_frame_equal(result, model_full.DT)

From 02fc42b28fa1821ad438dd3ffdfe0afe8e4338db Mon Sep 17 00:00:00 2001
From: Tom Palmer <remlapmot@hotmail.com>
Date: Mon, 27 Apr 2026 13:46:23 +0100
Subject: [PATCH 2/3] Bump version

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 3482867..de87979 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pySEQTarget"
-version = "0.13.1"
+version = "0.13.2"
 description = "Sequentially Nested Target Trial Emulation"
 readme = "README.md"
 license = {text = "MIT"}

From 2882d07baa72ea8190fe0014377754ad2cea9582 Mon Sep 17 00:00:00 2001
From: Tom Palmer <remlapmot@hotmail.com>
Date: Tue, 28 Apr 2026 10:11:34 +0100
Subject: [PATCH 3/3] Skip protocol censoring in expand_only mode to preserve
 all rows in the expanded dataset

---
 pySEQTarget/SEQuential.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pySEQTarget/SEQuential.py b/pySEQTarget/SEQuential.py
index 7b32a31..76b9df9 100644
--- a/pySEQTarget/SEQuential.py
+++ b/pySEQTarget/SEQuential.py
@@ -162,7 +162,7 @@ def expand(self):
             pl.col(self.id_col).cast(pl.Utf8).alias(self.id_col)
         )
 
-        if self.method != "ITT":
+        if self.method == "dose-response" or (self.method == "censoring" and not self.expand_only):
             _dynamic(self)
         if self.selection_random:
             _random_selection(self)