use crate::InterpolationMode;
use crate::copp::copp3::Topp3ProfileRef;
use crate::ffi::python::array::array_like_to_vec_f64;
use crate::ffi::python::error::to_py_err;
use numpy::{IntoPyArray, PyArray1};
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
pub(crate) fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PyProfile3rd>()?;
m.add_function(wrap_pyfunction!(a_to_b_topp2, m)?)?;
m.add_function(wrap_pyfunction!(s_to_t_topp2, m)?)?;
m.add_function(wrap_pyfunction!(t_to_s_topp2_uniform, m)?)?;
m.add_function(wrap_pyfunction!(t_to_s_topp2_samples, m)?)?;
m.add_function(wrap_pyfunction!(t_to_s_topp2, m)?)?;
m.add_function(wrap_pyfunction!(s_to_t_topp3, m)?)?;
m.add_function(wrap_pyfunction!(t_to_s_topp3_uniform, m)?)?;
m.add_function(wrap_pyfunction!(t_to_s_topp3_samples, m)?)?;
m.add_function(wrap_pyfunction!(t_to_s_topp3, m)?)?;
Ok(())
}
#[pyclass(name = "Profile3rd", module = "copp_py._native", skip_from_py_object)]
#[derive(Clone)]
pub(crate) struct PyProfile3rd {
a: Vec<f64>,
b: Vec<f64>,
num_stationary: (usize, usize),
}
#[pymethods]
impl PyProfile3rd {
#[new]
#[pyo3(
signature = (a, b, num_stationary = (0, 0)),
text_signature = "(a, b, num_stationary=(0, 0))"
)]
fn new(
a: &Bound<'_, PyAny>,
b: &Bound<'_, PyAny>,
num_stationary: (usize, usize),
) -> PyResult<Self> {
let a = array_like_to_vec_f64("a", a)?;
let b = array_like_to_vec_f64("b", b)?;
if a.len() != b.len() {
return Err(PyValueError::new_err(format!(
"`a` and `b` must have the same length, got {} and {}",
a.len(),
b.len()
)));
}
Ok(Self {
a,
b,
num_stationary,
})
}
#[getter]
fn a<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1<f64>> {
PyArray1::from_vec(py, self.a.clone())
}
#[getter]
fn b<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1<f64>> {
PyArray1::from_vec(py, self.b.clone())
}
#[getter]
fn num_stationary(&self) -> (usize, usize) {
self.num_stationary
}
#[getter]
fn num_stationary_start(&self) -> usize {
self.num_stationary.0
}
#[getter]
fn num_stationary_end(&self) -> usize {
self.num_stationary.1
}
#[getter]
fn len(&self) -> usize {
self.a.len()
}
fn __bool__(&self) -> bool {
!self.a.is_empty()
}
fn __len__(&self) -> usize {
self.a.len()
}
fn __repr__(&self) -> String {
format!(
"Profile3rd(len={}, num_stationary={:?})",
self.a.len(),
self.num_stationary
)
}
}
impl PyProfile3rd {
pub(crate) fn as_parts(&self) -> Topp3ProfileRef<'_> {
(&self.a, &self.b, self.num_stationary)
}
pub(crate) fn from_rust(profile: crate::copp::copp3::Topp3Profile) -> Self {
let (a, b, num_stationary) = profile.into_parts();
Self {
a,
b,
num_stationary,
}
}
}
#[pyfunction]
fn a_to_b_topp2<'py>(
py: Python<'py>,
s: &Bound<'py, PyAny>,
a: &Bound<'py, PyAny>,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
let s = array_like_to_vec_f64("s", s)?;
let a = array_like_to_vec_f64("a", a)?;
let b = crate::solver::topp2_ra::a_to_b_topp2(&s, &a).map_err(to_py_err)?;
Ok(b.into_pyarray(py))
}
#[pyfunction]
#[pyo3(signature = (s, a, t0 = 0.0))]
fn s_to_t_topp2<'py>(
py: Python<'py>,
s: &Bound<'py, PyAny>,
a: &Bound<'py, PyAny>,
t0: f64,
) -> PyResult<(f64, Bound<'py, PyArray1<f64>>)> {
let s = array_like_to_vec_f64("s", s)?;
let a = array_like_to_vec_f64("a", a)?;
let (t_final, t_s) = crate::solver::topp2_ra::s_to_t_topp2(&s, &a, t0).map_err(to_py_err)?;
Ok((t_final, t_s.into_pyarray(py)))
}
#[pyfunction]
#[pyo3(signature = (s, a, t_s, dt, *, t0 = 0.0, include_final = true))]
fn t_to_s_topp2_uniform<'py>(
py: Python<'py>,
s: &Bound<'py, PyAny>,
a: &Bound<'py, PyAny>,
t_s: &Bound<'py, PyAny>,
dt: f64,
t0: f64,
include_final: bool,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
let s = array_like_to_vec_f64("s", s)?;
let a = array_like_to_vec_f64("a", a)?;
let t_s = array_like_to_vec_f64("t_s", t_s)?;
let s_t = t_to_s_topp2_uniform_impl(&s, &a, &t_s, t0, dt, include_final)?;
Ok(s_t.into_pyarray(py))
}
#[pyfunction]
#[pyo3(signature = (s, a, t_s, t_sample))]
fn t_to_s_topp2_samples<'py>(
py: Python<'py>,
s: &Bound<'py, PyAny>,
a: &Bound<'py, PyAny>,
t_s: &Bound<'py, PyAny>,
t_sample: &Bound<'py, PyAny>,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
let s = array_like_to_vec_f64("s", s)?;
let a = array_like_to_vec_f64("a", a)?;
let t_s = array_like_to_vec_f64("t_s", t_s)?;
let t_sample = array_like_to_vec_f64("t_sample", t_sample)?;
let s_t = t_to_s_topp2_samples_impl(&s, &a, &t_s, &t_sample)?;
Ok(s_t.into_pyarray(py))
}
#[pyfunction]
#[pyo3(signature = (s, a, t_s, *, t0 = 0.0, dt = None, include_final = true, t_sample = None))]
fn t_to_s_topp2<'py>(
py: Python<'py>,
s: &Bound<'py, PyAny>,
a: &Bound<'py, PyAny>,
t_s: &Bound<'py, PyAny>,
t0: f64,
dt: Option<f64>,
include_final: bool,
t_sample: Option<&Bound<'py, PyAny>>,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
let s = array_like_to_vec_f64("s", s)?;
let a = array_like_to_vec_f64("a", a)?;
let t_s = array_like_to_vec_f64("t_s", t_s)?;
let s_t = match (dt, t_sample) {
(Some(dt), None) => t_to_s_topp2_uniform_impl(&s, &a, &t_s, t0, dt, include_final)?,
(None, Some(t_sample)) => {
let t_sample = array_like_to_vec_f64("t_sample", t_sample)?;
t_to_s_topp2_samples_impl(&s, &a, &t_s, &t_sample)?
}
(Some(_), Some(_)) => {
return Err(PyValueError::new_err(
"pass either `dt` for a uniform time grid or `t_sample` for a nonuniform grid, not both",
));
}
(None, None) => {
return Err(PyValueError::new_err(
"`dt` is required unless `t_sample` is provided",
));
}
};
Ok(s_t.into_pyarray(py))
}
fn t_to_s_topp2_uniform_impl(
s: &[f64],
a: &[f64],
t_s: &[f64],
t0: f64,
dt: f64,
include_final: bool,
) -> PyResult<Vec<f64>> {
crate::solver::topp2_ra::t_to_s_topp2(
s,
a,
t_s,
InterpolationMode::UniformTimeGrid(t0, dt, include_final),
)
.map_err(to_py_err)
}
fn t_to_s_topp2_samples_impl(
s: &[f64],
a: &[f64],
t_s: &[f64],
t_sample: &[f64],
) -> PyResult<Vec<f64>> {
crate::solver::topp2_ra::t_to_s_topp2(
s,
a,
t_s,
InterpolationMode::NonUniformTimeGrid(t_sample),
)
.map_err(to_py_err)
}
#[pyfunction]
#[pyo3(signature = (s, profile, t0 = 0.0))]
fn s_to_t_topp3<'py>(
py: Python<'py>,
s: &Bound<'py, PyAny>,
profile: PyRef<'py, PyProfile3rd>,
t0: f64,
) -> PyResult<(f64, Bound<'py, PyArray1<f64>>)> {
let s = array_like_to_vec_f64("s", s)?;
let (t_final, t_s) =
crate::solver::topp3_lp::s_to_t_topp3(&s, profile.as_parts(), t0).map_err(to_py_err)?;
Ok((t_final, t_s.into_pyarray(py)))
}
#[pyfunction]
#[pyo3(signature = (s, profile, t_s, dt, *, t0 = 0.0, include_final = true))]
fn t_to_s_topp3_uniform<'py>(
py: Python<'py>,
s: &Bound<'py, PyAny>,
profile: PyRef<'py, PyProfile3rd>,
t_s: &Bound<'py, PyAny>,
dt: f64,
t0: f64,
include_final: bool,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
let s = array_like_to_vec_f64("s", s)?;
let t_s = array_like_to_vec_f64("t_s", t_s)?;
let s_t = t_to_s_topp3_uniform_impl(&s, profile.as_parts(), &t_s, t0, dt, include_final)?;
Ok(s_t.into_pyarray(py))
}
#[pyfunction]
#[pyo3(signature = (s, profile, t_s, t_sample))]
fn t_to_s_topp3_samples<'py>(
py: Python<'py>,
s: &Bound<'py, PyAny>,
profile: PyRef<'py, PyProfile3rd>,
t_s: &Bound<'py, PyAny>,
t_sample: &Bound<'py, PyAny>,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
let s = array_like_to_vec_f64("s", s)?;
let t_s = array_like_to_vec_f64("t_s", t_s)?;
let t_sample = array_like_to_vec_f64("t_sample", t_sample)?;
let s_t = t_to_s_topp3_samples_impl(&s, profile.as_parts(), &t_s, &t_sample)?;
Ok(s_t.into_pyarray(py))
}
#[pyfunction]
#[pyo3(signature = (s, profile, t_s, *, t0 = 0.0, dt = None, include_final = true, t_sample = None))]
fn t_to_s_topp3<'py>(
py: Python<'py>,
s: &Bound<'py, PyAny>,
profile: PyRef<'py, PyProfile3rd>,
t_s: &Bound<'py, PyAny>,
t0: f64,
dt: Option<f64>,
include_final: bool,
t_sample: Option<&Bound<'py, PyAny>>,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
let s = array_like_to_vec_f64("s", s)?;
let t_s = array_like_to_vec_f64("t_s", t_s)?;
let s_t = match (dt, t_sample) {
(Some(dt), None) => {
t_to_s_topp3_uniform_impl(&s, profile.as_parts(), &t_s, t0, dt, include_final)?
}
(None, Some(t_sample)) => {
let t_sample = array_like_to_vec_f64("t_sample", t_sample)?;
t_to_s_topp3_samples_impl(&s, profile.as_parts(), &t_s, &t_sample)?
}
(Some(_), Some(_)) => {
return Err(PyValueError::new_err(
"pass either `dt` for a uniform time grid or `t_sample` for a nonuniform grid, not both",
));
}
(None, None) => {
return Err(PyValueError::new_err(
"`dt` is required unless `t_sample` is provided",
));
}
};
Ok(s_t.into_pyarray(py))
}
fn t_to_s_topp3_uniform_impl(
s: &[f64],
profile: Topp3ProfileRef<'_>,
t_s: &[f64],
t0: f64,
dt: f64,
include_final: bool,
) -> PyResult<Vec<f64>> {
crate::solver::topp3_lp::t_to_s_topp3(
s,
profile,
t_s,
InterpolationMode::UniformTimeGrid(t0, dt, include_final),
)
.map_err(to_py_err)
}
fn t_to_s_topp3_samples_impl(
s: &[f64],
profile: Topp3ProfileRef<'_>,
t_s: &[f64],
t_sample: &[f64],
) -> PyResult<Vec<f64>> {
crate::solver::topp3_lp::t_to_s_topp3(
s,
profile,
t_s,
InterpolationMode::NonUniformTimeGrid(t_sample),
)
.map_err(to_py_err)
}