main_plot.py

# Plot stability and feasibility scatter plots for community 1 and community 2,
# and stability and feasibility heatmaps for community 3.
# Modification: all stability (ev) plots use absolute values.

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import rcParams

rcParams['mathtext.fontset'] = 'stix'
rcParams['font.family'] = 'serif'
rcParams['font.serif'] = ['Times']
rcParams['font.size'] = 24

file_path = 'scan_balances_results.csv'
df = pd.read_csv(file_path)  # CSV is being read here

mask_bal1 = (~df['competition_cooperation_balance1'].isin([0, 1])) | df['competition_cooperation_balance1'].isna()
mask_bal2 = (~df['competition_cooperation_balance2'].isin([0, 1])) | df['competition_cooperation_balance2'].isna()
df = df[mask_bal1 & mask_bal2].copy()

POINT_COLOR = 'C0'
POINT_SIZE = 60
POINT_ALPHA = 0.4


def scatter_all_replicates_by_x(ax, x_vals_sorted, df_sub, x_col, y_col):
    for x in x_vals_sorted:
        ys = df_sub.loc[df_sub[x_col] == x, y_col].dropna().values
        if ys.size == 0:
            continue
        if y_col == 'ev':
            ys = np.abs(ys)

        ax.scatter(
            np.full(ys.shape, x, dtype=float),
            ys,
            s=POINT_SIZE,
            alpha=POINT_ALPHA,
            color=POINT_COLOR,
            edgecolors='none',
            linewidths=0
        )
    ax.grid(True, linestyle='--', alpha=0.3)


# ========= community1: two panels (stability / feasibility), plot all replicates for each x =========
df1 = df[df['community'] == 'community1'].copy()
x1_vals = np.sort(df1['competition_cooperation_balance1'].dropna().unique())

fig, axes = plt.subplots(1, 2, figsize=(11, 6), sharex=False)

# Left: stability (absolute ev)
scatter_all_replicates_by_x(
    axes[0], x1_vals, df1,
    x_col='competition_cooperation_balance1', y_col='ev'
)
axes[0].set_xlabel('Balance of community 1 ($b_1$)')
axes[0].set_ylabel('Stability')
axes[0].set_title(' ')

# Right: feasibility
scatter_all_replicates_by_x(
    axes[1], x1_vals, df1,
    x_col='competition_cooperation_balance1', y_col='feasibility'
)
axes[1].set_xlabel('Balance of community 1 ($b_1$)')
axes[1].set_ylabel('Feasibility')
axes[1].set_title(' ')

axes[0].set_xlim(0, 1)
axes[0].set_xticks(np.arange(0, 1.01, 0.2))
axes[1].set_xlim(0, 1)
axes[1].set_xticks(np.arange(0, 1.01, 0.2))

plt.tight_layout()
fig.savefig("community1.pdf", format="pdf", bbox_inches="tight", transparent=True)
plt.show()


# ========= community2: two panels (stability / feasibility), plot all replicates for each x =========
df2 = df[df['community'] == 'community2'].copy()
x2_vals = np.sort(df2['competition_cooperation_balance2'].dropna().unique())

fig, axes = plt.subplots(1, 2, figsize=(11, 6), sharex=False)

# Left: stability (absolute ev)
scatter_all_replicates_by_x(
    axes[0], x2_vals, df2,
    x_col='competition_cooperation_balance2', y_col='ev'
)
axes[0].set_xlabel('Balance of community 2 ($b_2$)')
axes[0].set_ylabel('Stability')
axes[0].set_title(' ')

# Right: feasibility
scatter_all_replicates_by_x(
    axes[1], x2_vals, df2,
    x_col='competition_cooperation_balance2', y_col='feasibility'
)
axes[1].set_xlabel('Balance of community 2 ($b_2$)')
axes[1].set_ylabel('Feasibility')
axes[1].set_title(' ')

axes[0].set_xlim(0, 1)
axes[0].set_xticks(np.arange(0, 1.01, 0.2))
axes[1].set_xlim(0, 1)
axes[1].set_xticks(np.arange(0, 1.01, 0.2))

plt.tight_layout()
fig.savefig("community2.pdf", format="pdf", bbox_inches="tight", transparent=True)
plt.show()


# ========= community3: 2D heatmaps (mean over replicates) =========
df3 = df[df['community'] == 'community3'].copy()

# Compute mean values for each (b1, b2) pair
agg = (
    df3.dropna(subset=['competition_cooperation_balance1', 'competition_cooperation_balance2'])
       .groupby(['competition_cooperation_balance1', 'competition_cooperation_balance2'])
       .agg(
           ev_mean=('ev', lambda x: np.mean(np.abs(x))),
           feas_mean=('feasibility', 'mean')
       )
       .reset_index()
)


def edges_from_centers(vals):
    vals = np.asarray(vals, dtype=float)
    if vals.size < 2:
        d = 0.025
        return np.array([vals[0] - d, vals[0] + d])
    d = np.diff(vals)
    left = vals[0] - d[0] / 2
    right = vals[-1] + d[-1] / 2
    return np.r_[left, vals[:-1] + d / 2, right]


def plot_surface_from_agg(agg_df, z_col, title):
    pivot = agg_df.pivot_table(
        index='competition_cooperation_balance2',
        columns='competition_cooperation_balance1',
        values=z_col
    )
    Xc = pivot.columns.values
    Yc = pivot.index.values
    Z = pivot.values.astype(float)

    if z_col == 'ev_mean':
        Z = np.abs(Z)

    Zm = np.ma.masked_invalid(Z)
    Xe = edges_from_centers(Xc)
    Ye = edges_from_centers(Yc)

    fig, ax = plt.subplots(figsize=(8, 6.5))

    # Use the default matplotlib colormap for both feasibility and stability
    im = ax.pcolormesh(Xe, Ye, Zm, shading='auto')

    cbar = fig.colorbar(im, ax=ax, shrink=0.85)
    if z_col == 'feas_mean':
        cbar.set_label('Feasibility')
    elif z_col == 'ev_mean':
        cbar.set_label('Stability')
    else:
        cbar.set_label(z_col)

    cbar.ax.tick_params(labelsize=14)

    ax.set_xlabel('Balance of community 1 ($b_1$)')
    ax.set_ylabel('Balance of community 2 ($b_2$)')
    ax.set_title(title)
    ax.grid(True, linestyle='--', alpha=0.25)
    ax.set_xlim(0, 1)
    ax.set_xticks(np.arange(0, 1.01, 0.2))
    ax.set_ylim(0, 1)
    ax.set_yticks(np.arange(0, 1.01, 0.2))

    plt.tight_layout()
    fig.savefig(f"community3_{z_col}.pdf", format="pdf", bbox_inches="tight", transparent=True)
    plt.show()


# Heatmap A: feasibility
plot_surface_from_agg(agg, 'feas_mean', ' ')

# Heatmap B: stability
plot_surface_from_agg(agg, 'ev_mean', ' ')