1"""
2Define a function with a field output: the viscous free fall example
3====================================================================
4"""
5# %%
6# In this example, we define a function which has a vector input and a field output. This is why we use the `PythonPointToFieldFunction` class to create the associated function and propagate the uncertainties through it.
7#
8# We consider a viscous free fall as explained :ref:`here <use-case-viscous-fall>`.
9
10# %%
11# Define the model
12# ----------------
13
14# %%
15from __future__ import print_function
16import openturns as ot
17import openturns.viewer as viewer
18from matplotlib import pylab as plt
19import numpy as np
20ot.Log.Show(ot.Log.NONE)
21
22# %%
23# We first define the time grid associated with the model.
24
25# %%
26tmin = 0.0  # Minimum time
27tmax = 12.  # Maximum time
28gridsize = 100  # Number of time steps
29mesh = ot.IntervalMesher([gridsize-1]).build(ot.Interval(tmin, tmax))
30
31# %%
32# The `getVertices` method returns the time values in this mesh.
33
34# %%
35vertices = mesh.getVertices()
36vertices[0:5]
37
38# %%
39# Creation of the input distribution.
40
41# %%
42distZ0 = ot.Uniform(100.0, 150.0)
43distV0 = ot.Normal(55.0, 10.0)
44distM = ot.Normal(80.0, 8.0)
45distC = ot.Uniform(0.0, 30.0)
46distribution = ot.ComposedDistribution([distZ0, distV0, distM, distC])
47
48# %%
49dimension = distribution.getDimension()
50dimension
51
52
53# %%
54# Then we define the Python function which computes the altitude at each time value. In order to compute all altitudes with a vectorized evaluation, we first convert the vertices into a `numpy` `array` and use the `numpy` function `exp` and `maximum`: this increases the evaluation performance of the script.
55
56# %%
57def AltiFunc(X):
58    g = 9.81
59    z0 = X[0]
60    v0 = X[1]
61    m = X[2]
62    c = X[3]
63    tau = m / c
64    vinf = - m * g / c
65    t = np.array(vertices)
66    z = z0 + vinf * t + tau * (v0 - vinf) * (1 - np.exp(- t / tau))
67    z = np.maximum(z, 0.)
68    return [[zeta[0]] for zeta in z]
69
70
71# %%
72# In order to create a `Function` from this Python function, we use the `PythonPointToFieldFunction` class. Since the altitude is the only output field, the third argument `outputDimension` is equal to `1`. If we had computed the speed as an extra output field, we would have set `2` instead.
73
74# %%
75outputDimension = 1
76alti = ot.PythonPointToFieldFunction(
77    dimension, mesh, outputDimension, AltiFunc)
78
79# %%
80# Sample trajectories
81# -------------------
82
83# %%
84# In order to sample trajectories, we use the `getSample` method of the input distribution and apply the field function.
85
86# %%
87size = 10
88inputSample = distribution.getSample(size)
89outputSample = alti(inputSample)
90
91# %%
92ot.ResourceMap.SetAsUnsignedInteger('Drawable-DefaultPalettePhase', size)
93
94# %%
95# Draw some curves.
96
97# %%
98graph = outputSample.drawMarginal(0)
99graph.setTitle('Viscous free fall: %d trajectories' % (size))
100graph.setXTitle(r'$t$')
101graph.setYTitle(r'$z$')
102view = viewer.View(graph)
103plt.show()
104# %%
105# We see that the object first moves up and then falls down. Not all objects, however, achieve the same maximum altitude. We see that some trajectories reach a higher maximum altitude than others. Moreover, at the final time :math:`t_{max}`, one trajectory hits the ground: :math:`z(t_{max})=0` for this trajectory.
106