plot_environment.py

#!/usr/bin/env python
"""Polished paired environmental representations figure (fig3)."""
import os, glob, 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.colors import LinearSegmentedColormap, TwoSlopeNorm
from matplotlib import gridspec
from mpl_toolkits.axes_grid1.inset_locator import inset_axes

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','Norway':'norway_ship_trajectory_datasets',
        'Piraeus':'Piraeus_ship_trajectory_datasets','NOAA':'NOAA_ship_trajectory_datasets'}

# chosen samples: (scene label, source, sample_id)
ROWS=[
 ('Harbour','Norway','259003010_000440_r00_c00_00025'),
 ('Constrained','DMA','218179000_000043_r00_c00_00065'),
 ('Nearshore','Piraeus','430429824_000029_r00_c00_00015'),
 ('Open water','DMA','205168000_000058_r04_c00_00040'),
]

def load(src):
    b=_leaf(src)
    sid=np.load(os.path.join(b,'context_v1/environment/rasters/train/sample_ids.npy'),allow_pickle=True)
    idx={s:i for i,s in enumerate(sid)}
    rdir=os.path.join(b,'context_v1/environment/rasters/train')
    _npz=os.path.join(rdir,'masks.npz')
    masks=np.load(_npz)['masks'] if os.path.exists(_npz) else np.load(os.path.join(rdir,'masks.npy'),mmap_mode='r')
    sdf=np.load(os.path.join(rdir,'signed_dist_shore.npy'),mmap_mode='r')
    _agg=os.path.join(b,'context_v1/environment/all_environment_descriptors.csv')
    if os.path.exists(_agg):
        ed=pd.read_csv(_agg).set_index('sample_id')
    else:  # release layout: concatenate per-split descriptor files
        _fs=sorted(glob.glob(os.path.join(b,'context_v1/environment/features/*/environment_descriptors.csv')))
        ed=pd.concat([pd.read_csv(f) for f in _fs],ignore_index=True).set_index('sample_id')
    traj=pd.read_csv(os.path.join(b,'train','part-000.csv.gz'),
                     usecols=['sample_id','hist_x_json','hist_y_json','fut_x_json','fut_y_json','hist_sog_json']).set_index('sample_id')
    return idx,masks,sdf,ed,traj

def raster_rgb(mk):
    land=mk[0].astype(bool); water=mk[1].astype(bool); nav=mk[2].astype(bool)
    manmade=mk[4].astype(bool); barrier=mk[5].astype(bool)
    H,W=land.shape
    img=np.ones((H,W,3))*np.array([0.93,0.88,0.79])
    img[water]=[0.74,0.86,0.96]
    img[nav]=[0.55,0.76,0.93]
    img[manmade]=[0.46,0.33,0.24]
    img[barrier]=[0.30,0.22,0.18]
    return img

# custom SDF colormap: brown(land)->white(0)->blue(water)
sdf_cmap=LinearSegmentedColormap.from_list('shore',
    [(0.0,'#6b4f2a'),(0.35,'#c8a97b'),(0.5,'#f7f4ee'),(0.62,'#9ec5f4'),(1.0,'#184f95')])

DESC=[('water ratio','water_ratio',1),('nav. ratio','navigable_ratio',1),
      ('shore dist','nearest_shore_dist_m',5000),('manmade dist','nearest_manmade_dist_m',5000),
      ('nat. bdry','natural_boundary_density',1),('manmade bdry','manmade_boundary_density',1),
      ('barrier den.','barrier_density',1),('env. quality','env_quality_score',1)]
DGROUP=['#2a78d6','#2a78d6','#1baf7a','#1baf7a','#eda100','#eda100','#e34948','#4a3aa7']

VMAX=3000.0  # fixed shared SDF display range (m); values beyond saturate
norm=TwoSlopeNorm(vmin=-VMAX,vcenter=0,vmax=VMAX)
fig=plt.figure(figsize=(WIDTH,8.4))
# columns: raster | sdf | spacer | descriptor ; shared horizontal colorbar under the two maps
gs=gridspec.GridSpec(4,4,width_ratios=[1,1,0.34,0.92],wspace=0.16,hspace=0.16,
                     left=0.115,right=0.885,top=0.93,bottom=0.115)
cache={}; im=None; raster_axes=[]; sdf_axes=[]
col_titles=['Six-channel OSM raster','Shore signed-distance field','Scene descriptor']
for r,(scene,src,sid) in enumerate(ROWS):
    if src not in cache: cache[src]=load(src)
    idx,masks,sdf,ed,traj=cache[src]
    i=idx[sid]
    # rasters are stored north-row-0; flip vertically so origin='lower' shows north up
    mk=np.asarray(masks[i])[:, ::-1, :]
    sd=np.asarray(sdf[i],dtype=float)[::-1, :]
    t=traj.loc[sid]
    hx=np.array(json.loads(t['hist_x_json']))/1000; hy=np.array(json.loads(t['hist_y_json']))/1000
    fx=np.array(json.loads(t['fut_x_json']))/1000; fy=np.array(json.loads(t['fut_y_json']))/1000
    spd=float(np.mean(json.loads(t['hist_sog_json'])))
    ext=[-5,5,-5,5]
    def draw_traj(ax):
        ax.plot(hx,hy,'-',color='#12347a',lw=1.7,solid_capstyle='round')
        ax.plot(fx,fy,'-',color='#0f8a3c',lw=1.7,solid_capstyle='round')
        ax.plot(hx[-1],hy[-1],'*',color='#d11',ms=11,mec='white',mew=0.6)
    # -- raster
    ax=fig.add_subplot(gs[r,0]); raster_axes.append(ax)
    ax.imshow(raster_rgb(mk),extent=ext,origin='lower',interpolation='nearest')
    draw_traj(ax)
    ax.set_xlim(-5,5); ax.set_ylim(-5,5); ax.set_xticks([-4,0,4]); ax.set_yticks([-4,0,4])
    ax.tick_params(length=2)
    ax.plot([1.6,3.6],[-4.4,-4.4],'-',color=INK,lw=2.4,solid_capstyle='butt')
    ax.text(2.6,-4.0,'2 km',ha='center',va='bottom',fontsize=7,color=INK)
    ax.set_ylabel(f'{scene}\n{spd:.1f} kn · {src}',weight='bold',labelpad=6)
    if r==0: ax.set_title(col_titles[0],weight='bold',pad=6)
    if r==len(ROWS)-1: ax.set_xlabel('local x (km)')
    # -- SDF
    ax=fig.add_subplot(gs[r,1]); sdf_axes.append(ax)
    im=ax.imshow(sd,extent=ext,origin='lower',cmap=sdf_cmap,norm=norm,interpolation='bilinear')
    gx=np.linspace(-5,5,sd.shape[1]); gy=np.linspace(-5,5,sd.shape[0])
    ax.contour(gx,gy,sd,levels=[0],colors='#222222',linewidths=0.8)
    draw_traj(ax)
    ax.set_xlim(-5,5); ax.set_ylim(-5,5); ax.set_xticks([-4,0,4]); ax.set_yticks([])
    ax.tick_params(length=2)
    if r==0: ax.set_title(col_titles[1],weight='bold',pad=6)
    if r==len(ROWS)-1: ax.set_xlabel('local x (km)')
    # -- descriptor (category labels on the left, values inside/beside bars)
    ax=fig.add_subplot(gs[r,3])
    vals=[]; labs=[]
    for lab,key,den in DESC:
        v=float(ed.loc[sid,key])/den; vals.append(min(max(v,0),1)); labs.append(lab)
    yp=np.arange(len(vals))[::-1]
    ax.barh(yp,vals,height=0.66,color=DGROUP,edgecolor='white',linewidth=0.5)
    for y,v in zip(yp,vals):
        ax.text(v-0.03 if v>0.16 else v+0.03,y,f'{v:.2f}',va='center',
                ha='right' if v>0.16 else 'left',fontsize=6.8,
                color='white' if v>0.16 else INK2)
    ax.set_yticks(yp); ax.set_yticklabels(labs)
    ax.set_xlim(0,1.0); ax.set_xticks([0,0.5,1.0]); ax.tick_params(length=2)
    ax.spines['top'].set_visible(False); ax.spines['right'].set_visible(False)
    if r==0: ax.set_title(col_titles[2],weight='bold',pad=6)
    if r==len(ROWS)-1: ax.set_xlabel('normalised value')

# one shared horizontal SDF colorbar under the two map columns
fig.canvas.draw()
x0=min(a.get_position().x0 for a in raster_axes)
x1=max(a.get_position().x1 for a in sdf_axes)
cax=fig.add_axes([x0,0.045,x1-x0,0.013])
cb=fig.colorbar(im,cax=cax,orientation='horizontal',extend='both')
cb.set_ticks([-VMAX,0,VMAX]); cb.set_ticklabels([f'$-${VMAX:.0f}','0',f'{VMAX:.0f}'])
cb.ax.tick_params(labelsize=7,length=2)
cb.set_label('signed distance to shore (m)  ·  brown: inland,  blue: offshore',fontsize=8)

# shared legend for trajectory
from matplotlib.lines import Line2D
handles=[Line2D([0],[0],color='#12347a',lw=2.2,label='history (10 min)'),
         Line2D([0],[0],color='#0f8a3c',lw=2.2,label='future (10 min)'),
         Line2D([0],[0],marker='*',color='w',markerfacecolor='#d11',markersize=11,label='anchor',markeredgecolor='gray')]
fig.legend(handles=handles,loc='upper center',ncol=3,frameon=False,fontsize=8.5,bbox_to_anchor=(0.5,0.99))
fig.savefig(os.path.join(OUT,'fig3_env_context.pdf'),bbox_inches='tight')
print("wrote fig3_env_context.pdf")