mini-ode 0.1.5

A minimalistic ODE solvers library built on top of PyTorch
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84

use crate::compute_jacobian;
use tch::{Device, IndexOp, Kind, Tensor};

#[test]
fn test_compute_jacobian_error_non_1d_input() {
    let x = Tensor::zeros(&[2, 2], (Kind::Float, Device::Cpu));

    let result = compute_jacobian(|_| Tensor::ones(&[1], (Kind::Float, Device::Cpu)), &x);

    assert!(result.is_err());

    let err = result.unwrap_err().to_string();
    assert!(
        err.contains("one-dimensional"),
        "Unexpected error message: {}",
        err
    );
}

#[test]
fn test_compute_jacobian_error_non_1d_output() {
    let x = Tensor::ones(&[1], (Kind::Float, Device::Cpu));

    let result = compute_jacobian(|_| Tensor::ones(&[2, 2], (Kind::Float, Device::Cpu)), &x);

    assert!(result.is_err());

    let err = result.unwrap_err().to_string();
    assert!(
        err.contains("one-dimensional"),
        "Unexpected error message: {}",
        err
    );
}

#[test]
fn test_compute_jacobian_linear() {
    let x = Tensor::from_slice(&[1.0f32, 2.0f32]);

    let jacobian = compute_jacobian(
        |y| {
            let a = Tensor::from_slice(&[2.0f32, 0.0f32, 0.0f32, 3.0f32]).reshape(&[2, 2]);

            a.matmul(y)
        },
        &x,
    )
    .unwrap();

    let expected = Tensor::from_slice(&[2.0f32, 0.0f32, 0.0f32, 3.0f32]).reshape(&[2, 2]);

    assert_eq!(jacobian, expected);
}

#[test]
fn test_compute_jacobian_nonlinear_scalar() {
    let x = Tensor::from_slice(&[2.0f32]);

    let jacobian = compute_jacobian(|y| y.pow_tensor_scalar(2), &x).unwrap();

    let expected = Tensor::from_slice(&[4.0f32]).reshape(&[1, 1]);

    assert_eq!(jacobian, expected);
}

#[test]
fn test_compute_jacobian_multi_dim_nonlinear() {
    let x = Tensor::from_slice(&[1.0f32, 2.0f32]);

    let jacobian = compute_jacobian(
        |y| {
            let y0 = y.i(0);
            let y1 = y.i(1);

            Tensor::stack(&[&(&y0.pow_tensor_scalar(2) + &y1), &(&y0 * &y1)], 0)
        },
        &x,
    )
    .unwrap();

    let expected = Tensor::from_slice(&[2.0f32, 1.0f32, 2.0f32, 1.0f32]).reshape(&[2, 2]);

    assert_eq!(jacobian, expected);
}