Search code examples
pythonmatplotlibmatplotlib-3d

How to color 3D BSplines based on weights

Given a Dataframe with recorded acceleration values on 3-axis that are used to determine position vectors. After applying BSpline, there is a smooth curve which looks quite well.

As a next step, I want to show also the extent of acceleration using different colors. For that, Matplotlib line color in 3D gives a good base, although the result is somehow messy, and I cannot identify the root cause:

Here is the record of acceleration values in CSV:

,Counter,Acc.X,Acc.Y,Acc.Z
18,18,9.71,0.51,4.41
19,19,9.69,0.65,4.34
20,20,9.67,0.6,4.3
21,21,9.6,0.59,4.33
22,22,9.34,0.11,4.63
23,23,6.01,-2.75,8.43
24,24,0.58,-5.01,6.18
25,25,2.16,-2.19,0.8
26,26,7.42,-1.84,-4.85
27,27,12.56,-5.04,-8.95
28,28,14.78,-9.62,-8.07
29,29,10.82,-7.58,-5.8
30,30,11.05,-8.45,-3.82

And the code:

import pandas as pd
import numpy as np
from scipy.interpolate import splrep, BSpline
from mpl_toolkits import mplot3d
%matplotlib inline
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from mpl_toolkits.mplot3d.art3d import Line3DCollection

from enum import Enum
class axis(Enum):
    X = 1
    Y = 2
    Z = 3

records = pd.read_csv(filepath_or_buffer='input.csv', index_col=0, sep=',')

i = range(0, len(records.index)-1)
x = (records.iloc[1:].values-records.iloc[:-1]).loc[:,'Acc.X'].to_numpy()
y = (records.iloc[1:].values-records.iloc[:-1]).loc[:,'Acc.Y'].to_numpy()
z = (records.iloc[1:].values-records.iloc[:-1]).loc[:,'Acc.Z'].to_numpy()

ii = np.linspace(i.start, i.stop);

tckx0 = splrep(i, x, s=0)
tckx = splrep(i, x, s=50)
tcky = splrep(i, y, s=50)
tckz = splrep(i, z, s=50)

ax = plt.figure(figsize=(50, 50)).add_subplot(projection='3d')
ax.plot(BSpline(*tckx)(ii),BSpline(*tcky)(ii),BSpline(*tckz)(ii), 'o--', color='green')
plt.show()

It results the following figure: enter image description here

Applying colors, firstly just a simple style such as 'magma':

# Coloring
points = np.array([BSpline(*tckx)(ii),BSpline(*tcky)(ii),BSpline(*tckz)(ii)]).transpose().reshape(-1,1,3)
segs = np.concatenate([points[:-1],points[1:]],axis=1)

lc = Line3DCollection(segs, cmap=plt.get_cmap('magma'))
lc.set_array(ii)

fig = plt.figure(figsize=(50, 50))
ax = fig.add_subplot(111, projection='3d')
ax.add_collection3d(lc)
plt.show()

The figure is completely messed: enter image description here

I cannot identify the root cause, could you please help me?

Solution

  • I think you are interpolating the wrong things.

    To interpolate a 3D curve you have to use scipy's splprep (see this question).

    To interpolate a 1D curve you have to use scipy's splrep.

    In the following code I also plotted the initial x, y, z data so that you can compare your result with what I come up.

    from scipy.interpolate import splrep, BSpline, splev, splprep
from mpl_toolkits import mplot3d
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from mpl_toolkits.mplot3d.art3d import Line3DCollection
import numpy as np
from enum import Enum
class axis(Enum):
    X = 1
    Y = 2
    Z = 3

from io import StringIO
import pandas as pd

mystr = StringIO("""Counter,Acc.X,Acc.Y,Acc.Z
18,18,9.71,0.51,4.41
19,19,9.69,0.65,4.34
20,20,9.67,0.6,4.3
21,21,9.6,0.59,4.33
22,22,9.34,0.11,4.63
23,23,6.01,-2.75,8.43
24,24,0.58,-5.01,6.18
25,25,2.16,-2.19,0.8
26,26,7.42,-1.84,-4.85
27,27,12.56,-5.04,-8.95
28,28,14.78,-9.62,-8.07
29,29,10.82,-7.58,-5.8
30,30,11.05,-8.45,-3.82
""")
records = pd.read_csv(mystr, sep=",")

i = range(0, len(records.index)-1)
x = (records.iloc[1:].values-records.iloc[:-1]).loc[:,'Acc.X'].to_numpy()
y = (records.iloc[1:].values-records.iloc[:-1]).loc[:,'Acc.Y'].to_numpy()
z = (records.iloc[1:].values-records.iloc[:-1]).loc[:,'Acc.Z'].to_numpy()

ii = np.linspace(i.start, i.stop);

tckx0 = splrep(i, x, s=0)
tckx = splrep(i, x, s=50)
tcky = splrep(i, y, s=50)
tckz = splrep(i, z, s=50)

fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.plot(x, y, z, 'ko', label="true data")
ax.plot(BSpline(*tckx)(ii),BSpline(*tcky)(ii),BSpline(*tckz)(ii),
        color='green', label="your interpolation")

# spline interpolation of a 3D curve
tck, u = splprep([x, y, z], s=2)
ii = np.linspace(0, 1)
xs, ys, zs = splev(ii, tck)

# spline interpolation of a 1D curve
a_x = records["Acc.X"].to_numpy()[:-1]
a_y = records["Acc.Y"].to_numpy()[:-1]
a_z = records["Acc.Z"].to_numpy()[:-1]
# color by magnitude of acceleration
a = np.sqrt(a_x**2 + a_y**2 + a_z**2)
tcka = splrep(np.linspace(0, 1, len(a)), a, s=0)
parameter = splev(ii, tcka)

ax.plot(xs, ys, zs, 'r.', label="new interpolation")

def get_segments(x, y, z):
    points = np.ma.array([x, y, z]).T.reshape(-1, 1, 3)
    return np.ma.concatenate([points[:-1], points[1:]], axis=1)

line = Line3DCollection(get_segments(xs, ys, zs), cmap="magma", array=parameter)
ax.add_collection(line)
fig.colorbar(line, label="acceleration")

ax.legend()
plt.show()

    enter image description here