component-model/examples/bouncing_ball_3d_pythonfmu.py at b97c86f7acb231341e5aa69eb869cd05e7cf4cc4 · dnv-opensource/component-model · GitHub

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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
from math import sqrt
from typing import Any

import numpy as np
from pythonfmu import Fmi2Causality, Fmi2Slave, Real  # type: ignore[import-untyped]


class BouncingBall3D(Fmi2Slave):
    """Another Python-based BouncingBall model, using PythonFMU to construct a FMU.

    Features:

    * The ball has a 3-D vector as position and speed
    * As output variable the model estimates the next bouncing point
    * As input variables, the restitution coefficient `e`, the gravitational acceleration `g`
       and the initial speed can be changed.
    * Internal units are assumed as SI (m,s,rad)
    """

    def __init__(self, **kwargs: Any):
        super().__init__(
            name="BouncingBall3D",
            description="Another Python-based BouncingBall model, using Model and Variable to construct a FMU",
            author="DNV, SEACo project",
            **kwargs,
        )

        # Register variable ports for the model
        self.posX = 0.0
        self.register_variable(Real("posX", causality=Fmi2Causality.output))

        self.posY = 0.0
        self.register_variable(Real("posY", causality=Fmi2Causality.output))

        self.posZ = 10.0
        self.register_variable(Real("posZ", causality=Fmi2Causality.output))

        self.speedX = 0.0
        self.register_variable(Real("speedX", causality=Fmi2Causality.output))

        self.speedY = 0.0
        self.register_variable(Real("speedY", causality=Fmi2Causality.output))

        self.speedZ = 0.0
        self.register_variable(Real("speedZ", causality=Fmi2Causality.output))

        self.p_bounceX = -1.0
        self.register_variable(Real("p_bounceX", causality=Fmi2Causality.output))

        self.p_bounceY = -1.0
        self.register_variable(Real("p_bounceY", causality=Fmi2Causality.output))
        self.t_bounce = -1.0

        self.g = 9.81  # Gravitational acceleration
        self.register_variable(Real("g", causality=Fmi2Causality.parameter))

        self.e = 0.0  # Coefficient of restitution
        self.register_variable(Real("e", causality=Fmi2Causality.parameter))

        # Internal states
        self.stopped = False
        self.min_speed_z = 1e-6
        self.accelerationX = 0.0
        self.accelerationY = 0.0
        self.accelerationZ = -self.g

    def do_step(self, current_time: float, step_size: float) -> bool:
        """Perform a simulation step from `self.time` to `self.time + step_size`.

        With respect to bouncing (self.t_bounce should be initialized to a negative value)
        .t_bounce <= .time: update .t_bounce
        .time < .t_bounce <= .time+step_size: bouncing happens within time step
        .t_bounce > .time+step_size: no bouncing. Just advance .pos and .speed
        """
        if self.t_bounce < self.time:  # calculate first bounce
            self.t_bounce, self.p_bounce = self.next_bounce()
        while self.t_bounce <= self.time + step_size:  # bounce happens within step or at border
            dt1 = self.t_bounce - self.time
            self.posX = self.p_bounceX
            self.posY = self.p_bounceY
            self.speedZ += self.accelerationZ * dt1  # speed before bouncing
            # speed reduction due to coefficient of restitution
            self.speedX *= self.e
            self.speedY *= self.e
            self.speedZ *= -self.e  # change also direction
            if self.speedZ < self.min_speed_z:
                self.stopped = True
                self.accelerationZ = 0.0
                self.speedZ = 0.0
                self.posZ = 0.0
            self.time += dt1  # jump to the exact bounce time
            step_size -= dt1
            self.t_bounce, self.p_bounce = self.next_bounce()  # update to the next bounce
            self.p_bounceX, self.p_bounceY = self.p_bounce[0:2]
        if step_size > 0:
            self.posX += self.speedX * step_size
            self.posY += self.speedY * step_size
            self.posZ += self.speedZ * step_size + 0.5 * self.accelerationZ * step_size**2
            self.speedZ += self.accelerationZ * step_size
            self.time += step_size
        if self.posZ < 0:  # avoid numeric issues
            self.posZ = 0
        return True

    def next_bounce(self) -> tuple[float, np.ndarray]:
        """Calculate time of next bounce and position where the ground will be hit,
        based on .time, .pos and .speed.
        """
        if self.stopped:  # stopped bouncing
            return (1e300, np.array((1e300, 1e300, 0), float))
        else:
            dt_bounce = (self.speedZ + sqrt(self.speedZ**2 + 2 * self.g * self.posZ)) / self.g
            p_bounceX = self.posX + self.speedX * dt_bounce  # linear in x and y!
            p_bounceY = self.posY + self.speedY * dt_bounce
            return (self.time + dt_bounce, np.array((p_bounceX, p_bounceY, 0), float))

    def setup_experiment(self, start_time: float, stop_time: float | None = None, tolerance: float | None = None):
        """Set initial (non-interface) variables."""
        super().setup_experiment(start_time, stop_time, tolerance)
        # print(f"SETUP_EXPERIMENT g={self.g}, e={self.e}")
        self.stopped = False
        self.time = start_time

    def exit_initialization_mode(self):
        """Initialize the model after initial variables are set."""
        super().exit_initialization_mode()
        self.accelerationZ = -self.g