Add backend support for visualizations

backend/app/experiments/dependencies.py

@@ -18,7 +18,7 @@
 from .schemas import (
     ArmPriors,
     ArmResponse,
-    ExperimentSample,
+    ExperimentResponse,
     ExperimentsEnum,
     NotificationsResponse,
     ObservationType,
@@ -30,7 +30,7 @@
 async def experiments_db_to_schema(
     experiments_db: list[ExperimentDB],
     asession: AsyncSession,
-) -> list[ExperimentSample]:
+) -> list[ExperimentResponse]:
     """
     Convert a list of ExperimentDB objects to a list of ExperimentResponse schemas.
     """
@@ -47,7 +47,7 @@ async def experiments_db_to_schema(
             )
         ]
         all_experiments.append(
-            ExperimentSample.model_validate(
+            ExperimentResponse.model_validate(
                 {
                     **exp_dict,
                     "notifications": [
@@ -269,7 +269,7 @@ async def update_arm_parameters(
     treatments: Union[list[float], None],
 ) -> None:
     """Update the arm parameters based on the reward type and outcome"""
-    experiment_data = ExperimentSample.model_validate(experiment.to_dict())
+    experiment_data = ExperimentResponse.model_validate(experiment.to_dict())
     if experiment_data.reward_type == RewardLikelihood.BERNOULLI:
         Outcome(rewards[0])  # Check if reward is 0 or 1
     params = update_arm(

backend/app/experiments/models.py

@@ -262,6 +262,14 @@ class DrawDB(Base):
     context_val: Mapped[Optional[list[float]]] = mapped_column(
         ARRAY(Float), nullable=True
     )
+    current_alpha: Mapped[Optional[float]] = mapped_column(Float, nullable=True)
+    current_beta: Mapped[Optional[float]] = mapped_column(Float, nullable=True)
+    current_mu: Mapped[Optional[list[float]]] = mapped_column(
+        ARRAY(Float), nullable=True
+    )
+    current_covariance: Mapped[Optional[list[float]]] = mapped_column(
+        ARRAY(Float), nullable=True
+    )
 
     # Relationships
     arm: Mapped[ArmDB] = relationship("ArmDB", back_populates="draws", lazy="joined")
@@ -708,6 +716,10 @@ async def save_observation_to_db(
     draw.observed_datetime_utc = datetime.now(timezone.utc)
     draw.observation_type = observation_type
     draw.reward = reward
+    draw.current_alpha = draw.arm.alpha
+    draw.current_beta = draw.arm.beta
+    draw.current_mu = draw.arm.mu
+    draw.current_covariance = draw.arm.covariance
 
     await asession.commit()
     await asession.refresh(draw)

backend/app/experiments/routers.py

@@ -42,24 +42,26 @@
     ContextType,
     DrawResponse,
     Experiment,
-    ExperimentSample,
+    ExperimentResponse,
     ExperimentsEnum,
     ObservationType,
     Outcome,
+    PlottingData,
 )
+from .visualization_utils import get_required_plotting_data
 
 router = APIRouter(prefix="/experiment", tags=["Experiments"])
 
 logger = setup_logger(__name__)
 
 
 # --- POST experiments routers ---
-@router.post("/", response_model=ExperimentSample)
+@router.post("/", response_model=ExperimentResponse)
 async def create_experiment(
     experiment: Experiment,
     user_db: Annotated[UserDB, Depends(require_admin_role)],
     asession: AsyncSession = Depends(get_async_session),
-) -> ExperimentSample:
+) -> ExperimentResponse:
     """
     Create a new experiment in the current user's workspace.
     """
@@ -87,15 +89,15 @@ async def create_experiment(
 
     experiment_dict = experiment_db.to_dict()
     experiment_dict["notifications"] = [n.to_dict() for n in notifications]
-    return ExperimentSample.model_validate(experiment_dict)
+    return ExperimentResponse.model_validate(experiment_dict)
 
 
 # -- GET experiment routers ---
-@router.get("/", response_model=list[ExperimentSample])
+@router.get("/", response_model=list[ExperimentResponse])
 async def get_all_experiments(
     user_db: Annotated[UserDB, Depends(get_verified_user)],
     asession: AsyncSession = Depends(get_async_session),
-) -> list[ExperimentSample]:
+) -> list[ExperimentResponse]:
     """
     Retrieve all experiments for the current user's workspace.
     """
@@ -119,12 +121,12 @@ async def get_all_experiments(
     return all_experiments
 
 
-@router.get("/type/{experiment_type}", response_model=list[ExperimentSample])
+@router.get("/type/{experiment_type}", response_model=list[ExperimentResponse])
 async def get_all_experiments_by_type(
     experiment_type: ExperimentsEnum,
     user_db: Annotated[UserDB, Depends(get_verified_user)],
     asession: AsyncSession = Depends(get_async_session),
-) -> list[ExperimentSample]:
+) -> list[ExperimentResponse]:
     """
     Retrieve all experiments for the current user's workspace.
     """
@@ -149,12 +151,12 @@ async def get_all_experiments_by_type(
     return all_experiments
 
 
-@router.get("/id/{experiment_id}", response_model=ExperimentSample)
+@router.get("/id/{experiment_id}", response_model=ExperimentResponse)
 async def get_experiment_by_id(
     experiment_id: int,
     user_db: Annotated[UserDB, Depends(get_verified_user)],
     asession: AsyncSession = Depends(get_async_session),
-) -> ExperimentSample:
+) -> ExperimentResponse:
     """
     Retrieve a specific experiment by ID for the current user's workspace.
     """
@@ -330,7 +332,7 @@ async def draw_experiment_arm(
         )
 
     # -- Perform the draw ---
-    experiment_data = ExperimentSample.model_validate(experiment.to_dict())
+    experiment_data = ExperimentResponse.model_validate(experiment.to_dict())
 
     # Validate contexts input
     if contexts:
@@ -476,7 +478,77 @@ async def get_rewards(
                 "arm": [arm for arm in experiment.arms if arm.arm_id == draw.arm_id][0],
                 "reward": draw.reward,
                 "context_val": draw.context_val,
+                "current_alpha": draw.current_alpha,
+                "current_beta": draw.current_beta,
+                "current_mu": draw.current_mu,
+                "current_covariance": draw.current_covariance,
             }
         )
         for draw in draws
     ]
+
+
+@router.get("/{experiment_id}/plotting", response_model=PlottingData)
+async def get_plotting_data(
+    experiment_id: int,
+    workspace_db: WorkspaceDB = Depends(authenticate_workspace_key),
+    asession: AsyncSession = Depends(get_async_session),
+) -> PlottingData:
+    """
+    Retrieve the data required for plotting.
+    """
+    experiment = await get_experiment_by_id_from_db(
+        workspace_id=workspace_db.workspace_id,
+        experiment_id=experiment_id,
+        asession=asession,
+    )
+
+    if not experiment:
+        raise HTTPException(
+            status_code=404, detail=f"Experiment with id {experiment_id} not found"
+        )
+    draws = await get_draws_by_experiment_id(
+        experiment_id=experiment_id, asession=asession
+    )
+
+    if not draws:
+        raise HTTPException(
+            status_code=404,
+            detail=f"No draws found for experiment with id {experiment_id}",
+        )
+
+    experiment_data = ExperimentResponse.model_validate(experiment.to_dict())
+    draw_data = [
+        DrawResponse.model_validate(
+            {
+                "draw_id": draw.draw_id,
+                "draw_datetime_utc": str(draw.draw_datetime_utc),
+                "observed_datetime_utc": str(draw.observed_datetime_utc),
+                "arm": [arm for arm in experiment.arms if arm.arm_id == draw.arm_id][0],
+                "reward": draw.reward,
+                "context_val": draw.context_val,
+                "current_alpha": draw.current_alpha,
+                "current_beta": draw.current_beta,
+                "current_mu": draw.current_mu,
+                "current_covariance": draw.current_covariance,
+            }
+        )
+        for draw in draws
+    ]
+    try:
+        plotting_data = get_required_plotting_data(
+            experiment=experiment_data, draws=draw_data
+        )
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error retrieving plotting data: {str(e)}",
+        ) from e
+
+    return PlottingData(
+        prior_samples=plotting_data["prior_samples"],
+        posterior_samples=plotting_data["posterior_samples"],
+        volumes=plotting_data["volumes"],
+        posterior_means=plotting_data["posterior_means"],
+        posterior_stds=plotting_data["posterior_stds"],
+    )

backend/app/experiments/sampling_utils.py

@@ -7,7 +7,7 @@
 from .schemas import (
     ArmPriors,
     ContextLinkFunctions,
-    ExperimentSample,
+    ExperimentResponse,
     ExperimentsEnum,
     Outcome,
     RewardLikelihood,
@@ -167,7 +167,7 @@ def objective(theta: np.ndarray) -> float:
 # ------------- Import functions ----------------
 # --- Choose arm function ---
 def choose_arm(
-    experiment: ExperimentSample, context: Optional[Union[list, np.ndarray, None]]
+    experiment: ExperimentResponse, context: Optional[Union[list, np.ndarray, None]]
 ) -> int:
     """
     Choose arm based on posterior using Thompson Sampling.
@@ -212,7 +212,7 @@ def choose_arm(
 
 # --- Update arm parameters ---
 def update_arm(
-    experiment: ExperimentSample,
+    experiment: ExperimentResponse,
     rewards: list[float],
     arm_to_update: Optional[int] = None,
     context: Optional[Union[list, np.ndarray, None]] = None,

backend/app/experiments/schemas.py

@@ -333,6 +333,22 @@ class DrawResponse(BaseModel):
         description="Context values associated with the draw",
         default=None,
     )
+    current_alpha: Optional[Union[float, None]] = Field(
+        description="Current alpha value of the arm",
+        default=None,
+    )
+    current_beta: Optional[Union[float, None]] = Field(
+        description="Current beta value of the arm",
+        default=None,
+    )
+    current_mu: Optional[List[Union[float, None]]] = Field(
+        description="Current mean value of the arm",
+        default=None,
+    )
+    current_covariance: Optional[List[List[Union[float, None]]]] = Field(
+        description="Current covariance matrix of the arm",
+        default=None,
+    )
     arm: ArmResponse
     client: Optional[Client] = None
 
@@ -516,33 +532,30 @@ def check_contexts(self) -> Self:
 
 class ExperimentResponse(ExperimentBase):
     """
-    Pydantic model for a response for experiment creation
+    Pydantic model for experiments for drawing and updating arms.
     """
 
     experiment_id: int
     n_trials: int
     last_trial_datetime_utc: Optional[str] = None
+    observation_type: ObservationType = ObservationType.USER
 
     arms: list[ArmResponse]
-    notifications: list[NotificationsResponse]
     contexts: Optional[list[ContextResponse]] = None
     clients: Optional[list[Client]] = None
 
     model_config = ConfigDict(from_attributes=True)
 
 
-class ExperimentSample(ExperimentBase):
+class PlottingData(BaseModel):
     """
-    Pydantic model for experiments for drawing and updating arms.
+    Pydantic model for the data required for plotting.
     """
 
-    experiment_id: int
-    n_trials: int
-    last_trial_datetime_utc: Optional[str] = None
-    observation_type: ObservationType = ObservationType.USER
-
-    arms: list[ArmResponse]
-    contexts: Optional[list[ContextResponse]] = None
-    clients: Optional[list[Client]] = None
+    posterior_means: list
+    posterior_stds: list
+    volumes: list
+    prior_samples: list
+    posterior_samples: list
 
     model_config = ConfigDict(from_attributes=True)