plot_social.py

#!/usr/bin/env python
"""Combined social figure: top row = 4 neighbourhood examples, bottom row = interaction statistics."""
import os, json, numpy as np, pandas as pd
import matplotlib as mpl; mpl.use('Agg')
import figstyle
from figstyle import INK, INK2, WIDTH
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from matplotlib.patches import Circle
from matplotlib import gridspec

DATA=os.environ.get("MARIS_DATA","/path/to/maris")  # set $MARIS_DATA to your Zenodo/HF download root
REL="multi_type_mini_bench_build/standard_track_v1"
def _leaf(region):
    """Resolve a leaf dir, preferring the Zenodo/HF release layout, else the build layout."""
    r=region.lower()
    cands=[os.path.join(DATA,"track_a_short-term_Cross-domain_Datasets",f"{r}_track_v1"),
           os.path.join(DATA,SRCDIR[region],REL)]
    for c in cands:
        if os.path.isdir(c): return c
    return cands[0]

OUT=os.path.join(os.path.dirname(os.path.abspath(__file__)),"out"); os.makedirs(OUT,exist_ok=True)
SRCDIR={'DMA':'DMA_ship_trajectory_datasets','NOAA':'NOAA_ship_trajectory_datasets',
        'Piraeus':'Piraeus_ship_trajectory_datasets','Norway':'norway_ship_trajectory_datasets'}
ORDER=['DMA','NOAA','Piraeus','Norway']
SC={'DMA':'#2a78d6','NOAA':'#1baf7a','Piraeus':'#eb6834','Norway':'#4a3aa7'}

def sfeat(src):
    return pd.read_csv(os.path.join(_leaf(src),'context_v1/social/features/train/social_features.csv'))

# examples from DMA
dma=sfeat('DMA').set_index('sample_id')
SAMPLES=['305203000_000001_r00_c00_02120','305299000_000000_r00_c00_00195',
         '220345000_000068_r00_c00_00080','211169080_000013_r05_c00_00120']
# stats from all sources
social={n:sfeat(n) for n in ORDER}
for n in ORDER: social[n]['neighbor_count_used']=social[n]['neighbor_count_used'].fillna(0).astype(int)

allsp=[]
for sid in SAMPLES:
    r=dma.loc[sid]
    vx=np.array(json.loads(r['neighbor_rel_vx_mps_json'])); vy=np.array(json.loads(r['neighbor_rel_vy_mps_json']))
    allsp.append(np.hypot(vx,vy))
vmax=np.percentile(np.concatenate(allsp),95)
cmap=plt.get_cmap('viridis'); norm=mpl.colors.Normalize(vmin=0,vmax=vmax)

fig=plt.figure(figsize=(WIDTH,5.7))
gs=gridspec.GridSpec(2,12,height_ratios=[1.0,0.92],hspace=0.46,wspace=1.7,
                     left=0.065,right=0.895,top=0.905,bottom=0.095)
# ---- top row: 4 examples (each spans 3 cols)
letters='abcd'
for k,sid in enumerate(SAMPLES):
    ax=fig.add_subplot(gs[0,k*3:(k+1)*3])
    r=dma.loc[sid]
    rx=np.array(json.loads(r['neighbor_rel_x_m_json']))/1000
    ry=np.array(json.loads(r['neighbor_rel_y_m_json']))/1000
    vx=np.array(json.loads(r['neighbor_rel_vx_mps_json'])); vy=np.array(json.loads(r['neighbor_rel_vy_mps_json']))
    sp=np.hypot(vx,vy)
    ax.add_patch(Circle((0,0),3,fill=True,fc='#f2f6fb',ec='#c9d3e0',lw=1.0,ls=(0,(5,3)),zorder=0))
    L=0.62
    for x,y,ux,uy in zip(rx,ry,vx,vy):
        m=np.hypot(ux,uy)
        if m>1e-6: ax.annotate('',xy=(x+ux/m*L,y+uy/m*L),xytext=(x,y),
                    arrowprops=dict(arrowstyle='-|>',color='#6b7688',lw=1.0,mutation_scale=8),zorder=2)
    scc=ax.scatter(rx,ry,c=sp,cmap=cmap,norm=norm,s=78,edgecolor='white',linewidth=1.0,zorder=3)
    ax.scatter([0],[0],marker='*',s=210,c='#d11',edgecolor='white',linewidth=0.8,zorder=4)
    ax.set_xlim(-3.4,3.4); ax.set_ylim(-3.4,3.4); ax.set_aspect('equal')
    ax.set_xticks([-2,0,2]); ax.set_yticks([-2,0,2]); ax.tick_params(labelsize=8,length=2)
    ax.set_xlabel('rel. x (km)',fontsize=8.5)
    if k==0: ax.set_ylabel('rel. y (km)',fontsize=8.5)
    n=int(r['neighbor_count_used'])
    ax.set_title(f'{letters[k]}  N={n} · {r["ship_type"]}',fontsize=9.3,weight='bold',loc='left')
    for s2 in ['top','right']: ax.spines[s2].set_visible(False)
# colorbar for examples (thin, right of top row)
cax=fig.add_axes([0.925,0.585,0.011,0.29])
cb=fig.colorbar(scc,cax=cax); cb.set_label('rel. speed (m s$^{-1}$)',fontsize=8); cb.ax.tick_params(labelsize=7.2)

# ---- bottom row: 3 stat panels (each spans 4 cols)
# (e) interaction density
ax=fig.add_subplot(gs[1,0:4])
frac=[(social[n]['neighbor_count_used']>0).mean()*100 for n in ORDER]
xb=np.arange(len(ORDER))
ax.bar(xb,frac,width=0.62,color=[SC[n] for n in ORDER],edgecolor='white',linewidth=0.7)
for x,f in zip(xb,frac): ax.text(x,f+0.5,f'{f:.1f}%',ha='center',va='bottom',fontsize=8,weight='bold',color=INK)
ax.set_xticks(xb); ax.set_xticklabels(ORDER,fontsize=8.5); ax.set_ylim(0,30)
ax.set_ylabel('samples with $\\geq$1 neighbour (%)',fontsize=8.3)
ax.set_title('e  Interaction density',loc='left',weight='bold',fontsize=9.3)
for s2 in ['top','right']: ax.spines[s2].set_visible(False);
ax.tick_params(length=2)
# (f) CPA
ax=fig.add_subplot(gs[1,4:8])
bins=np.linspace(0,1500,40)
for n in ORDER:
    s=social[n]; v=s.loc[s['neighbor_count_used']>0,'min_cpa_m'].dropna().values; v=v[(v>=0)&(v<=1500)]
    if len(v)>30: ax.hist(v,bins=bins,density=True,histtype='step',lw=1.5,color=SC[n],label=n)
ax.set_xlim(0,1500); ax.set_xlabel('minimum CPA (m)',fontsize=8.5); ax.set_ylabel('density',fontsize=8.5)
ax.set_title('f  Closest point of approach',loc='left',weight='bold',fontsize=9.3)
ax.legend(frameon=False,fontsize=8,handlelength=1.1)
for s2 in ['top','right']: ax.spines[s2].set_visible(False)
ax.tick_params(length=2,labelsize=8); ax.ticklabel_format(axis='y',style='sci',scilimits=(0,0))
# (g) TCPA
ax=fig.add_subplot(gs[1,8:12])
bins=np.linspace(0,900,40)
for n in ORDER:
    s=social[n]; v=s.loc[s['neighbor_count_used']>0,'min_abs_tcpa_s'].dropna().values; v=v[(v>=0)&(v<=900)]
    if len(v)>30: ax.hist(v,bins=bins,density=True,histtype='step',lw=1.5,color=SC[n],label=n)
ax.set_xlim(0,900); ax.set_xlabel('minimum |TCPA| (s)',fontsize=8.5); ax.set_ylabel('density',fontsize=8.5)
ax.set_title('g  Time to CPA',loc='left',weight='bold',fontsize=9.3)
for s2 in ['top','right']: ax.spines[s2].set_visible(False)
ax.tick_params(length=2,labelsize=8); ax.ticklabel_format(axis='y',style='sci',scilimits=(0,0))

# legend for top row (place between rows)
handles=[Line2D([0],[0],marker='*',color='w',markerfacecolor='#d11',markersize=12,markeredgecolor='gray',label='ego (target) vessel'),
         Line2D([0],[0],marker='o',color='w',markerfacecolor='#4c9a8a',markersize=8,markeredgecolor='white',label='neighbour'),
         Line2D([0],[0],color='#6b7688',lw=1.1,marker='>',markersize=5,label='relative-motion direction'),
         Line2D([0],[0],color='#c9d3e0',lw=1.1,ls=(0,(5,3)),label='3 km social radius')]
fig.legend(handles=handles,loc='center',ncol=4,frameon=False,fontsize=8.3,bbox_to_anchor=(0.47,0.515))
fig.savefig(os.path.join(OUT,'fig4_social_context.pdf'))
print("wrote combined fig4_social_context.pdf")