File size: 4,341 Bytes
34c30a6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """Synthetic fleet generator for the SCALING evaluation.
Fleet at scale k = k time-shifted clones of each real building. Each clone
building gets ONE shift (preserves within-building death correlation,
decorrelates clone sets). Ground-truth EOL is DERIVED from the shifted
telemetry with the official mechanism (7d-median < 2.40, validated 81/81
by the team repo) - labels stay mechanically consistent with the data.
Travel: distance between clones = distance between base buildings;
same-base clone pairs get the median off-diagonal hop.
Writes /tmp/c5/fleet_k{K}.pkl with: timeseries, locations, eol, travel fn.
"""
import sys, pickle, time
from pathlib import Path
import numpy as np
import pandas as pd
TEAM = Path('/tmp/teamrepo')
OURS = Path('/Users/THIF/repos/BatterySwapAI2026-Example')
sys.path.insert(0, str(TEAM))
sys.path.insert(0, str(OURS))
from batteryswap_public.utils import load_dataset
from features import daily_panel, eol_from_panel
K = int(sys.argv[1]) if len(sys.argv) > 1 else 2
SEED = 7
locations, ts, eol_times, scenarios = load_dataset(OURS / 'dataset/train')
df = ts.reset_index()
df['device_id'] = df['device_id'].astype(str)
dev = pd.read_csv(OURS / 'dataset/train/devices.csv').drop_duplicates('device_id')
span_lo = df['end_time'].min()
span_hi = df['end_time'].max()
rng = np.random.default_rng(SEED)
frames, loc_rows = [], []
for c in range(K):
for b, sub in dev.groupby('building_id'):
# clone 0 of each building = the untouched original (scale anchor)
shift = 0 if c == 0 else int(rng.integers(-120, 121))
sb = b if c == 0 else f'{b}_c{c}'
for _, r in sub.iterrows():
sd = r['device_id'] if c == 0 else f"{r['device_id']}_c{c}"
d = df[df['device_id'] == r['device_id']].copy()
d['device_id'] = sd
d['end_time'] = d['end_time'] + pd.Timedelta(days=shift)
d = d[(d['end_time'] >= span_lo) & (d['end_time'] <= span_hi)]
if len(d) < 50:
continue
frames.append(d)
loc_rows.append(dict(battery=sd, building=sb,
room=(r['room_id'] if c == 0
else f"{r['room_id']}_c{c}"),
base_building=b,
start_time=pd.Timestamp(r['start_time'])
+ pd.Timedelta(days=shift),
end_time=pd.Timestamp(r['end_time'])
+ pd.Timedelta(days=shift)))
sts = pd.concat(frames, ignore_index=True)
locs = pd.DataFrame(loc_rows)
for col in ('start_time', 'end_time'):
locs[col] = pd.to_datetime(locs[col], utc=True)
t0 = time.time()
panel = daily_panel(sts.set_index(['device_id', 'end_time']))
eol = eol_from_panel(panel)
print(f'k={K}: {locs.battery.nunique()} devices, '
f'{locs.building.nunique()} buildings, {len(eol)} EOLs '
f'(panel {time.time()-t0:.0f}s)')
# sanity: clone-0 EOLs must match the official file exactly
official = pd.read_csv(OURS / 'dataset/train/eol_times.csv')
col = [c for c in official.columns if c != 'device_id'][0]
truth = pd.to_datetime(official.set_index('device_id')[col]).dt.tz_localize(None)
c0 = eol[~eol.index.str.contains('_c')]
both = truth.dropna().index.intersection(c0.index)
mism = int((pd.to_datetime(c0.reindex(both)).dt.normalize()
!= truth.reindex(both).dt.normalize()).sum())
print(f'clone-0 EOL check: {len(both)} deaths, {mism} mismatches')
# travel-cost template: base 24x24 matrix from scenario 0
tc0 = scenarios[0]['travel_costs']
base_h = tc0.set_index(['from', 'to'])['hours']
med_hop = tc0.loc[tc0['from'] != tc0['to'], 'hours'].median()
blds = sorted(locs['building'].unique())
base_of = locs.drop_duplicates('building').set_index('building')['base_building']
rows = []
for i in blds:
for j in blds:
if i == j:
h = 0.0
else:
bi, bj = base_of[i], base_of[j]
h = med_hop if bi == bj else float(base_h.get((bi, bj), med_hop))
rows.append(dict(**{'from': i, 'to': j}, hours=h))
travel = pd.DataFrame(rows)
with open(f'/tmp/c5/fleet_k{K}.pkl', 'wb') as f:
pickle.dump(dict(ts=sts, locations=locs, eol=pd.to_datetime(eol),
travel=travel), f)
print(f'wrote /tmp/c5/fleet_k{K}.pkl')
|