# Updating the AutoRA Workflow

Now that we configured all the pieces, let's integrate them into a single workflow. The AutoRA workflow will generate the experiment, upload the experiment conditions to the Firestore database, and download the results of the experiment. 

## Workflow Overview

You can learn more about AutoRA workflows in the [AutoRA Tutorial](../../tutorials/basic/Tutorial III Functional Workflow.ipynb). The figure below illustrates the workflow for our closed-loop example.

![Setup](img/workflow_overview.png)

The workflow involves the **following steps**:
1. Seed the experiment with two initial conditions.
2. Repeat the following steps three times:
   1. Collect data from an experiment runner that communicates with our web-based experiment. Here, we will use the ``experiment-runnner-firebase-prolific``.
   2. Identify models that characterize the data using an AutoRA theorist. Here, we will construct two models, a simple Logistic Regression from the ``sklearn`` package and the Bayesian Machine Scientist (BMS) from the ``autora-theorist-bms`` package.
   3. We will compare the two models to identify two new experiments that differentiate between these models. For this part, we will use the ``autora-experimentalist-model-disagreement`` package.

AutoRA passes around a ``state`` object that contains all relevant data for the workflow. This object is updated at each step of the workflow (see the Figure above). 


## Updating the Workflow

- Replace the code in the ``autora-workflow.py`` file with the following code:

```python
"""
Basic Workflow
    Two Independent Variables, One Dependent Variable
    Theorist: Logistic Regression, Bayesian Machine Scientist
    Experimentalist: Random Sampling, Model Disagreement
    Runner: Firebase Runner (no prolific recruitment)
"""

import json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from autora.variable import VariableCollection, Variable
from autora.theorist.bms import BMSRegressor
from autora.experimentalist.grid import grid_pool
from autora.experimentalist.random import random_sample
from autora.experimentalist.model_disagreement import model_disagreement_sample
from autora.experiment_runner.firebase_prolific import firebase_runner
from autora.state import StandardState, on_state, Delta

from sklearn.linear_model import LogisticRegression
from sklearn.base import BaseEstimator, ClassifierMixin

from trial_sequence import trial_sequence
from stimulus_sequence import stimulus_sequence
from preprocessing import trial_list_to_experiment_data

# Some study parameters
num_cycles = 3
num_trials = 20
num_conditions_per_cycle = 2

# *** Set up variables *** #
# The independent variables correspond to dot numbers for left and right stimulus.
# The dependent variable is accuracy ranging from 0 (0 percent) to 1 (100 percent)
variables = VariableCollection(
    independent_variables=[
        Variable(name="dots_left", allowed_values=np.linspace(1, 100, 100)),
        Variable(name="dots_right", allowed_values=np.linspace(1, 100, 100)),
        ],
    dependent_variables=[Variable(name="accuracy", value_range=(0, 1))])

# Here, we enlist the entire design space, by combining all possibles values 
# of each of the two independent variables
allowed_conditions = grid_pool(variables)
# We remove the conditions where the number of dots is the same on both sides
allowed_conditions = allowed_conditions[allowed_conditions['dots_left'] != allowed_conditions['dots_right']]

# *** State *** #
# With the variables, we can set up a state. The state object represents the state of our
# closed-loop experiment.
state = StandardState(
    variables=variables,
)

# *** Components/Agents *** #
# Components are functions that run on the state. The main components are:
# - theorist
# - experiment-runner
# - experimentalist
# Learn more about components here: 
# https://autoresearch.github.io/autora/tutorials/basic/Tutorial%20I%20Components/


# ** Theorist ** #
# Here we use two different theorists:
# - Logistic Regression for discovering a linear regression model
# - Bayesian Machine Scientist for discovering a (non-linear) model
# Later, we will compare the models generated by both theorists 
# For more information about AutoRA theorists, see 
# https://autoresearch.github.io/autora/theorist/

# First, we define a custom LogisticRegressor class that inherits from BaseEstimator and ClassifierMixin
# We override the predict method to return the probability of the '1' (accurate) class,
# to align its output with that of the BMSRegressor
class LogisticRegressor(BaseEstimator, ClassifierMixin):
    def __init__(self, *args, **kwargs):
        self.model = LogisticRegression(*args, **kwargs)  # Initialize the LogisticRegression model with any passed arguments

    def fit(self, X, y):
        """Fit the LogisticRegression model."""
        self.model.fit(X, y)
        return self

    def predict(self, X):
        """Override the predict method to return the probability of the '1' (accurate) class."""
        return self.model.predict_proba(X)[:, 1].reshape(-1, 1) 

# Initialize the BMSRegressor and LogisticRegressor
bms_theorist = BMSRegressor(epochs=500)
lr_theorist = LogisticRegressor()

# To use the theorists on the state object, we wrap it with the on_state functionality and return a
# Delta object, indicating how we want to modify the state.
# Note: The if the input arguments of the theorist_on_state function are state-fields like
# experiment_data, variables, ..., then calling this function on a state object will automatically
# update those state fields. You can learn more about states and the on_state function here:
# https://autoresearch.github.io/autora/tutorials/basic/Tutorial%20IV%20Customization/

@on_state()
def theorist_on_state(experiment_data, variables):
    ivs = [iv.name for iv in variables.independent_variables]
    dvs = [dv.name for dv in variables.dependent_variables]
    x = experiment_data[ivs]
    y = experiment_data[dvs]
    return Delta(models=[bms_theorist.fit(x, y), lr_theorist.fit(x, y)])


# ** Experimentalists ** #
# We will seed our study with randomly sampled experiment conditions, 
# using a random pool experimentalist. Here, we use a random pool and use the wrapper
# to that operates on the state object.
@on_state()
def initialize_state(allowed_conditions, num_samples):
    return Delta(conditions=random_sample(allowed_conditions, num_samples))

# After collecting our first data, we will use a model disagreement experimentalist
# to determine the next experiment conditions.
@on_state()
def experimentalist_on_state(allowed_conditions, models_to_compare, num_samples):
    return Delta(conditions=model_disagreement_sample(allowed_conditions, models_to_compare, num_samples))


# ** Experiment Runner ** #
# We will run our experiment on firebase and need credentials. You will find them here:
# (https://console.firebase.google.com/)
#   -> project -> project settings -> service accounts -> generate new private key

firebase_credentials = {
  "type": "service_account",
  "project_id": "closed-loop-study",
  "private_key_id": "YOURKEYID",
  "private_key": "-----BEGIN PRIVATE KEY-----\nYOURCREDENTIALS\n-----END PRIVATE KEY-----\n",
  "client_email": "firebase-adminsdk-y7hnh@closed-loop-study.iam.gserviceaccount.com",
  "client_id": "YOURLIENTID",
  "auth_uri": "https://accounts.google.com/o/oauth2/auth",
  "token_uri": "https://oauth2.googleapis.com/token",
  "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
  "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/firebase-adminsdk-y7hnh%40closed-loop-study.iam.gserviceaccount.com",
  "universe_domain": "googleapis.com"
}

# Simple experiment runner that runs the experiment on firebase
# The runner defines a timeout of 100 seconds, which means that a participant
# has 5 *minutes* to complete an experiment. Afterward, it will be freed for another participant.
# The sleep time is set to 3 *seconds*, which means that the runner will check every 5 seconds for data.
experiment_runner = firebase_runner(
    firebase_credentials=firebase_credentials,
    time_out=5,
    sleep_time=3)

# Again, we need to wrap the experiment runner to use it on the state.
# Specifically, the runner compiles the identified conditions (i.e., number of tested dots)
# into an experiment sequence, which is then used to generate the full JavaScript experiment
# to be sent to Firebase. Once the data are collected, the runner will then pre-process it 
# and store it in the experimental_data object
@on_state()
def runner_on_state(conditions):
    res = []
    for idx, c in conditions.iterrows():
        iv_1 = c['dots_left']
        iv_2 = c['dots_right']
        # get a timeline via sweetPea
        timeline = trial_sequence(iv_1, iv_2, num_trials)[0]
        print("Generated counterbalanced trial sequence.")
        # get js code via sweetBeaan
        js_code = stimulus_sequence(timeline, iv_1, iv_2)
        print("Compiled experiment.")
        res.append(js_code)
    
    # prepare conditions to send
    conditions_to_send = conditions.copy()
    conditions_to_send['experiment_code'] = res
    # upload and run the experiment:
    print("Uploading the experiment...")
    data_raw = experiment_runner(conditions_to_send)
    print("Collected experimental data.")

    # process the experiment data
    experiment_data = pd.DataFrame()
    for item in data_raw:
        _lst = json.loads(item)['trials']
        _df = trial_list_to_experiment_data(_lst)
        experiment_data = pd.concat([experiment_data, _df], axis=0)
    print("Preprocessed experimental data.")
    return Delta(experiment_data=experiment_data)


# *** AutoRA Workflow *** #
# Next, we specify our actual workflow using the components defined above.
# We begin with sampling two random initial conditions for the dots, resulting in two experiments.
# We then iterate three times through the following steps:
# 1) Collect data.
# 2) Fit models.
# 3) Identify novel experiment conditions.

state = initialize_state(state, allowed_conditions=allowed_conditions, num_samples=num_conditions_per_cycle)

# Now, we can run our components in a loop
for _ in range(num_cycles):
    state = runner_on_state(state)
    print("Finished data collection and preprocessing.")
    state = theorist_on_state(state)
    print("Fitted models.")
    models_to_compare = [state.models[-1], state.models[-2]]
    state = experimentalist_on_state(state, 
                                     allowed_conditions=allowed_conditions, 
                                     models_to_compare=models_to_compare, 
                                     num_samples=num_conditions_per_cycle)  
    print("Determined experiment conditions.")


# *** Plot *** #
# This plot visualizes the fit of the final logistic regression model and the model identified by the
# Bayesian Machine Scientist. 
ivs = [iv.name for iv in variables.independent_variables]
dvs = [dv.name for dv in variables.dependent_variables]
X = state.experiment_data[ivs]
y = state.experiment_data[dvs]

# Create a meshgrid for plotting the logistic regression decision boundary
iv1_range = variables.independent_variables[0].allowed_values
iv2_range = variables.independent_variables[1].allowed_values
iv1_grid, iv2_grid = np.meshgrid(iv1_range, iv2_range)
iv_grid = np.c_[iv1_grid.ravel(), iv2_grid.ravel()]

# Predict probabilities for the meshgrid (for the surface plot)
dv_pred_lr = state.models[-1].predict(iv_grid).reshape(iv1_grid.shape)  # Probability of the positive class
dv_pred_bms = state.models[-2].predict(iv_grid).reshape(iv1_grid.shape)  # Probability of the positive class

# Create the 3D plot
fig = plt.figure(figsize=(14, 6))

ax1 = fig.add_subplot(121, projection='3d')
# Plot the actual data points
ax1.scatter(X['dots_left'], X['dots_right'], y, color='red', label='Data Points')
# Plot the logistic regression surface
ax1.plot_surface(iv1_grid, iv2_grid, dv_pred_lr, cmap='viridis', alpha=0.6)
# Label the axes
ax1.set_xlabel('Number of Dots (left)')
ax1.set_ylabel('Number of Dots (right)')
ax1.set_zlabel('Accuracy')
ax1.set_zlim(0, 1)
ax1.set_title("Logistic Regression")

ax2 = fig.add_subplot(122, projection='3d')
# Plot the actual data points
ax2.scatter(X['dots_left'], X['dots_right'], y, color='red', label='Data Points')
# Plot the logistic regression surface
ax2.plot_surface(iv1_grid, iv2_grid, dv_pred_bms, cmap='viridis', alpha=0.6)
# Label the axes
ax2.set_xlabel('Number of Dots (left)')
ax2.set_ylabel('Number of Dots (right)')
ax2.set_zlabel('Accuracy')
ax2.set_zlim(0, 1)
ax2.set_title("BMS Equation: " + state.models[-2].repr())

# Show the plot
plt.savefig('model_comparison.png')
plt.show()
```

- Next, you should be able to execute the workflow in the ``researcher_hub``:

```bash
python autora_workflow.py
```

- Once the workflow uploads the experiment, you should be able to participate in the experiment, helping AutoRA to collect data.

!!! hint
    If you forgot the website address for your experiment, you can find it in the [Firebase Console](https://console.firebase.google.com/). One the console, select your project and then click on ``Hosting`` in the left navigation menu. ![hosting.png](../img/hosting.png) You should then see the link to your experiment in the section ``Domains''. It's the upper one, in this case ``https://closed-loop-study.web.app``. ![domain.png](../img/domain.png)


- Try to run the workflow for three cycles. Once completed, you should see the plot (also stored in the file ``model_comparison.png``) that compares the logistic regression model with the Bayesian Machine Scientist model.

!!! hint
    Note that you need to wait until the experiment is finished until you see a page with white background. If you end the experiment before, the ``firebase_runner`` will wait the minutes specified in ``time_out`` before it will be available for the next participant, i.e., run. If no more slots are currently available, you should see something like "We are sorry, there has been an unexpected technical issue. Thank you for your understanding." 

!!! hint
    You can check which experiments were successfully completed by looking into the Firestore database. In your project on the [Firebase Console](https://console.firebase.google.com/), simply navigate to ``FireStore Database``. The fields in ``autora`` > ``autora_out`` > ``observations`` list all the conditions. "null" means that no data has been collected for that condition yet.

- **Congratulations**, you just set up and ran a closed-loop behavioral study!

Below, we provide some more detailed explanations for the code above.

### Explanation

#### Study Parameters
```python
num_cycles = 3
num_trials = 20
num_conditions_per_cycle = 2
```

In the first part of the code, we set on some simulation parameters. Here, the workflow is set to run for three cycles, with 20 trials per experiment. In each cycle, we will generate two new experimental conditions, which, in this case also results in two experiments.

#### Variables
```python
variables = VariableCollection(
    independent_variables=[
        Variable(name="dots_left", allowed_values=np.linspace(1, 100, 100)),
        Variable(name="dots_right", allowed_values=np.linspace(1, 100, 100)),
        ],
    dependent_variables=[Variable(name="accuracy", value_range=(0, 1))])
```

Here, we define the variables for our study. We have two independent variables, the number of dots on the left and right side, and one dependent variable, the accuracy of the participant's response. The independent variables range from 1 to 100, and the dependent variable ranges from 0 to 1.

#### Allowed Conditions

```python
allowed_conditions = grid_pool(variables)
allowed_conditions = allowed_conditions[allowed_conditions['dots_left'] != allowed_conditions['dots_right']]
```

We generate all possible conditions for the study by combining all possible values of the independent variables. We then remove the conditions where the number of dots is the same on both sides. The latter is needed because our experimental design script will consider conditions one of the two values is set for both the left and right side of the dots.

#### State

```python
state = StandardState(
    variables=variables,
)
```

The state object represents the state of our closed-loop experiment. It contains all relevant data for the workflow (e.g., ``experiment_data`` or generated ``models``) and is updated at each step of the workflow.

#### Theorist

```python
class LogisticRegressor(BaseEstimator, ClassifierMixin):
    def __init__(self, *args, **kwargs):
        self.model = LogisticRegression(*args, **kwargs)  # Initialize the LogisticRegression model with any passed arguments

    def fit(self, X, y):
        """Fit the LogisticRegression model."""
        self.model.fit(X, y)
        return self

    def predict(self, X):
        """Override the predict method to return the probability of the '1' (accurate) class."""
        return self.model.predict_proba(X)[:, 1].reshape(-1, 1) 
```

Here, we define one of the two theorists, the Logistic Regression. Specifically, we define a custom class ``LogisticRegressor`` that inherits from ``BaseEstimator`` and ``ClassifierMixin``. We override the ``predict`` method to return the probability of the '1' (accurate) class, aligning its output with that of the BMSRegressor. This is necessary because the predictions of both theorists are compared later on, so they are required to have the same output format.

```python
bms_theorist = BMSRegressor(epochs=500)
lr_theorist = LogisticRegressor()
```

We then initialize the Bayesian Machine Scientist (BMS) and the Logistic Regression (LR) theorist. The BMSRegressor is a custom class from the AutoRA package, while the LogisticRegressor is the custom class we defined above. The BMSRegressor is initialized to run for 500 epochs.

```python
@on_state()
def theorist_on_state(experiment_data, variables):
    ivs = [iv.name for iv in variables.independent_variables]
    dvs = [dv.name for dv in variables.dependent_variables]
    x = experiment_data[ivs]
    y = experiment_data[dvs]
    return Delta(models=[bms_theorist.fit(x, y), lr_theorist.fit(x, y)])
```

We then define a function ``theorist_on_state`` that fits the models to the data. The function is wrapped with the ``on_state`` decorator, which allows the function to update the state object of our workflow. The function takes the experiment data and variables as input and returns a ``Delta`` object that updates the state with the two identified models.

#### Experimentalists

```python
@on_state()
def initialize_state(allowed_conditions, num_samples):
    return Delta(conditions=random_sample(allowed_conditions, num_samples))
```

Here, we define the function ``initialize_state`` that seeds our study with randomly sampled experiment conditions. The function is wrapped with the ``on_state`` decorator, allowing it to update the state object. The function takes the allowed conditions and the number of samples as input and returns a ``Delta`` object that updates the state with the randomly sampled conditions using the [random experimentalist](https://autoresearch.github.io/autora/core/docs/experimentalists/random/).

```python
@on_state()
def experimentalist_on_state(allowed_conditions, models_to_compare, num_samples):
    return Delta(conditions=model_disagreement_sample(allowed_conditions, models_to_compare, num_samples))
```

We then define the function ``experimentalist_on_state`` that determines the next experiment conditions for every cycle. The function is wrapped with the ``on_state`` decorator, allowing it to update the state object. The function takes the allowed conditions, the models to compare, and the number of samples as input and returns a ``Delta`` object that updates the state with the conditions determined by the [model disagreement experimentalist](https://autoresearch.github.io/autora/user-guide/experimentalists/model-disagreement/).

#### Firebase Credentials

```python
firebase_credentials = {
  "type": "service_account",
  "project_id": "closed-loop-study",
  "private_key_id": "YOURKEYID",
  "private_key": "-----BEGIN PRIVATE KEY-----\nYOURCREDENTIALS\n-----END PRIVATE KEY-----\n",
  "client_email": "firebase-adminsdk-y7hnh@closed-loop-study.iam.gserviceaccount.com",
  "client_id": "YOURLIENTID",
  "auth_uri": "https://accounts.google.com/o/oauth2/auth",
  "token_uri": "https://oauth2.googleapis.com/token",
  "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
  "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/firebase-adminsdk-y7hnh%40closed-loop-study.iam.gserviceaccount.com",
  "universe_domain": "googleapis.com"
}
```

Here, we define the Firebase credentials required to run the experiment on Firebase. You can find these credentials in the Firebase console under ``project -> project settings -> service accounts -> generate new private key``.

#### Experiment Runner

```python
experiment_runner = firebase_runner(
    firebase_credentials=firebase_credentials,
    time_out=5,
    sleep_time=3)
```

We then define the experiment runner that runs the experiment on Firebase. The runner is wrapped with the ``on_state`` decorator, allowing it to update the state object. The runner takes the Firebase credentials, the timeout, and the sleep time as input and returns a ``Delta`` object that updates the state with the experiment data. The ``time_out`` determines the amount of time **in minutes** a participant has available to complete the experiment, while the ``sleep_time`` determines how many **seconds** the runner waits run another check for experimental data.

```python
@on_state()
def runner_on_state(conditions):
    res = []
```

We then define the function ``runner_on_state`` that uploads the experiment to Firebase and collects the experimental data. The function is wrapped with the ``on_state`` decorator, allowing it to update the state object. The function takes the conditions as input and returns a ``Delta`` object that updates the state with the preprocessed experiment data.

```python
    for idx, c in conditions.iterrows():
        iv_1 = c['dots_left']
        iv_2 = c['dots_right']
        # get a timeline via sweetPea
        timeline = trial_sequence(iv_1, iv_2, num_trials)[0]
        print("Generated counterbalanced trial sequence.")
        # get js code via sweetBeaan
        js_code = stimulus_sequence(timeline, iv_1, iv_2)
        print("Compiled experiment.")
        res.append(js_code)
```

First, we iterate through all the sampled conditions (in this case, there are two). Each conditions consists of a pair of values for the independent variables ``dots_left`` and ``dots_right``. For each condition, we generate a counterbalanced trial sequence using the ``trial_sequence`` function that we programmed [in a prior step](experiment.md).
We then compile the experiment using the ``stimulus_sequence`` function. The compiled experiment is then appended to the list ``res``.

```python
    conditions_to_send = conditions.copy()
    conditions_to_send['experiment_code'] = res
    # upload and run the experiment:
    print("Uploading the experiment...")
    data_raw = experiment_runner(conditions_to_send)
    print("Collected experimental data.")
```

We then update the conditions with the compiled experiment code and send the conditions to the experiment runner. The runner uploads the experiment to Firebase and collects the experimental data. The raw data is then stored in the variable ``data_raw``.

```python
    # process the experiment data
    experiment_data = pd.DataFrame()
    for item in data_raw:
        _lst = json.loads(item)['trials']
        _df = trial_list_to_experiment_data(_lst)
        experiment_data = pd.concat([experiment_data, _df], axis=0)
    print("Preprocessed experimental data.")
    return Delta(experiment_data=experiment_data)
```

Finally, we preprocess the experimental data by converting the raw data into a pandas DataFrame. Here, we iterate through data from each conducted experiment. The preprocessed data is then stored in the variable ``experiment_data``.

#### AutoRA Workflow

```python
state = initialize_state(state, allowed_conditions=allowed_conditions, num_samples=num_conditions_per_cycle)
```

We first seed the study with randomly sampled conditions using the ``initialize_state`` function.

```python
for _ in range(num_cycles):
    state = runner_on_state(state)
    print("Finished data collection and preprocessing.")
    state = theorist_on_state(state)
    print("Fitted models.")
    models_to_compare = [state.models[-1], state.models[-2]]
    state = experimentalist_on_state(state, 
                                     allowed_conditions=allowed_conditions, 
                                     models_to_compare=models_to_compare, 
                                     num_samples=num_conditions_per_cycle)  
    print("Determined experiment conditions.")
```

We then run the workflow for the specified number of cycles. In each cycle, we collect data, fit the two models, and identify novel experiment conditions using the wrapper functions defined above. 

Note that we can retrieve the two models using ``state.models[-1]`` and ``state.models[-2]``. As determined in the ``theorist_on_state`` function, the BMS model is stored at the second-to-last index of the models list, while the Logistic Regression model is stored at thelast index.

#### Plot

The remainder of the script simply plots the variables. Note that we can extract the relevant data from the state: 
```python
ivs = [iv.name for iv in variables.independent_variables]
dvs = [dv.name for dv in variables.dependent_variables]
X = state.experiment_data[ivs]
y = state.experiment_data[dvs]
```

This includes the set of allowed values which is convenient for plotting the entire space of possible conditions. 

```python

# Create a meshgrid for plotting the logistic regression decision boundary
iv1_range = variables.independent_variables[0].allowed_values
iv2_range = variables.independent_variables[1].allowed_values
iv1_grid, iv2_grid = np.meshgrid(iv1_range, iv2_range)
iv_grid = np.c_[iv1_grid.ravel(), iv2_grid.ravel()]
```

Predictions can also be obtained directly from the models stored in the state:

```python
# Predict probabilities for the meshgrid (for the surface plot)
dv_pred_lr = state.models[-1].predict(iv_grid).reshape(iv1_grid.shape)  # Probability of the positive class
dv_pred_bms = state.models[-2].predict(iv_grid).reshape(iv1_grid.shape)  # Probability of the positive class
```

The rest of the code generates the 3D plot comparing the predictions of the two models from the last cycle.

```python
# Create the 3D plot
fig = plt.figure(figsize=(14, 6))

ax1 = fig.add_subplot(121, projection='3d')
# Plot the actual data points
ax1.scatter(X['dots_left'], X['dots_right'], y, color='red', label='Data Points')
# Plot the logistic regression surface
ax1.plot_surface(iv1_grid, iv2_grid, dv_pred_lr, cmap='viridis', alpha=0.6)
# Label the axes
ax1.set_xlabel('Number of Dots (left)')
ax1.set_ylabel('Number of Dots (right)')
ax1.set_zlabel('Accuracy')
ax1.set_zlim(0, 1)
ax1.set_title("Logistic Regression")

ax2 = fig.add_subplot(122, projection='3d')
# Plot the actual data points
ax2.scatter(X['dots_left'], X['dots_right'], y, color='red', label='Data Points')
# Plot the logistic regression surface
ax2.plot_surface(iv1_grid, iv2_grid, dv_pred_bms, cmap='viridis', alpha=0.6)
# Label the axes
ax2.set_xlabel('Number of Dots (left)')
ax2.set_ylabel('Number of Dots (right)')
ax2.set_zlabel('Accuracy')
ax2.set_zlim(0, 1)
ax2.set_title("BMS Equation: " + state.models[-2].repr())

# Show the plot
plt.savefig('model_comparison.png')
plt.show()
```

### Next Steps

[Next: Connect your closed-loop experiment with Prolific in order to recruit real participants.](prolific.md)
