Parkinson’s disease in the spinal cord: an exploratory study to establish T2*w, MTR and diffusion-weighted imaging metric values

import pandas as pd
import numpy as np
import plotly.graph_objects as go
from plotly.offline import plot
from plotly.subplots import make_subplots
import base64
import os
import plotly.io as pio
pio.renderers.default = "plotly_mimetype"

# List of metrics to include
DWI_metrics = ['FA', 'MD', 'AD', 'RD', 'ODI', 'FISO', 'FICVF', 'T2star', 'MTR']

# Labels and vertebral levels
labels = ['spinal cord', 'white matter', 'gray matter', 'WM/GM', 'dorsal columns', 'ventral funiculi', 'lateral funiculi']
spinal_levels = ['2', '3', '4', '5']

# Load all datasets and compute WM/GM
datasets = {}
for metric in DWI_metrics:
    df = pd.read_csv(f'../../data/parkinsons-spinalcord-mri-metrics/data/{metric}.csv')

    # Compute WM/GM
    data_WM = df[df['Label'] == 'white matter'].copy()
    data_GM = df[df['Label'] == 'gray matter'].copy()
    data_WMGM = data_WM.copy()
    data_WMGM['Label'] = 'WM/GM'
    data_WMGM['WA'] = data_WM['WA'].values / data_GM['WA'].values
    df = pd.concat([df, data_WMGM], ignore_index=True)

    datasets[metric] = df

# Create figure with subplots
fig = make_subplots(
    rows=5, cols=8,
    horizontal_spacing=0.035,
    vertical_spacing=0.07
)

# Store trace visibility info per metric
trace_visibility_by_metric = {metric: [] for metric in DWI_metrics}
all_traces = []

# Add scatter traces for each metric
for metric_index, metric in enumerate(DWI_metrics):
    df = datasets[metric]
    for i, spinal_level in enumerate(spinal_levels):
        for j, label in enumerate(labels):
            row = i + 2 # Start at row 2 to leave the first row empty for the template image
            col = j + 2 # Same for the columns

            filtered = df[(df['Label'] == label) & (df['SpinalLevel'] == spinal_level)]
            PD_data = filtered[filtered['CTRL_or_PD'].str.contains('PD', na=False)]

            # Add trace (only for subjects with Parkinson's disease (PD))
            trace = go.Scatter(
                x=PD_data['UPDRSIII_total'],
                y=PD_data['WA'],
                mode='markers',
                marker=dict(color='#B4464F', opacity=0.8, size=5),
                name='PD',
                showlegend=(row == 2 and col == 2),
                legendgroup='PD', # Link all traces with the same legend group (so they are toggled together)
                visible=(metric_index == 0)
            )
            fig.add_trace(trace, row=row, col=col)
            trace_visibility_by_metric[metric].append(len(all_traces))
            all_traces.append(trace)

            # Add regression line
            if len(PD_data) > 1:
                coef = np.polyfit(PD_data['UPDRSIII_total'], PD_data['WA'], 1)
                x_range = np.linspace(PD_data['UPDRSIII_total'].min(), PD_data['UPDRSIII_total'].max(), 100)
                y_fit = coef[0] * x_range + coef[1]
                trace_fit = go.Scatter(
                    x=x_range,
                    y=y_fit,
                    mode='lines',
                    line=dict(color='#B4464F', width=2),
                    name='PD',
                    showlegend=False,
                    legendgroup='PD',
                    visible=(metric_index == 0)
                )
                fig.add_trace(trace_fit, row=row, col=col)
                trace_visibility_by_metric[metric].append(len(all_traces))
                all_traces.append(trace_fit)

            # Axes titles and layout
            fig.update_xaxes(tickfont=dict(size=11), title=dict(text='UPDRSIII', font=dict(size=14, family='Arial')), row=i+2, col=j+2, title_standoff=0)
            fig.update_yaxes(tickfont=dict(size=11), row=i+2, col=j+2, title_standoff=0, tickformat=".3f")

            fig.update_yaxes(fixedrange=True)
            fig.update_xaxes(fixedrange=True)

# Add dropdown menu
dropdown_buttons = []
for metric in DWI_metrics:
    visibility = [False] * len(all_traces)
    for trace_index in trace_visibility_by_metric[metric]:
        visibility[trace_index] = True

    button = dict(
        label=metric,
        method="update",
        args=[{"visible": visibility}]
    )
    dropdown_buttons.append(button)

# Load static background image (from the "templates_for_figures" folder) and encode as base64
with open("../templates_for_figures/suppl_figures1and2_template.png", "rb") as image_file:
    encoded_image = base64.b64encode(image_file.read()).decode()

fig.add_layout_image(
    dict(
        source="data:image/png;base64," + encoded_image,
        xref="paper",
        yref="paper",
        x=-0.01,  # Aligns the image to the left edge of the figure
        y=1.02,  # Aligns the image to the top edge of the figure
        sizex=1.1,  # Adjust the size relative to the figure
        sizey=1.1,  # Adjust the size relative to the figure
        xanchor="left",  # Anchors the image position to the left
        yanchor="top",   # Anchors the image position to the top
        opacity=1,
        layer="below"
    )
)

# Update layout with dropdown
fig.update_layout(
    updatemenus=[dict(
        active=0,
        buttons=dropdown_buttons,
        x=0.01, xanchor='left',
        y=1.0, yanchor='top',
        font=dict(
            size=16,    
            color='black', 
            family='Arial' 
        )
    )],
    height=800, width=1520,
    legend=dict(
        x=0.06, # Horizontal position (0 = left, 1 = right)
        y=0.93, # Vertical position (0 = bottom, 1 = top)
        xanchor='right',  # Anchor relative to the x-position
        yanchor='top',    # Anchor relative to the y-position
        bgcolor='rgba(255,255,255, 0)',  # Transparent background
        itemsizing='constant',
        font=dict(size=16, family='Arial', color='white'),
        )
)

fig.show()