aqacf-ml-main/firealarm_ml/eval_seq2seq.py

#!/usr/bin/env python
import os
import json
import click
import numpy as np
import xarray as xr
from tqdm import tqdm
import torch
import matplotlib.pyplot as plt

from epa_seq2seq import (
    EpaDataset, EpaSeq2Seq,
    eval_model, mse_nan_loss,
)


@click.command()
@click.argument('modelfile')
@click.argument('datafile')
@click.argument('outputfile')
def main(modelfile, datafile, outputfile):

    device = torch.device('cpu')
    checkpoint = torch.load(modelfile, map_location=device)
    config = checkpoint['config']

    window_size = config['sequence_length']
    in_channels = config['in_channels']
    out_channels = config['out_channels']
    model_params = config['model_params']
    model_params['device'] = device

    dataset_names = ('test',)
    datasets = {}

    with xr.open_zarr(datafile) as ds:
        for s in dataset_names:
            ds_sub = ds.sel(time=slice(*config[f'{s}_date_range']))
            datasets[s] = EpaDataset(ds_sub, window_size, in_channels, out_channels)

    model = EpaSeq2Seq(
        in_channels=len(in_channels),
        out_channels=len(out_channels),
        frame_size=datasets['test'].frame_size,
        **model_params
    )
    model.load_state_dict(checkpoint['model_state_dict'])
    model.to(device)
    model.eval()

    test = datasets['test']

    all_times = []
    predictions = []

    for i in tqdm(list(range(len(test)))):
        X, y = test[i]
        times = test.get_time(i)
        X = torch.tensor(X[np.newaxis, ...])
        y = torch.tensor(y[np.newaxis, ...])
        ypred = np.squeeze(model(X).detach().numpy())
        predictions.append(ypred[-1])
        all_times.append(times[-1])

    predictions = np.stack(predictions, axis=0)

    dataset = xr.Dataset(
        data_vars={
            'PM25': (
                ['time', 'lat', 'lon'],
                predictions,
                test.y['PM25'].attrs
            )
        },
        coords={
            'time': xr.Variable('time', np.array(all_times)),
            'lat': test.dataset.lat,
            'lon': test.dataset.lon,
        }
    ).chunk(time=8)

    dataset.to_zarr(outputfile)


if __name__ == '__main__':
    main()