aqacf-ml-main/firealarm_ml/deploy.py

import json
import requests
import numpy as np
from datetime import datetime
from collections import defaultdict
import torch

from epa_seq2seq_model import EpaSeq2Seq


def construct_url(endpoint, bounds, dataset, start_time, end_time):
    lat_lo, lat_hi = bounds['latitude']
    lon_lo, lon_hi = bounds['longitude']
    stime = start_time.strftime('%Y-%m-%dT%H:%M:%SZ')
    etime = end_time.strftime('%Y-%m-%dT%H:%M:%SZ')
    bstr = f'{lon_lo},{lat_lo},{lon_hi},{lat_hi}'
    args = f'ds={dataset}&b={bstr}&startTime={stime}&endTime={etime}'
    return f'{endpoint}?{args}&lowPassFilter=False'


def grid_data(data):
    latitudes = set([])
    longitudes = set([])
    times = set([])
    datadict = defaultdict(lambda: np.nan)
    for item in data:
        lat = item['latitude']
        lon = item['longitude']
        time = item['time']
        value = item['data'][0]['variable']
        latitudes.add(lat)
        longitudes.add(lon)
        times.add(time)
        datadict[lat, lon, time] = value

    latitudes = sorted(latitudes)
    longitudes = sorted(longitudes)
    times = sorted(times)

    D = np.array([
        [
            [
                datadict[lat, lon, time]
                for lon in longitudes
            ]
            for lat in latitudes
        ]
        for time in times
    ], dtype=np.float32)

    return D, (times, latitudes, longitudes)


def fetch_data(config):
    data = {}

    start = datetime(2020, 9, 15)
    end = datetime(2020, 9, 16)
    for dname, dinfo in config['datasets'].items():
        url = construct_url(
            config['endpoint'], config['bounds'],
            dinfo['dataset'], start, end
        )
        print(f'Fetching "{url}"...')
        response = requests.get(url)
        raw = response.json()
        print(f'...done.')
        gridded, axes = grid_data(raw)
        if 'factor' in dinfo:
            gridded *= dinfo['factor']
        data[dname] = (gridded, axes)

    return data


def format_input(config, in_channels, window_size):
    datasets = fetch_data(config)
    relevant = {
        k: v for k, v in datasets.items() if k in in_channels
    }

    common_times = set.intersection(*[
        set(ax[0]) for _, ax in datasets.values()
    ])

    assert len(common_times) >= window_size

    latitudes = None
    longitudes = None
    times = sorted(common_times)
    raw = {}
    for k, (data, axes) in relevant.items():
        time, lat, lon = axes

        if latitudes is None:
            latitudes = lat
            longitudes = lon
        else:
            assert np.allclose(lat, latitudes)
            assert np.allclose(lon, longitudes)

        # Sub-select common times
        idx = np.array([ti in common_times for ti in time])
        raw[k] = data[idx, ...]

    X = np.stack([
        raw[c][-window_size:]
        for c in in_channels
    ], dtype=np.float32)
    frame_size = X.shape[-2:]
    return X, frame_size, (times[-window_size:], latitudes, longitudes)


def load_config(configpath):
    with open(configpath, 'r') as f:
        return json.load(f)


def predict(event, context):

    body = event['body']
    config_path = body['configuration_path']
    deploy_config = load_config(config_path)
    model_path = deploy_config['model_path']

    # Load model
    device = torch.device('cpu')
    checkpoint = torch.load(model_path, 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

    X, frame_size, axes = format_input(deploy_config, in_channels, window_size)
    X = torch.tensor(X[np.newaxis, ...])

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

    ypred = model(X).detach().numpy()[0, -1, ...]

    out_data = {
        'lat': axes[1],
        'lon': axes[2],
        'variables': {
            k: y.tolist()
            for k, y in zip(out_channels, ypred)
        }
    }

    result = {
        'data': out_data,
    }

    return {
        'statusCode': 200,
        'body': result,
    }


if __name__ == '__main__':
    event = {
        'body': {
            'configuration_path': 'config/deploy_config.json',
        }
    }
    response = predict(event, None)
    pred = response['body']['data']['variables']['PM25']
    import matplotlib.pyplot as plt
    Y = np.array(pred)
    plt.imshow(np.flipud(Y))
    plt.show()