use crate::algorithms::offline::graph_search::{Path, Paths};
use crate::algorithms::offline::OfflineOptions;
use crate::config::{Config, IntegralConfig};
use crate::cost::{Cost, CostFn, SingleCostFn};
use crate::model::{ModelOutputFailure, ModelOutputSuccess};
use crate::problem::{
IntegralSimplifiedSmoothedConvexOptimization, Problem,
SimplifiedSmoothedConvexOptimization,
};
use crate::result::{Failure, Result};
use crate::schedule::{IntegralSchedule, Schedule};
use crate::utils::{assert, is_pow_of_2};
use noisy_float::prelude::*;
use pyo3::prelude::*;
use serde_derive::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
pub struct Vertice(i32, i32);
#[pyclass(name = "OptimalGraphSearch1dOptions")]
#[derive(Clone)]
pub struct Options {
#[pyo3(get, set)]
pub x_start: i32,
}
impl Default for Options {
fn default() -> Self {
Options { x_start: 0 }
}
}
impl Options {
pub fn new(x_start: i32) -> Self {
Options { x_start }
}
}
#[pymethods]
impl Options {
#[new]
fn constructor(x_start: i32) -> Self {
Options { x_start }
}
}
pub fn optimal_graph_search<'a, C, D>(
mut p: IntegralSimplifiedSmoothedConvexOptimization<'a, C, D>,
Options { x_start }: Options,
OfflineOptions { inverted, alpha, l }: OfflineOptions,
) -> Result<Path>
where
C: ModelOutputSuccess + 'a,
D: ModelOutputFailure + 'a,
{
assert(l.is_none(), Failure::UnsupportedLConstrainedMovement)?;
assert(p.d == 1, Failure::UnsupportedProblemDimension(p.d))?;
if !is_pow_of_2(p.bounds[0]) {
p = make_pow_of_2(p)?;
}
let k_init = if p.bounds[0] > 2 {
(p.bounds[0] as f64).log(2.) as u32 - 2
} else {
0
};
let mut path =
find_schedule(&p, select_initial_rows(&p), alpha, inverted, x_start);
for k in (0..k_init).rev() {
path = find_schedule(
&p,
select_next_rows(&p, &path.xs, k),
alpha,
inverted,
x_start,
);
}
Ok(path)
}
pub fn make_pow_of_2<'a, C, D>(
p: IntegralSimplifiedSmoothedConvexOptimization<'a, C, D>,
) -> Result<IntegralSimplifiedSmoothedConvexOptimization<C, D>>
where
C: ModelOutputSuccess + 'a,
D: ModelOutputFailure + 'a,
{
let prev_m = p.bounds[0];
let m = 2_i32.pow((prev_m as f64).log(2.).ceil() as u32);
Ok(SimplifiedSmoothedConvexOptimization {
d: p.d,
t_end: p.t_end,
bounds: vec![m],
switching_cost: p.switching_cost.clone(),
hitting_cost: CostFn::new(
1,
SingleCostFn::certain(move |t, x: IntegralConfig| {
if x[0] <= prev_m {
p.hit_cost(t, x)
} else {
let Cost { cost, output } =
p.hit_cost(t, Config::new(vec![prev_m]));
Cost::new(n64(x[0] as f64) * (cost + f64::EPSILON), output)
}
}),
),
})
}
fn select_initial_rows<'a, C, D>(
p: &'a IntegralSimplifiedSmoothedConvexOptimization<'a, C, D>,
) -> impl Fn(i32) -> Vec<i32> + 'a {
move |_| (0..=4).map(|e| e * p.bounds[0] / 4).collect()
}
fn select_next_rows<'a, C, D>(
p: &'a IntegralSimplifiedSmoothedConvexOptimization<'a, C, D>,
xs: &'a IntegralSchedule,
k: u32,
) -> impl Fn(i32) -> Vec<i32> + 'a {
move |t| {
(-2..=2)
.map(|e| xs[t as usize - 1][0] + e * 2_i32.pow(k))
.filter(|&j| 0 <= j && j <= p.bounds[0])
.collect()
}
}
fn find_schedule<C, D>(
p: &IntegralSimplifiedSmoothedConvexOptimization<'_, C, D>,
select_rows: impl Fn(i32) -> Vec<i32>,
alpha: f64,
inverted: bool,
x_start: i32,
) -> Path
where
C: ModelOutputSuccess,
D: ModelOutputFailure,
{
let mut paths: Paths<Vertice> = HashMap::new();
let initial_vertice = Vertice(0, x_start);
let initial_path = Path {
xs: Schedule::empty(),
cost: 0.,
};
paths.insert(initial_vertice, initial_path);
let mut prev_rows = vec![x_start];
for t in 1..=p.t_end {
let rows = select_rows(t);
for &j in &rows {
find_shortest_subpath(
p, &mut paths, t, &prev_rows, j, alpha, inverted,
);
}
prev_rows = rows;
}
prev_rows.iter().fold(
Path {
xs: Schedule::empty(),
cost: f64::INFINITY,
},
|result, &i| {
let path = &paths[&Vertice(p.t_end, i)];
let cost = alpha
* p.movement(Config::single(i), Config::single(0), inverted)
.raw();
let picked_cost = path.cost + cost;
if picked_cost < result.cost {
Path {
xs: path.xs.clone(),
cost: picked_cost,
}
} else {
result
}
},
)
}
fn find_shortest_subpath<C, D>(
p: &IntegralSimplifiedSmoothedConvexOptimization<'_, C, D>,
paths: &mut Paths<Vertice>,
t: i32,
from: &Vec<i32>,
to: i32,
alpha: f64,
inverted: bool,
) where
C: ModelOutputSuccess,
D: ModelOutputFailure,
{
let mut picked_source = from[0];
let mut picked_cost = f64::INFINITY;
for &source in from {
let prev_cost = paths[&Vertice(t - 1, source)].cost;
let cost = build_cost(p, t, source, to, alpha, inverted);
let new_cost = prev_cost + cost;
if new_cost < picked_cost {
picked_source = source;
picked_cost = new_cost;
};
}
update_paths(paths, t, picked_source, to, picked_cost);
}
fn build_cost<C, D>(
p: &IntegralSimplifiedSmoothedConvexOptimization<'_, C, D>,
t: i32,
i: i32,
j: i32,
alpha: f64,
inverted: bool,
) -> f64
where
C: ModelOutputSuccess,
D: ModelOutputFailure,
{
let hitting_cost = p.hit_cost(t, Config::single(j)).cost.raw();
let movement = p
.movement(Config::single(i), Config::single(j), inverted)
.raw();
let switching_cost = alpha * movement;
hitting_cost + switching_cost
}
fn update_paths(paths: &mut Paths<Vertice>, t: i32, i: i32, j: i32, cost: f64) {
let u = Vertice(t - 1, i);
let v = Vertice(t, j);
let prev_xs = &paths[&u].xs;
let xs = prev_xs.extend(Config::single(j));
paths.insert(v, Path { xs, cost });
}