plotting node_wise_regression

[MeaghanPerdue]

Hello!
I'm trying to use AFQ-Insight for combat harmonization and regression analysis of my multi-site pyAFQ tract profiles dataset.
I'm running AFQ-Insight v0.7.1 in a conda env with python 3.12

Following the examples in the documentation, I'm able to load my pyAFQ tract profiles data and work with it to run OLS regression, etc.
I'm wondering if there is a way I can add a color key to indicate the groups for a regression as in the example below, which i updated to fit my own data based on the documentation example? I'd like to indicate the colors of the lines for males and females, but i'm not sure which is which in the plots by default.

# example here: <https://tractometry.org/AFQ-Insight/auto_examples/plot_als_comparison.html#sphx-glr-auto-examples-plot-als-comparison-py>

import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from afqinsight import AFQDataset
from afqinsight.parametric import node_wise_regression
from afqinsight.plot import plot_regression_profiles

# When reading in data, target_cols[0] gets renamed to "group", so the first listed variable must be the variable of interest for nodewise regression
afqdata = AFQDataset.from_files(
    fn_nodes="tract_profiles.csv",
    fn_subjects="subjects.csv",
    dwi_metrics=["dti_md", "dti_fa"],
    target_cols=["sex", "mri_age", "ga_days", "bw"]
    )

# Specify the tracts of interest
# ------------------------------
# Many times we have a specific set of tracts we want to analyze. We can specify
# those tracts using a list with each tract of interest

tracts = ["Left Arcuate", "Right Arcuate", "Left Corticospinal", "Right Corticospinal"]

# Set up plot, run linear models, and show the results
# ----------------------------------------------------
# With the next few lines of code, we fit a linear model in order to predict
# group differences in the diffusion properties of our tracts of interest. The
# function `node_wise_regression` does this by looping through each node of a
# tract and predicting our diffusion metric, in this case FA, as a function of
# group. The initializer `OLS.from_formula` takes `R-style formulas
# <https://www.statsmodels.org/dev/example_formulas.html>`_ as its model specification.
# Here, we are using the `"group"` variable as a categorical variable in the model.
# We can also specify linear-mixed effects models for more complex phenotypic data
# by passing in the appropriate model formula and setting `lme=True`.
#
# Because we conducted 100 comparisons, we need to correct the p-values that
# we obtained for the potential for a false discovery. There are multiple
# ways to conduct multiple comparison correction, and we will not get into
# the considerations in selecting a method here. The function `node_wise_regression`
# uses Benjamini/Hochberg FDR controlling method. This returns a boolean array for
# the p-values that are rejected at a specified alpha level (after correction),
# as well as an array of the corrected p-values.

num_cols = 2

# Define the figure and grid
fig, axes = plt.subplots(nrows=2, ncols=num_cols, figsize=(10, 6))


# Loop through the data and generate plots
for i, tract in enumerate(tracts):
    tract_dict = node_wise_regression(afqdata, tract, "dti_fa", "dti_fa ~ C(group)")

    row = i // num_cols
    col = i % num_cols
    axes[row][col].set_title(tract)

    plot_regression_profiles(tract_dict, axes[row][col])

plt.tight_layout()
plt.show()