aqacf-ml-main/firealarm_ml/train_classifiers.py

#!/usr/bin/env python
import os
import click
from pathlib import Path
import json
import pickle
import numpy as np
from tqdm import tqdm
from collections import defaultdict
from torch.utils.data import Dataset
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import (
    balanced_accuracy_score,
    accuracy_score,
)


class SmokeDataset(Dataset):


    def __init__(self, dates, X, y, distance_bins_km, density_categories,
            sequence_length, look_ahead):

        self.dates = dates
        self.X = X
        self.y = y
        self.distance_bins_km = distance_bins_km
        self.density_categories = density_categories

        self.sequence_length = sequence_length
        self.look_ahead = look_ahead

        self._len = max(0, len(X) - (sequence_length + look_ahead) + 1)


    def __len__(self):
        return self._len


    def __getitem__(self, i):
        midpoint = (i + self.sequence_length)
        xx = self.X[i : midpoint].ravel()
        yy = self.y[midpoint : midpoint + self.look_ahead].ravel()
        return xx, yy


    def get_dates(self, i):
        midpoint = (i + self.sequence_length)
        dx = self.dates[i : midpoint].ravel()
        dy = self.dates[midpoint : midpoint + self.look_ahead].ravel()
        return dx, dy


def load_training_config(configfile):
    with open(configfile, 'r') as f:
        return json.load(f)


def load_datasets(datafile, config):
    data = np.load(datafile)

    dtype = 'datetime64[D]'
    train_range = np.array(config['train_date_range'], dtype=dtype)
    val_range = np.array(config['validation_date_range'], dtype=dtype)
    test_range = np.array(config['test_date_range'], dtype=dtype)

    dates = data['dates']
    dist_bins = data['distance_bins_km']
    density_categories = data['density_categories']

    get_idx = lambda d_range: np.logical_and(
        d_range[0] <= dates,
        dates <= d_range[1]
    )

    idx = {
        'train': get_idx(train_range),
        'validation': get_idx(val_range),
        'test': get_idx(test_range),
    }

    datasets = []
    for Xi, yi in zip(data['features'], data['densities']):
        datasets.append({
            subset: SmokeDataset(
                dates[jj], Xi[jj], yi[jj],
                dist_bins, density_categories,
                config['sequence_length'],
                config['look_ahead'],
            )
            for subset, jj in idx.items()
        })

    site_info = {
        'latitudes': data['latitudes'],
        'longitudes': data['longitudes'],
        'names': data['names'],
    }

    return site_info, datasets


def train_eval_classifier(dataset):
    datasets = {
        dname: tuple(map(np.vstack, zip(*
            [dataset[dname][i] for i in range(len(dataset[dname]))]
        )))
        for dname in ('train', 'validation', 'test')
    }

    cls = RandomForestClassifier()
    cls.fit(*datasets['train'])

    statistics = defaultdict(dict)

    for dname, (X, y) in datasets.items():
        print(f'Evaluationg {dname}')
        ypred = cls.predict(X)

        statistics[dname]['balanced_accuracy'] = [
            balanced_accuracy_score(y[:, i], ypred[:, i], adjusted=True)
            for i in range(y.shape[1])
        ]

        statistics[dname]['smoke_accuracy'] = [
            balanced_accuracy_score(y[:, i] > 0, ypred[:, i] > 0, adjusted=True)
            for i in range(y.shape[1])
        ]

    print(statistics)

    return cls, dict(statistics)


@click.command()
@click.argument('datafile', type=click.Path(
    path_type=Path, exists=True
))
@click.argument('configfile', type=click.Path(
    path_type=Path, exists=True
))
@click.argument('outputfile', type=click.Path(
    path_type=Path
))
def main(datafile, configfile, outputfile):

    config = load_training_config(configfile)

    site_info, datasets = load_datasets(datafile, config)

    result = {
        'configuration': config,
        'classifiers': [],
        'statistics': [],
    }

    for name, dset in zip(site_info['names'], datasets):
        print(f'Training classifier for {name}...')
        cls, statistics = train_eval_classifier(dset)
        result['classifiers'].append(cls)
        result['statistics'].append(statistics)

    with open(outputfile, 'wb') as f:
        pickle.dump(result, f)

if __name__ == '__main__':
    main()