1use crate::errors::RocheError;
2use crate::{Star, Vec3};
3use pyo3::prelude::*;
4use std::f64::consts::PI;
5
6#[pyfunction]
25pub fn rpot_val(
26 q: f64,
27 star: Star,
28 spin: f64,
29 earth: &Vec3,
30 p: &Vec3,
31 lam: f64,
32) -> Result<f64, RocheError> {
33 let r: Vec3 = *p + lam * *earth;
34 Ok(match star {
35 Star::Primary => rpot1(q, spin, &r)?,
36 Star::Secondary => rpot2(q, spin, &r)?,
37 })
38}
39
40#[pyfunction]
61pub fn rpot_val_grad(
62 q: f64,
63 star: Star,
64 spin: f64,
65 earth: &Vec3,
66 p: &Vec3,
67 lam: f64,
68) -> Result<(f64, f64, f64), RocheError> {
69 let r: Vec3 = *p + lam * *earth;
70 let d: Vec3 = match star {
71 Star::Primary => drpot1(q, spin, &r)?,
72 Star::Secondary => drpot2(q, spin, &r)?,
73 };
74 let rpot: f64 = match star {
75 Star::Primary => rpot1(q, spin, &r)?,
76 Star::Secondary => rpot2(q, spin, &r)?,
77 };
78 let ed: Vec3 = Vec3::new(earth.y, -earth.x, 0.);
79 let dphi: f64 = 2. * PI * lam * d.dot(&ed);
80 let dlam: f64 = d.dot(earth);
81
82 Ok((rpot, dphi, dlam))
83}
84
85#[pyfunction]
105pub fn rpot_grad(
106 q: f64,
107 star: Star,
108 spin: f64,
109 earth: &Vec3,
110 p: &Vec3,
111 lam: f64,
112) -> Result<(f64, f64), RocheError> {
113 let r: Vec3 = *p + lam * *earth;
114
115 let d: Vec3 = match star {
116 Star::Primary => drpot1(q, spin, &r)?,
117 Star::Secondary => drpot2(q, spin, &r)?,
118 };
119
120 let ed: Vec3 = Vec3::new(earth.y, -earth.x, 0.);
121
122 let dphi: f64 = 2. * PI * lam * d.dot(&ed);
123 let dlam: f64 = d.dot(earth);
124
125 Ok((dphi, dlam))
126}
127
128#[pyfunction]
141pub fn rpot(q: f64, p: &Vec3) -> Result<f64, RocheError> {
142 if q <= 0. {
143 let message = format!("q = {} <= 0", q);
144 return Err(RocheError::ParameterError(message));
145 }
146 let mu: f64 = q / (1. + q);
147 let comp: f64 = 1. - mu;
148 let x2y2: f64 = p.x * p.x + p.y * p.y;
149 let z2: f64 = p.z * p.z;
150 let r1sq: f64 = x2y2 + z2;
151 let r1: f64 = r1sq.sqrt();
152 let r2: f64 = (r1sq + 1. - 2. * p.x).sqrt();
153 Ok(-comp / r1 - mu / r2 - (x2y2 + mu * (mu - 2.0 * p.x)) / 2.0)
154}
155
156#[pyfunction]
170pub fn rpot1(q: f64, spin: f64, p: &Vec3) -> Result<f64, RocheError> {
171 if q <= 0. {
172 let message = format!("q = {} <= 0", q);
173 return Err(RocheError::ParameterError(message));
174 }
175 let mu: f64 = q / (1. + q);
176 let comp: f64 = 1. - mu;
177 let x2y2: f64 = p.x * p.x + p.y * p.y;
178 let z2: f64 = p.z * p.z;
179 let r1sq: f64 = x2y2 + z2;
180 let r1: f64 = r1sq.sqrt();
181 let r2: f64 = (r1sq + 1. - 2. * p.x).sqrt();
182 Ok(-comp / r1 - mu / r2 - spin * spin * x2y2 / 2.0 + mu * p.x)
183}
184
185#[pyfunction]
199pub fn rpot2(q: f64, spin: f64, p: &Vec3) -> Result<f64, RocheError> {
200 if q <= 0. {
201 let message = format!("q = {} <= 0", q);
202 return Err(RocheError::ParameterError(message));
203 }
204 let mu: f64 = q / (1. + q);
205 let comp: f64 = 1. - mu;
206 let x2y2: f64 = p.x * p.x + p.y * p.y;
207 let z2: f64 = p.z * p.z;
208 let r1sq: f64 = x2y2 + z2;
209 let r1: f64 = r1sq.sqrt();
210 let r2: f64 = (r1sq + 1. - 2. * p.x).sqrt();
211 Ok(-comp / r1 - mu / r2 - spin * spin * (0.5 + 0.5 * x2y2 - p.x) - comp * p.x)
212}
213
214#[pyfunction]
228pub fn drpot(q: f64, p: &Vec3) -> Result<Vec3, RocheError> {
229 if q <= 0. {
230 let message = format!("q = {} <= 0", q);
231 return Err(RocheError::ParameterError(message));
232 }
233 let r1sq: f64 = p.sqr();
234 let r1: f64 = r1sq.sqrt();
235 let r2sq: f64 = r1sq + 1. - 2. * p.x;
236 let r2: f64 = r2sq.sqrt();
237 let mu: f64 = q / (1. + q);
238 let mu1: f64 = mu / r2 / r2sq;
239 let comp: f64 = (1. - mu) / r1 / r1sq;
240 Ok(Vec3::new(
241 comp * p.x + mu1 * (p.x - 1.) - p.x + mu,
242 comp * p.y + mu1 * p.y - p.y,
243 comp * p.z + mu1 * p.z,
244 ))
245}
246
247#[pyfunction]
262pub fn drpot1(q: f64, spin: f64, p: &Vec3) -> Result<Vec3, RocheError> {
263 if q <= 0. {
264 let message = format!("q = {} <= 0", q);
265 return Err(RocheError::ParameterError(message));
266 }
267 let r1sq: f64 = p.sqr();
268 let r1: f64 = r1sq.sqrt();
269 let r2sq: f64 = r1sq + 1. - 2. * p.x;
270 let r2: f64 = r2sq.sqrt();
271 let mu: f64 = q / (1. + q);
272 let mu1: f64 = mu / r2 / r2sq;
273 let comp: f64 = (1. - mu) / r1 / r1sq;
274 let ssq: f64 = spin * spin;
275 Ok(Vec3::new(
276 comp * p.x + mu1 * (p.x - 1.) - ssq * p.x + mu,
277 comp * p.y + mu1 * p.y - ssq * p.y,
278 comp * p.z + mu1 * p.z,
279 ))
280}
281
282#[pyfunction]
297pub fn drpot2(q: f64, spin: f64, p: &Vec3) -> Result<Vec3, RocheError> {
298 if q <= 0. {
299 let message = format!("q = {} <= 0", q);
300 return Err(RocheError::ParameterError(message));
301 }
302 let r1sq: f64 = p.sqr();
303 let r1: f64 = r1sq.sqrt();
304 let r2sq: f64 = r1sq + 1. - 2. * p.x;
305 let r2: f64 = r2sq.sqrt();
306 let mu: f64 = q / (1. + q);
307 let mu1: f64 = mu / r2 / r2sq;
308 let comp: f64 = (1. - mu) / r1 / r1sq;
309 let ssq: f64 = spin * spin;
310 Ok(Vec3::new(
311 comp * p.x + mu1 * (p.x - 1.) - ssq * (p.x - 1.) + mu - 1.,
312 comp * p.y + mu1 * p.y - ssq * p.y,
313 comp * p.z + mu1 * p.z,
314 ))
315}
316
317#[cfg(test)]
318mod tests {
319 use super::*;
320
321 #[test]
322 fn rpot_test() -> Result<(), RocheError> {
323 let r = Vec3::new(0.3, 0.3, 0.0);
325 assert_eq!(rpot(0.2, &r)?, -2.2369184469510586);
326 assert!(rpot(-0.2, &r).is_err());
327 Ok(())
328 }
329
330 #[test]
331 fn rpot1_test() -> Result<(), RocheError> {
332 let r = Vec3::new(0.3, 0.3, 0.0);
334 assert_eq!(rpot1(0.2, 0.5, &r)?, -2.15552955806217);
335 assert!(rpot1(-0.2, 0.5, &r).is_err());
336 Ok(())
337 }
338
339 #[test]
340 fn rpot2_test() -> Result<(), RocheError> {
341 let r = Vec3::new(0.3, 0.3, 0.0);
343 assert_eq!(rpot2(0.2, 0.5, &r)?, -2.5055295580621695);
344 assert!(rpot2(-0.2, 0.5, &r).is_err());
345 Ok(())
346 }
347
348 #[test]
349 fn drpot_test() -> Result<(), RocheError> {
350 let r = Vec3::new(0.3, 0.3, 0.0);
352 assert_eq!(
353 drpot(0.2, &r)?,
354 Vec3::new(2.876187037097281, 3.0868377062344154, 0.0)
355 );
356 assert!(drpot(-0.2, &r).is_err());
357 Ok(())
358 }
359
360 #[test]
361 fn drpot1_test() -> Result<(), RocheError> {
362 let r = Vec3::new(0.3, 0.3, 0.0);
364 assert_eq!(
365 drpot1(0.2, 0.5, &r)?,
366 Vec3::new(3.1011870370972807, 3.311837706234415, 0.0)
367 );
368 assert!(drpot1(-0.2, 0.5, &r).is_err());
369 Ok(())
370 }
371
372 #[test]
373 fn drpot2_test() -> Result<(), RocheError> {
374 let r = Vec3::new(0.3, 0.3, 0.0);
376 assert_eq!(
377 drpot2(0.2, 0.5, &r)?,
378 Vec3::new(2.3511870370972807, 3.311837706234415, 0.0)
379 );
380 assert!(drpot2(-0.2, 0.5, &r).is_err());
381 Ok(())
382 }
383}