Test/testlib/resultset.py

# SPDX-License-Identifier: Apache-2.0
# -----------------------------------------------------------------------------
# Copyright 2020-2026 Arm Limited
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy
# of the License at:
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
# -----------------------------------------------------------------------------
'''
A ResultSet stores a set of results about the performance of a TestSet. Each
set keeps result Records for each image and block size tested, storing the
PSNR and coding time.

ResultSets are often backed by a CSV file on disk, and a ResultSet can be
compared against a set of reference results created by an earlier test run.
'''

from __future__ import annotations

import csv
import enum
from pathlib import Path
from typing import Optional

import numpy


@enum.unique
class ResultStatus(enum.IntEnum):
    '''
    An enumeration of test result status.

    Attributes:
        NOT_RUN: The test has not been run.
        PASS: The test passed.
        WARN: The quality was below pass threshold but above fail threshold.
        FAIL: The quality was below the fail threshold.
    '''
    NOT_RUN = 0
    PASS = 1
    WARN = 2
    FAIL = 3


class Record:
    '''
    A single result record, holding results for a single image and block size.

    Attributes:
        block_size: The block size.
        name: The image name.
        psnr: The image quality (PSNR dB)
        total_time: The total time.
        coding_time: The coding time.
        coding_rate: The coding rate.
        status: The test result status.
        psnr_rel: The relative image quality (diff PSNR dB, >= 0 good)
        total_time_rel: The relative total time vs ref (scale, <1 fast).
        coding_time_rel: The relative coding time vs ref (scale, <1 fast).
        coding_rate_rel: The relative coding rate vs ref (scale, <1 fast).
    '''

    def __init__(self, block_size: str, name: str, psnr: float,
                 total_time: float, coding_time: float, coding_rate: float):
        '''
        Create a result record, initially in the NOT_RUN status.

        Args:
            block_size (str): The block size.
            name (str): The test image name.
            psnr (float): The image quality PSNR, in dB.
            total_time (float): The total compression time, in seconds.
            coding_time (float): The coding compression time, in seconds.
            coding_rate (float): The coding compression rate, in MPix/s.
        '''
        self.block_size = block_size
        self.name = name
        self.psnr = psnr
        self.total_time = total_time
        self.coding_time = coding_time
        self.coding_rate = coding_rate
        self.status = ResultStatus.NOT_RUN

        self.psnr_rel: Optional[float] = None
        self.total_time_rel: Optional[float] = None
        self.coding_time_rel: Optional[float] = None
        self.coding_rate_rel: Optional[float] = None

    def set_status(self, result: ResultStatus) -> None:
        '''
        Set the result status.

        Args:
            result: The test result status to set.
        '''
        self.status = result

    def set_relative_to_reference(self, reference: Record) -> None:
        '''
        Set relative results compared to an existing reference score.

        Args:
            reference: The reference result to compare against.
        '''
        self.psnr_rel = self.psnr - reference.psnr

        try:
            self.total_time_rel = reference.total_time / self.total_time
        except ZeroDivisionError:
            self.total_time_rel = float('NaN')

        try:
            self.coding_time_rel = reference.coding_time / self.coding_time
        except ZeroDivisionError:
            self.coding_time_rel = float('NaN')


class ResultSet:
    '''
    A set of results for a TestSet, across one or more block sizes.

    Attributes:
        test_set_name: The name of the test set that generated these results.
        records: The list of test result records.
    '''

    def __init__(self, test_set_name: str):
        '''
        Create a new empty ResultSet.

        Args:
            test_set_name: The test set these results were generated by.
        '''
        self.test_set_name = test_set_name
        self.records: list[Record] = []

    def add_record(self, record: Record) -> None:
        '''
        Add a new test result record to this result set.

        Args:
            record: The test record to add.
        '''
        self.records.append(record)

    def get_matching_record(self, other: Record) -> Record:
        '''
        Get a record matching the config of another record.

        Args:
            other: The pattern record to match.

        Return:
            The result, if present.

        Raise:
            KeyError: No match could be found.
        '''
        for record in self.records:
            if record.block_size == other.block_size and \
               record.name == other.name:
                return record

        raise KeyError()

    def get_results_summary(self) -> ResultSummary:
        '''
        Get a results summary of all the records in this result set.

        Return:
            ResultSummary: The result summary.
        '''
        summary = ResultSummary()

        for record in self.records:
            summary.add_record(record)

        return summary

    def save_to_file(self, file_path: Path):
        '''
        Save this result set to a CSV file.

        Args:
            file_path: The output file path.
        '''
        dir_path = file_path.parent
        dir_path.mkdir(parents=True, exist_ok=True)

        with open(file_path, 'w', encoding='utf=8', newline='') as handle:
            writer = csv.writer(handle)
            self._save_header(writer)
            for record in self.records:
                self._save_record(writer, record)

    @staticmethod
    def _save_header(writer) -> None:
        '''
        Write the header to a result CSV file.

        Args:
            writer: The CSV writer.
        '''
        row = [
            'Image Set',
            'Block Size',
            'Name',
            'PSNR',
            'Total Time',
            'Coding Time',
            'Coding Rate'
        ]

        writer.writerow(row)

    def _save_record(self, writer, record: Record) -> None:
        '''
        Write a record to the CSV file.

        Args:
            writer (csv.writer): The CSV writer.
            record (Record): The record to write.
        '''
        row = [
            self.test_set_name,
            record.block_size,
            record.name,
            f'{record.psnr:0.4f}',
            f'{record.total_time:0.4f}',
            f'{record.coding_time:0.4f}',
            f'{record.coding_rate:0.4f}'
        ]

        writer.writerow(row)

    def load_from_file(self, file_path: Path) -> None:
        '''
        Load a reference result set from a CSV file on disk.

        Args:
            file_path: The input file path.
        '''
        with open(file_path, 'r', encoding='utf-8') as handle:
            reader = csv.reader(handle)

            # Skip the header
            next(reader)

            # Read the data rows
            for row in reader:
                assert row[0] == self.test_set_name

                record = Record(
                    row[1],
                    row[2],
                    float(row[3]),
                    float(row[4]),
                    float(row[5]),
                    float(row[6])
                )

                self.add_record(record)


class ResultSummary:
    '''
    A summary of a set of results.

    Attributes:
        not_runs: The number of tests that did not run.
        passes: The number of tests that passed.
        warnings: The number of tests that produced a warning.
        fails: The number of tests that failed.

        psnr: The image quality (PSNR dB)
        total_time: The total time.
        coding_time: The coding time.

        status: The test result status.
        psnr_rel: The relative image quality (diff PSNR dB, >= 0 good)
        total_time_rel: The relative total time vs ref (scale, <1 fast).
        coding_time_rel: The relative coding time vs ref (scale, <1 fast)..
    '''

    def __init__(self):
        '''
        Create a new result summary.
        '''
        # Pass fail metrics
        self.not_runs = 0
        self.passes = 0
        self.warnings = 0
        self.fails = 0

        # Absolute results
        self.psnr: list[float] = []
        self.total_time: list[float] = []
        self.coding_time: list[float] = []

        # Relative results
        self.psnr_rel: list[float] = []
        self.total_time_rel: list[float] = []
        self.coding_time_rel: list[float] = []

    def add_record(self, record: Record) -> None:
        '''
        Add a record to this summary.

        Args:
            record (Record): The Record to add.
        '''
        # Record pass/fail status
        if record.status == ResultStatus.PASS:
            self.passes += 1

        elif record.status == ResultStatus.WARN:
            self.warnings += 1

        elif record.status == ResultStatus.FAIL:
            self.fails += 1

        else:
            self.not_runs += 1

        # Store absolute results
        self.total_time.append(record.total_time)
        self.coding_time.append(record.coding_time)
        self.psnr.append(record.psnr)

        # Store relative results, if we have them
        if record.total_time_rel is not None:
            assert record.coding_time_rel is not None
            assert record.psnr_rel is not None

            self.total_time_rel.append(record.total_time_rel)
            self.coding_time_rel.append(record.coding_time_rel)
            self.psnr_rel.append(record.psnr_rel)

    def get_worst_result(self) -> ResultStatus:
        '''
        Get the worst result in this set.

        Return:
            The worst test result.
        '''
        if self.fails:
            return ResultStatus.FAIL

        if self.warnings:
            return ResultStatus.WARN

        if self.passes:
            return ResultStatus.PASS

        return ResultStatus.NOT_RUN

    def __str__(self) -> str:
        lines = ['']

        # Emit overall summary of the run
        overall = self.get_worst_result().name

        line = [
            f'Set Status: {overall}',
            f' (Pass: {self.passes}',
            f' | Warn: {self.warnings}',
            f' | Fail: {self.fails}',
        ]

        lines.append(''.join(line))

        # If no relative scores just return the summary
        if not self.total_time_rel:
            return '\n'.join(lines)

        lines.append('')

        def format_line(name, value, units) -> str:
            mean_value = numpy.mean(value)
            std_value = numpy.std(value)

            line = [
                f'{name:<15}',
                f'Mean: {mean_value:>+7.3f} {units:<3}',
                f' Std: {std_value:>+6.3f} {units:<3}'
            ]

            return ''.join(line)

        # Performance summaries
        lines.append(format_line('Total speed:', self.total_time_rel, 'x'))

        lines.append(format_line('Coding speed:', self.coding_time_rel, 'x'))

        lines.append(format_line('Quality diff:', self.psnr_rel, 'dB'))

        lines.append(format_line('Coding time:', self.coding_time, 's'))

        # Filter out 999 dB images from the results ...
        mod_psnr = [x for x in self.psnr if x < 999.0]
        lines.append(format_line('Quality:', mod_psnr, 'dB'))

        return '\n'.join(lines)