plot_composition_geographic_transfer.py

#!/usr/bin/env python
"""Aggregate cross-domain statistics figures for the EnvShip-Bench data descriptor.
Outputs: fig_stats.pdf, fig_geo.pdf, fig_social_stats.pdf, fig7_xdomain.pdf"""
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.ticker import PercentFormatter
from scipy.stats import gaussian_kde

DATA=os.environ.get("MARIS_DATA","/path/to/maris")  # set $MARIS_DATA to your Zenodo/HF download root
OUT=os.path.join(os.path.dirname(os.path.abspath(__file__)),"out"); os.makedirs(OUT,exist_ok=True)
os.makedirs(OUT, exist_ok=True)
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,SRC[region],REL)]
    for c in cands:
        if os.path.isdir(c): return c
    return cands[0]

SRC={'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']
# validated categorical palette (blue / aqua / orange / violet)
SC={'DMA':'#2a78d6','NOAA':'#1baf7a','Piraeus':'#eb6834','Norway':'#4a3aa7'}
REGION={'DMA':'Danish waters','NOAA':'U.S. coastal waters','Piraeus':'Piraeus / Saronic Gulf','Norway':'Norwegian coast'}

GRIDC='#e6e6e2'

def base(name): return _leaf(name)

# ---------------------------------------------------------------- load
print("loading...")
comp={}      # ship-class counts
scene={}     # scene counts
speed={}     # mean SOG per sample (sampled)
disp={}      # hist displacement
anchors={}   # lat/lon
social={}    # neighbor_count, cpa, tcpa
SHIP_CLASSES=['cargo','tanker','passenger','fishing','tug','service','sailing_leisure','unknown']
SCENES=['open_water','nearshore','constrained','harbor']
for n in ORDER:
    b=base(n)
    head=pd.read_csv(os.path.join(b,'train','part-000.csv.gz'),nrows=1)
    clscol='ship_class_unified' if 'ship_class_unified' in head.columns else 'ship_class'
    df=pd.read_csv(os.path.join(b,'train','part-000.csv.gz'),
                   usecols=[clscol,'hist_displacement_m','hist_sog_json'])
    cc=df[clscol].fillna('unknown')
    cc=cc.where(cc.isin(SHIP_CLASSES),'unknown')
    comp[n]=cc.value_counts()
    disp[n]=df['hist_displacement_m'].dropna().values
    samp=df['hist_sog_json'].dropna().sample(min(9000,len(df)),random_state=0)
    speed[n]=samp.apply(lambda s:float(np.mean(json.loads(s)))).values
    e=pd.read_csv(os.path.join(b,'context_v1/environment/all_environment_descriptors.csv'),usecols=['scene_type'])
    scene[n]=e['scene_type'].value_counts()
    a=pd.read_csv(os.path.join(b,'context_v1/environment/anchors/all_anchors.csv'),usecols=['anchor_lat','anchor_lon'])
    anchors[n]=a
    sf=pd.read_csv(os.path.join(b,'context_v1/social/features/train/social_features.csv'),
                   usecols=['neighbor_count_used','min_cpa_m','min_abs_tcpa_s'])
    sf['neighbor_count_used']=sf['neighbor_count_used'].fillna(0).astype(int)
    social[n]=sf
    print(" ",n,"done")

# ================================================================ FIG STATS (2x2)
fig,axes=plt.subplots(2,2,figsize=(WIDTH,5.4))
def style(ax):
    ax.spines['top'].set_visible(False); ax.spines['right'].set_visible(False)
    ax.tick_params(length=2,color=INK2)

# (a) ship-class composition stacked horizontal
ax=axes[0,0]
CLASS_COL=['#2a78d6','#1baf7a','#eda100','#4a3aa7','#e34948','#e87ba4','#008300','#c2bfb6']
props=np.zeros((len(ORDER),len(SHIP_CLASSES)))
for i,n in enumerate(ORDER):
    tot=comp[n].sum()
    for j,c in enumerate(SHIP_CLASSES):
        props[i,j]=comp[n].get(c,0)/tot*100
ypos=np.arange(len(ORDER))[::-1]
left=np.zeros(len(ORDER))
for j,c in enumerate(SHIP_CLASSES):
    ax.barh(ypos,props[:,j],left=left,height=0.62,color=CLASS_COL[j],
            edgecolor='white',linewidth=0.6,label=c.replace('_',' '))
    left+=props[:,j]
ax.set_yticks(ypos); ax.set_yticklabels(ORDER)
ax.set_xlim(0,100); ax.set_xlabel('share of samples (%)')
ax.set_title('a  Vessel-class composition',loc='left',weight='bold',fontsize=9.5)
ax.legend(ncol=4,fontsize=6.3,frameon=False,loc='upper center',bbox_to_anchor=(0.5,-0.16),
          handlelength=1.0,columnspacing=0.9,handletextpad=0.4)
style(ax)

# (b) speed regime: smooth KDE curves (no fill), median dashed lines
ax=axes[0,1]
xs=np.linspace(0,22,240)
for n in ORDER:
    v=speed[n]; v=v[np.isfinite(v)]
    ax.plot(xs,gaussian_kde(v)(xs),color=SC[n],lw=1.7,label=n)
    ax.axvline(np.median(v),color=SC[n],lw=0.9,ls=(0,(3,2)),alpha=0.75)
ax.set_xlim(0,22); ax.set_ylim(bottom=0)
ax.set_xlabel('mean speed over history window (kn)'); ax.set_ylabel('density')
ax.set_title('b  Speed regime',loc='left',weight='bold',fontsize=9.5)
ax.legend(frameon=False,handlelength=1.1); style(ax)

# (c) motion scale: light histogram + overlaid KDE, y as share of samples (%)
ax=axes[1,0]
bins=np.linspace(0,9000,31); bw=bins[1]-bins[0]; xs=np.linspace(0,9000,240)
for n in ORDER:
    d=disp[n]; d=d[(d>=0)&(d<=9000)]
    ax.hist(d,bins=bins,weights=np.ones(len(d))*100.0/len(d),color=SC[n],alpha=0.20,
            edgecolor=SC[n],linewidth=0.4)
    ax.plot(xs,gaussian_kde(d)(xs)*bw*100,color=SC[n],lw=1.6,label=n)
ax.set_xlim(0,9000); ax.set_xticks([0,2000,4000,6000,8000]); ax.set_ylim(bottom=0)
ax.set_xlabel('history displacement (m)'); ax.set_ylabel('share of samples (%)')
ax.set_title('c  Motion scale',loc='left',weight='bold',fontsize=9.5)
ax.legend(frameon=False,handlelength=1.1); style(ax)

# (d) scene composition stacked horizontal
ax=axes[1,1]
SCENE_COL={'open_water':'#9ec5f4','nearshore':'#1baf7a','constrained':'#eda100','harbor':'#e34948'}
sprops=np.zeros((len(ORDER),len(SCENES)))
for i,n in enumerate(ORDER):
    tot=scene[n].sum()
    for j,s in enumerate(SCENES): sprops[i,j]=scene[n].get(s,0)/tot*100
left=np.zeros(len(ORDER))
for j,s in enumerate(SCENES):
    ax.barh(ypos,sprops[:,j],left=left,height=0.62,color=SCENE_COL[s],
            edgecolor='white',linewidth=0.6,label=s.replace('_',' '))
    left+=sprops[:,j]
ax.set_yticks(ypos); ax.set_yticklabels(ORDER)
ax.set_xlim(0,100); ax.set_xlabel('share of samples (%)')
ax.set_title('d  Scene composition',loc='left',weight='bold',fontsize=9.5)
ax.legend(ncol=4,fontsize=6.6,frameon=False,loc='upper center',bbox_to_anchor=(0.5,-0.16),
          handlelength=1.0,columnspacing=0.9,handletextpad=0.4)
style(ax)
fig.tight_layout(w_pad=2.4,h_pad=3.2)
fig.savefig(os.path.join(OUT,'fig_stats.pdf')); plt.close(fig); print("wrote fig_stats")

# ================================================================ FIG GEO (2x2 uniform rectangles)
from matplotlib.colors import LogNorm, LinearSegmentedColormap
ocean=LinearSegmentedColormap.from_list('ocean',
    ['#dcecf8','#a9d1ec','#6fb0dd','#3f8aca','#215fa6','#123f73','#0a2850'])
# per-source 2-D density on a robust extent
hists={}
for n in ORDER:
    a=anchors[n]; lon=a['anchor_lon'].values; lat=a['anchor_lat'].values
    lo_lon,hi_lon=np.percentile(lon,[1,99]); lo_lat,hi_lat=np.percentile(lat,[1,99])
    m=(lon>=lo_lon)&(lon<=hi_lon)&(lat>=lo_lat)&(lat<=hi_lat)
    H,xe,ye=np.histogram2d(lon[m],lat[m],bins=75)
    hists[n]=(H,xe,ye)
vmax=max(h[0].max() for h in hists.values())
fig,axes=plt.subplots(2,2,figsize=(WIDTH,5.15))
for ax,n in zip(axes.flat,ORDER):
    H,xe,ye=hists[n]
    Hm=np.ma.masked_where(H.T<=0,H.T)
    ax.set_facecolor('#f3f8fc')
    im=ax.imshow(Hm,extent=[xe[0],xe[-1],ye[0],ye[-1]],origin='lower',aspect='auto',
                 cmap=ocean,norm=LogNorm(vmin=1,vmax=vmax),interpolation='nearest')
    ax.set_title(f'{n} — {REGION[n]}',fontsize=9.5,weight='bold',color=INK,pad=4)
    ax.set_xlabel('longitude (°E)',fontsize=8.5)
    ax.set_ylabel('latitude (°N)',fontsize=8.5)
    ax.tick_params(length=2,color=INK2,labelsize=8)
    for sp in ax.spines.values(): sp.set_edgecolor('#b8c4d0')
fig.subplots_adjust(left=0.075,right=0.885,top=0.95,bottom=0.085,wspace=0.24,hspace=0.34)
cax=fig.add_axes([0.905,0.20,0.016,0.60])
cb=fig.colorbar(im,cax=cax); cb.set_label('anchors per cell (log scale)',fontsize=8.5)
cb.ax.tick_params(labelsize=7.5)
fig.savefig(os.path.join(OUT,'fig_geo.pdf')); plt.close(fig); print("wrote fig_geo")

# ================================================================ FIG SOCIAL STATS (1x3)
fig,axes=plt.subplots(1,3,figsize=(11.0,3.1))
# (a) % with >=1 neighbour + interaction candidate
ax=axes[0]
frac=[ (social[n]['neighbor_count_used']>0).mean()*100 for n in ORDER]
xb=np.arange(len(ORDER))
bars=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=7.8,weight='bold',color=INK)
ax.set_xticks(xb); ax.set_xticklabels(ORDER)
ax.set_ylabel('samples with $\\geq$1 neighbour (%)'); ax.set_ylim(0,30)
ax.set_title('a  Interaction density',loc='left',weight='bold',fontsize=9.5)
for sp in ['top','right']: ax.spines[sp].set_visible(False)
ax.tick_params(length=2,color=INK2)
# (b) CPA distribution (samples with neighbours)
ax=axes[1]
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.6,color=SC[n],label=n)
ax.set_xlim(0,1500); ax.set_xlabel('minimum CPA (m)'); ax.set_ylabel('density')
ax.set_title('b  Closest point of approach',loc='left',weight='bold',fontsize=9.5)
ax.legend(frameon=False,fontsize=7.3,handlelength=1.1)
for sp in ['top','right']: ax.spines[sp].set_visible(False)
ax.tick_params(length=2,color=INK2); ax.ticklabel_format(axis='y',style='sci',scilimits=(0,0))
# (c) TCPA distribution
ax=axes[2]
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.6,color=SC[n],label=n)
ax.set_xlim(0,900); ax.set_xlabel('minimum |TCPA| (s)'); ax.set_ylabel('density')
ax.set_title('c  Time to CPA',loc='left',weight='bold',fontsize=9.5)
for sp in ['top','right']: ax.spines[sp].set_visible(False)
ax.tick_params(length=2,color=INK2); ax.ticklabel_format(axis='y',style='sci',scilimits=(0,0))
fig.tight_layout(w_pad=2.0)
fig.savefig(os.path.join(OUT,'fig_social_stats.pdf')); plt.close(fig); print("wrote fig_social_stats")

# ================================================================ FIG7 cross-domain heatmap
jurs=['DMA','NOAA','Piraeus','Norway']
matrices={
'TCN':np.array([[87.8,91.2,163.2,124.9],[np.nan,94.1,162.9,128.6],[89.6,np.nan,163.1,125.5],[87.2,91.8,np.nan,113.0],[87.7,91.0,162.6,np.nan]]),
'LSTM+Env-SDF':np.array([[92.0,103.8,144.8,131.8],[np.nan,109.9,156.1,130.2],[89.5,np.nan,147.6,128.8],[95.1,104.2,np.nan,120.5],[85.5,95.6,145.1,np.nan]]),
'LSTM+Soc+Env-SDF':np.array([[91.9,100.6,149.5,128.5],[np.nan,106.3,156.6,131.7],[92.0,np.nan,152.9,132.1],[91.9,100.5,np.nan,123.8],[90.4,97.5,149.1,np.nan]])}
rows=['Combined','no-DMA','no-NOAA','no-Piraeus','no-Norway']
from matplotlib.colors import Normalize
cmap=plt.get_cmap('YlGnBu'); norm=Normalize(vmin=80,vmax=200)
fig,axes=plt.subplots(1,3,figsize=(WIDTH,2.95),constrained_layout=True)
for k,(ax,(title,mat)) in enumerate(zip(axes,matrices.items())):
    im=ax.imshow(mat,cmap=cmap,norm=norm)
    ax.set_title(title,weight='bold',fontsize=10)
    ax.set_xticks(np.arange(4)); ax.set_xticklabels(jurs,fontsize=9)
    ax.set_yticks(np.arange(5)); ax.set_yticklabels(rows if k==0 else [],fontsize=9)
    if k==0: ax.set_ylabel('training pool',fontsize=9.5)
    for i in range(mat.shape[0]):
        for j in range(mat.shape[1]):
            if np.isnan(mat[i,j]):
                ax.text(j,i,'—',ha='center',va='center',color='#b6b3ab',fontsize=11,weight='bold')
            else:
                r,g,b,_=cmap(norm(mat[i,j])); lum=0.299*r+0.587*g+0.114*b
                ax.text(j,i,f'{mat[i,j]:.0f}',ha='center',va='center',fontsize=9.2,
                        color='white' if lum<0.55 else INK,
                        weight='bold' if lum<0.55 else 'normal')
    ax.tick_params(length=0)
    for sp in ax.spines.values(): sp.set_visible(False)
cb=fig.colorbar(im,ax=axes.tolist(),shrink=0.9,pad=0.012,aspect=30)
cb.set_label('ADE (m)',fontsize=9); cb.ax.tick_params(labelsize=8)
fig.savefig(os.path.join(OUT,'fig9_xdomain.pdf')); plt.close(fig); print("wrote fig9_xdomain")
print("ALL DONE")