Skip to content

Treatment

A Treatment object describes the causal treatment variable and its assignment mechanism for an experiment, quasi-experiment, or observational study.

In FeatureByte, treatments are essential for building causal models and conducting uplift analysis. They provide structured metadata that downstream causal estimators can use to select appropriate identification and modeling strategies.

Treatment objects are used in collaboration with Context and Use Case to define causal inference scenarios.

Current Support

At this stage, FeatureByte supports causal modeling for randomized binary treatments only. Support for multi-arm and continuous treatments, as well as observational studies, will be added in future releases.

Creating a Treatment

A Treatment is created using the create() method. The treatment specification includes several components that describe the experimental or observational design:

Basic Treatment Components

When creating a treatment, you must specify:

  1. Name: A descriptive name for the treatment
  2. Data Type: The database variable type (e.g., INT, VARCHAR)
  3. Treatment Type: The scale of the treatment variable
  4. Assignment Source: How units are assigned to treatment
  5. Assignment Design: Specific design within the chosen source

Treatment Types

The treatment_type parameter defines the scale of the treatment variable:

  • Binary: Two-level treatment (e.g., exposed vs control, coupon vs no coupon)
  • Multi-arm: Discrete treatment with more than two levels (e.g., dosage tiers, variants)
  • Continuous: Numeric treatment representing a dose, price, spend, or intensity. Suitable for dose-response estimators such as DR-learner and continuous treatment DML or orthogonal ML.

Assignment Source

The source parameter determines the high-level source of treatment assignment, which influences the identification strategy and whether uplift modeling is valid:

  • Randomized: Treatment is randomly assigned in a controlled experiment
  • Observational: Treatment assignment is determined by business rules or natural variation

Assignment Design

The design parameter specifies the specific assignment mechanism within the chosen source:

For randomized sources: - simple-randomization: Units randomly assigned to treatment/control - stratified-randomization: Randomization within strata - cluster-randomization: Clusters of units assigned together

For observational sources: - business-rule: Treatment assigned by business logic - self-selection: Units choose their own treatment - natural-experiment: Treatment determined by external factors

Treatment Specification Details

Time and Time Structure

The time parameter indicates when treatment is defined:

  • Static: Treatment assignment is fixed at a single point in time
  • Time-varying: Treatment can change over time

The time_structure parameter provides detailed temporal information:

  • None: No specific temporal structure
  • Instantaneous: Treatment effect is immediate
  • Delayed: Treatment effect occurs after a delay
  • Persistent: Treatment effect continues over time

Interference

The interference parameter describes whether units can affect each other (violations of SUTVA - Stable Unit Treatment Value Assumption):

  • None: No interference between units
  • Spatial: Units interfere based on geographical proximity
  • Network: Units interfere based on network connections
  • Temporal: Earlier treatments affect later outcomes

Treatment Labels and Control

For binary and multi-arm treatments, you must specify:

  • treatment_labels: List of raw treatment values used in the dataset (e.g., [0, 1] or ["A", "B", "C"])
  • control_label: The value representing the control or baseline arm (must be one of the treatment_labels)

For continuous treatments, these parameters must be None.

Propensity

The optional propensity parameter describes how treatment assignment probabilities are known or estimated:

propensity = fb.Propensity(
    granularity="global",  # or "unit" for individual-level
    knowledge="design-known",  # or "estimated"
    p_global=0.5,  # optional: known global probability
)

Examples

Example 1: Simple A/B Test

A basic randomized experiment with 50/50 assignment:

treatment = fb.Treatment.create(
    name="Churn Campaign A/B test",
    dtype=fb.DBVarType.INT,
    treatment_type=fb.TreatmentType.BINARY,
    source="randomized",
    design="simple-randomization",
    time="static",
    time_structure="instantaneous",
    interference="none",
    treatment_labels=[0, 1],
    control_label=0,
    propensity=fb.Propensity(
        granularity="global",
        knowledge="design-known",
        p_global=0.5,
    ),
)

Example 2: Multi-Arm Randomized Experiment

Testing multiple treatment variants:

multi_arm_treatment = fb.Treatment.create(
    name="Discount Level Experiment",
    dtype=fb.DBVarType.VARCHAR,
    treatment_type=fb.TreatmentType.MULTI_ARM,
    source="randomized",
    design="simple-randomization",
    time="static",
    time_structure="instantaneous",
    interference="none",
    treatment_labels=["no_discount", "10pct_discount", "20pct_discount", "30pct_discount"],
    control_label="no_discount",
    propensity=fb.Propensity(
        granularity="global",
        knowledge="design-known",
    ),
)

Example 3: Observational Treatment

A treatment assigned by business rules rather than randomization:

observational_treatment = fb.Treatment.create(
    name="Premium Feature Access",
    dtype=fb.DBVarType.INT,
    treatment_type=fb.TreatmentType.BINARY,
    source="observational",
    design="business-rule",
    time="static",
    time_structure="none",
    interference="none",
    treatment_labels=[0, 1],
    control_label=0,
    propensity=fb.Propensity(
        granularity="unit",
        knowledge="estimated",  # Requires propensity score modeling
    ),
)

Example 4: Continuous Treatment

A numeric treatment representing dose or intensity:

continuous_treatment = fb.Treatment.create(
    name="Marketing Spend Amount",
    dtype=fb.DBVarType.FLOAT,
    treatment_type=fb.TreatmentType.CONTINUOUS,
    source="observational",
    design="business-rule",
    time="static",
    time_structure="none",
    interference="none",
    treatment_labels=None,  # Not applicable for continuous
    control_label=None,  # Not applicable for continuous
    propensity=None,  # Not applicable for continuous
)

Working with Treatments

Getting Treatment Information

You can retrieve detailed information about a treatment using the info() method:

treatment = catalog.get_treatment("Churn Campaign A/B test")  # or fb.Treatment.get("...")
treatment_info = treatment.info()

The info dictionary contains:

  • author: The user who created the treatment
  • name: Name of the treatment
  • created_at: Creation timestamp
  • updated_at: Last update timestamp
  • description: Treatment description
  • dtype: Data type
  • treatment_type: Scale of the treatment variable
  • source: High-level source of treatment assignment
  • design: Specific assignment design
  • time: Temporal definition of treatment
  • time_structure: Detailed temporal structure
  • interference: Structure of interference between units
  • treatment_labels: List of treatment values
  • control_label: Control or baseline arm value
  • propensity: Treatment assignment probability specification

Updating Treatment Description

After creating a treatment, you can add or update its description using the update_description() method:

treatment.update_description("A/B test of churn prevention campaign targeting high-risk customers")

Listing Treatments

You can view all treatments in the catalog using the list() method:

treatments_df = fb.Treatment.list()

This returns a DataFrame with treatment names, types, and creation dates.

Retrieving a Treatment

To retrieve an existing treatment from the catalog, use the get() method:

treatment = fb.Treatment.get("Churn Campaign A/B test")

Deleting a Treatment

A treatment can be deleted from the persistent data store using the delete() method, only if it is not used in any Context:

treatment.delete()

Using Treatments in Causal Analysis

Treatment objects are used in combination with Context and Use Case to define causal inference scenarios:

  1. Create a Treatment: Define the treatment variable and its assignment mechanism
  2. Create a Context: Define the circumstances where features are served by specifying a primary entity and optional description. For causal modeling, associate the context with the treatment to define the experimental or observational setting.
  3. Create a Use Case: Combine the context with a Target to build and evaluate causal models

Example: Complete Causal Analysis Workflow

Here's a complete example showing how to use a treatment in causal analysis:

# Step 1: Create a treatment
treatment = fb.Treatment.create(
    name="Churn Campaign A/B test",
    dtype=fb.DBVarType.INT,
    treatment_type=fb.TreatmentType.BINARY,
    source="randomized",
    design="simple-randomization",
    time="static",
    time_structure="instantaneous",
    interference="none",
    treatment_labels=[0, 1],
    control_label=0,
    propensity=fb.Propensity(
        granularity="global",
        knowledge="design-known",
        p_global=0.5,
    ),
)

# Step 2: Create a context with the treatment
context = fb.Context.create(
    name="Active Customers in Churn Campaign",
    primary_entity=["customer"],
    description="Active customers who were part of the churn prevention A/B test",
    treatment_name="Churn Campaign A/B test"
)

# Step 3: Create a target for the outcome
invoice_view = catalog.get_view("GROCERYINVOICE")
target = invoice_view["Amount"].forward_aggregate(
    method=fb.AggFunc.SUM,
    target_name="Revenue_next_30d",
    window='30d',
    target_type=fb.TargetType.REGRESSION,
)
target.save()

# Step 4: Create a use case combining context and target
use_case = fb.UseCase.create(
    name="Churn Campaign Revenue Impact",
    context_name="Active Customers in Churn Campaign",
    target_name="Revenue_next_30d",
    description="Measure the causal impact of the churn campaign on customer revenue"
)

This structured approach ensures that:

  • The correct causal identification strategy is applied
  • Appropriate statistical methods are selected
  • Model evaluation metrics align with the experimental design
  • Results are interpretable in the context of the assignment mechanism

Treatment Properties

All treatment properties can be accessed as attributes:

treatment = fb.Treatment.get("Churn Campaign A/B test")

# Access treatment properties
print(treatment.dtype)              # DBVarType
print(treatment.treatment_type)     # TreatmentType
print(treatment.source)             # AssignmentSource
print(treatment.design)             # AssignmentDesign
print(treatment.time)               # TreatmentTime
print(treatment.time_structure)     # TreatmentTimeStructure
print(treatment.interference)       # TreatmentInterference
print(treatment.treatment_labels)   # List of labels
print(treatment.control_label)      # Control value
print(treatment.propensity)         # Propensity specification

See Also