use crate::{Evaluator, Genotype};
use std::marker::PhantomData;
pub trait NumLike:
Copy
+ Send
+ Sync
+ PartialOrd
+ std::ops::Add<Output = Self>
+ std::ops::Sub<Output = Self>
+ std::ops::Mul<Output = Self>
+ std::ops::Div<Output = Self>
+ std::fmt::Debug
{
const ZERO: Self;
const ONE: Self;
fn from_f32(v: f32) -> Self;
fn to_f32(self) -> f32;
fn clamp_to(self, lo: Self, hi: Self) -> Self {
if self < lo {
lo
} else if self > hi {
hi
} else {
self
}
}
}
impl NumLike for f32 {
const ZERO: Self = 0.0;
const ONE: Self = 1.0;
fn from_f32(v: f32) -> Self {
v
}
fn to_f32(self) -> f32 {
self
}
}
impl NumLike for f64 {
const ZERO: Self = 0.0;
const ONE: Self = 1.0;
fn from_f32(v: f32) -> Self {
v as f64
}
fn to_f32(self) -> f32 {
self as f32
}
}
pub trait Scorer<S, O: NumLike>: Send + Sync {
fn score(&self, state: &S) -> O;
}
#[derive(Clone, Debug, Default)]
pub struct Trajectory {
pub start: [f32; 3],
pub end: [f32; 3],
pub up: [f32; 3],
pub max_height: f32,
pub energy_used: f32,
pub final_descriptor: Vec<f32>,
}
#[derive(Clone, Copy, Debug, Default)]
pub struct Displacement;
impl Scorer<Trajectory, f32> for Displacement {
fn score(&self, t: &Trajectory) -> f32 {
let dx = t.end[0] - t.start[0];
let dy = t.end[1] - t.start[1];
let dz = t.end[2] - t.start[2];
(dx * dx + dy * dy + dz * dz).sqrt()
}
}
#[derive(Clone, Copy, Debug, Default)]
pub struct UpAlignment;
impl Scorer<Trajectory, f32> for UpAlignment {
fn score(&self, t: &Trajectory) -> f32 {
t.up[1].clamp(0.0, 1.0)
}
}
#[derive(Clone, Copy, Debug, Default)]
pub struct Height;
impl Scorer<Trajectory, f32> for Height {
fn score(&self, t: &Trajectory) -> f32 {
t.max_height
}
}
#[derive(Clone, Copy, Debug, Default)]
pub struct EnergyEfficiency;
impl Scorer<Trajectory, f32> for EnergyEfficiency {
fn score(&self, t: &Trajectory) -> f32 {
if t.energy_used > 0.0 {
Displacement.score(t) / t.energy_used
} else {
0.0
}
}
}
#[derive(Clone, Copy, Debug)]
pub struct Const<O: NumLike>(pub O);
impl<S, O: NumLike> Scorer<S, O> for Const<O> {
fn score(&self, _state: &S) -> O {
self.0
}
}
#[derive(Clone, Copy, Debug)]
pub struct Multiply<A, B>(pub A, pub B);
impl<S, O, A, B> Scorer<S, O> for Multiply<A, B>
where
O: NumLike,
A: Scorer<S, O>,
B: Scorer<S, O>,
{
fn score(&self, state: &S) -> O {
self.0.score(state) * self.1.score(state)
}
}
#[derive(Clone, Copy, Debug)]
pub struct Sum<A, B>(pub A, pub B);
impl<S, O, A, B> Scorer<S, O> for Sum<A, B>
where
O: NumLike,
A: Scorer<S, O>,
B: Scorer<S, O>,
{
fn score(&self, state: &S) -> O {
self.0.score(state) + self.1.score(state)
}
}
#[derive(Clone, Copy, Debug)]
pub struct Penalize<A, B> {
pub base: A,
pub penalty: B,
}
impl<S, O, A, B> Scorer<S, O> for Penalize<A, B>
where
O: NumLike,
A: Scorer<S, O>,
B: Scorer<S, O>,
{
fn score(&self, state: &S) -> O {
self.base.score(state) - self.penalty.score(state)
}
}
#[derive(Clone, Copy, Debug)]
pub struct Normalize<A, O: NumLike> {
pub inner: A,
pub max: O,
}
impl<S, O, A> Scorer<S, O> for Normalize<A, O>
where
O: NumLike,
A: Scorer<S, O>,
{
fn score(&self, state: &S) -> O {
assert!(
self.max > O::ZERO,
"Normalize::max must be > 0 (division by zero otherwise)"
);
let raw = self.inner.score(state) / self.max;
raw.clamp_to(O::ZERO, O::ONE)
}
}
pub struct CompositeEvaluator<G, S, F>
where
G: Genotype,
S: Send + Sync,
F: Fn(&G) -> S + Send + Sync,
{
sim: F,
fitness_scorer: Box<dyn Scorer<S, f32> + Send + Sync>,
objective_scorers: Vec<Box<dyn Scorer<S, f32> + Send + Sync>>,
descriptor_scorers: Vec<Box<dyn Scorer<S, f32> + Send + Sync>>,
_marker: PhantomData<G>,
}
impl<G, S, F> CompositeEvaluator<G, S, F>
where
G: Genotype,
S: Send + Sync,
F: Fn(&G) -> S + Send + Sync,
{
pub fn new(sim: F, fitness_scorer: Box<dyn Scorer<S, f32> + Send + Sync>) -> Self {
Self {
sim,
fitness_scorer,
objective_scorers: Vec::new(),
descriptor_scorers: Vec::new(),
_marker: PhantomData,
}
}
pub fn with_objectives(mut self, scorers: Vec<Box<dyn Scorer<S, f32> + Send + Sync>>) -> Self {
self.objective_scorers = scorers;
self
}
pub fn with_descriptors(mut self, scorers: Vec<Box<dyn Scorer<S, f32> + Send + Sync>>) -> Self {
self.descriptor_scorers = scorers;
self
}
}
impl<G, S, F> Evaluator<G> for CompositeEvaluator<G, S, F>
where
G: Genotype,
S: Send + Sync,
F: Fn(&G) -> S + Send + Sync,
{
fn evaluate(&self, genotype: &G) -> (f32, Vec<f32>, Vec<f32>) {
let state = (self.sim)(genotype);
let fitness = self.fitness_scorer.score(&state);
let objectives = self
.objective_scorers
.iter()
.map(|s| s.score(&state))
.collect();
let descriptor = self
.descriptor_scorers
.iter()
.map(|s| s.score(&state))
.collect();
(fitness, objectives, descriptor)
}
}