134 lines
3.9 KiB
Python
134 lines
3.9 KiB
Python
from abc import ABC, abstractmethod
|
||
from typing import Callable
|
||
|
||
import numpy as np
|
||
from numpy.typing import NDArray
|
||
|
||
from .chromosome import Chromosome
|
||
|
||
type FitnessFn = Callable[
|
||
[
|
||
Chromosome,
|
||
NDArray[np.float64],
|
||
Callable[[NDArray[np.float64]], NDArray[np.float64]],
|
||
],
|
||
float,
|
||
]
|
||
type TargetFunction = Callable[[NDArray[np.float64]], NDArray[np.float64]]
|
||
type TestPointsFn = Callable[[], NDArray[np.float64]]
|
||
|
||
|
||
class BaseFitness(ABC):
|
||
def __init__(self, target_fn: TargetFunction, test_points_fn: TestPointsFn):
|
||
self.target_function = target_fn
|
||
self.test_points_fn = test_points_fn
|
||
|
||
@abstractmethod
|
||
def fitness_fn(
|
||
self,
|
||
chromosome: Chromosome,
|
||
predicted: NDArray[np.float64],
|
||
true_values: NDArray[np.float64],
|
||
) -> float: ...
|
||
|
||
def __call__(self, chromosome: Chromosome) -> float:
|
||
test_points = self.test_points_fn()
|
||
|
||
context = {t: test_points[:, i] for i, t in enumerate(chromosome.terminals)}
|
||
predicted = chromosome.root.eval(context)
|
||
true_values = self.target_function(test_points)
|
||
|
||
return self.fitness_fn(chromosome, predicted, true_values)
|
||
|
||
|
||
class MSEFitness(BaseFitness):
|
||
"""Среднеквадратичная ошибка"""
|
||
|
||
def fitness_fn(
|
||
self,
|
||
chromosome: Chromosome,
|
||
predicted: NDArray[np.float64],
|
||
true_values: NDArray[np.float64],
|
||
) -> float:
|
||
return float(np.mean((predicted - true_values) ** 2))
|
||
|
||
|
||
class RMSEFitness(BaseFitness):
|
||
"""Корень из среднеквадратичной ошибки"""
|
||
|
||
def fitness_fn(
|
||
self,
|
||
chromosome: Chromosome,
|
||
predicted: NDArray[np.float64],
|
||
true_values: NDArray[np.float64],
|
||
) -> float:
|
||
return float(np.sqrt(np.mean((predicted - true_values) ** 2)))
|
||
|
||
|
||
class MAEFitness(BaseFitness):
|
||
"""Средняя абсолютная ошибка"""
|
||
|
||
def fitness_fn(
|
||
self,
|
||
chromosome: Chromosome,
|
||
predicted: NDArray[np.float64],
|
||
true_values: NDArray[np.float64],
|
||
) -> float:
|
||
return float(np.mean(np.abs(predicted - true_values)))
|
||
|
||
|
||
class NRMSEFitness(BaseFitness):
|
||
"""Нормализованный RMSE (масштаб-инвариантен)"""
|
||
|
||
def fitness_fn(
|
||
self,
|
||
chromosome: Chromosome,
|
||
predicted: NDArray[np.float64],
|
||
true_values: NDArray[np.float64],
|
||
) -> float:
|
||
denom = np.std(true_values)
|
||
if denom == 0:
|
||
return 1e6
|
||
return float(np.sqrt(np.mean((predicted - true_values) ** 2)) / denom)
|
||
|
||
|
||
class PenalizedFitness(BaseFitness):
|
||
"""Фитнес со штрафом за размер и глубину дерева: ошибка + λ * (размер + depth_weight * глубина)"""
|
||
|
||
def __init__(
|
||
self,
|
||
target_fn: TargetFunction,
|
||
test_points_fn: TestPointsFn,
|
||
base_fitness: BaseFitness,
|
||
lambda_: float = 0.001,
|
||
depth_weight: float = 0.2,
|
||
scale_penalty: bool | None = None,
|
||
):
|
||
super().__init__(target_fn, test_points_fn)
|
||
self.base_fitness = base_fitness
|
||
self.lambda_ = lambda_
|
||
self.depth_weight = depth_weight
|
||
|
||
# Масштабировать штраф необязательно, если функция фитнеса нормализована
|
||
if scale_penalty is None:
|
||
scale_penalty = not isinstance(base_fitness, NRMSEFitness)
|
||
self.scale_penalty = scale_penalty
|
||
|
||
def fitness_fn(
|
||
self,
|
||
chromosome: Chromosome,
|
||
predicted: NDArray[np.float64],
|
||
true_values: NDArray[np.float64],
|
||
) -> float:
|
||
base = self.base_fitness.fitness_fn(chromosome, predicted, true_values)
|
||
|
||
size = chromosome.root.get_size()
|
||
depth = chromosome.root.get_depth()
|
||
|
||
penalty = self.lambda_ * (size + self.depth_weight * depth)
|
||
|
||
if self.scale_penalty:
|
||
penalty *= base
|
||
|
||
return float(base + penalty)
|