Fitting taylor expansion to particle trajectories

Imports

from IPython.display import display, HTML, clear_output
import base64

import tensorflow as tf
gpu = tf.config.get_visible_devices("GPU")[0]
tf.config.set_visible_devices(gpu, device_type="GPU")
tf.config.experimental.set_memory_growth(gpu, True)

import imageio
import pyvista
import numpy as np
from tqdm.notebook import tqdm
import matplotlib.pyplot as plt
from pint import Quantity

import sys
sys.path.insert(0, "../../")
sys.path.append("..")
import cmrsim
import local_functions

Simple example

Construct trajectories

n_p = 50
n_t = 50

r0 = Quantity(np.random.normal(0., 0.001, size=(n_p, 3)), "m")
t = Quantity(np.linspace(0, 100, n_t), "ms")
v = Quantity(np.random.uniform(1, 2, size=(n_p, 3)), "cm/s")
a = Quantity(np.random.uniform(1, 2, size=(n_p, 3)), "cm/s**2")

r = a.reshape(n_p, 1, 3) * t.reshape(1, n_t, 1)**2 +  v.reshape(n_p, 1, 3) * t.reshape(1, n_t, 1) + r0.reshape(n_p, 1, 3)

Fit trajectory Module

fig, ax = plt.subplots(1, 3, figsize=(20, 5))

for order in range(2, 5):
    trajectory_module = cmrsim.trajectory.TaylorTrajectoryN(order=order, time_grid=t.m_as("ms"), particle_trajectories=r.m_as("m"))
    taylor_positions, _ = trajectory_module(t.m_as("ms").astype(np.float32))
    m1 = np.mean(taylor_positions - r.m_as("m"), axis=0)
    s1 = np.std(taylor_positions - r.m_as("m"), axis=0)
    for i in range(3):
        ax[i].plot(t.m_as("ms"), m1[:,i], color=f"C{order}")
        ax[i].fill_between(t.m_as("ms"), m1[:,i]-s1[:,i], m1[:,i]+s1[:,i], color=f"C{order}", alpha=0.5)

Visualize results

# Animate trajectories
plotter = pyvista.Plotter(off_screen=True, window_size=(600, 600))
local_functions.animate_trajectories(plotter, np.swapaxes(r.m_as("m"), 0, 1), t.m_as("ms"), filename="particles_test.gif",
                                     scalars=np.swapaxes(taylor_positions-r.m_as("m"), 0, 1),
                                     text_kwargs=dict(color="black", font_size=16),
                                     mesh_kwargs=dict(show_scalar_bar=True, color="gray", render_points_as_spheres=True, cmap="viridis",
                                                      scalar_bar_args={"title":"Displacement to step [m]", "position_y":0.05}),
                                     visualize_trajectories=True, trajectory_kwargs=dict(cmap="hot", scalar_bar_args={"title":"Displacement per step [m]",  "position_y":0.9}))

# Show final state
pyvista.close_all()
t_idx = -1
plot_mesh = pyvista.PolyData(taylor_positions[:, t_idx].numpy())
ref_mesh = pyvista.PolyData(r[:, t_idx].m_as("m"))
plot_mesh["disp"] = ref_mesh.points - plot_mesh.points
glyphs = plot_mesh.glyph(orient="disp", geom=pyvista.Line())
plotter = pyvista.Plotter(off_screen=True, window_size=(600, 600))
plotter.add_mesh(plot_mesh, render_points_as_spheres=True, show_scalar_bar=False, color="white")
plotter.add_mesh(ref_mesh, render_points_as_spheres=True, show_scalar_bar=False, color="steelblue")
plotter.add_mesh(glyphs, cmap="hot", show_scalar_bar=True, scalar_bar_args=dict(title="Difference to ref (m)"))
plotter.add_title("Final Difference to ref")
img = plotter.screenshot()


# Combine figures and display:
gif1 = imageio.v3.imread(f"particles_test.gif")
combined_gif = imageio.get_writer(f'particles_testc.gif')
for f1 in tqdm(gif1, total=gif1.shape[0]):
    combined_gif.append_data(np.concatenate([f1, img], axis=1))
combined_gif.close()
clear_output()
b64 = base64.b64encode(open("particles_testc.gif",'rb').read()).decode('ascii')
display(HTML(f'<img src="data:image/gif;base64,{b64}" />'))

LV-mesh example

Load determinisitic trajectories

The node-trajectories of a sliced contracting left ventricle mesh are used to define the input for the trajectory module.

files = [f'../../../cmr-random-diffmaps/notebooks/example_data/example_mesh/Displacements/Displ{i:03}.vtk' for i in range(0, 100, 1)]
timing = Quantity(np.loadtxt("../../../cmr-random-diffmaps/notebooks/example_data/example_mesh/PV_loop_reordered.txt")[:, 1], "ms")

original_mesh = cmrsim.datasets.CardiacMeshDataset.from_list_of_meshes(files, timing[:len(files)],
                                                                       mesh_numbers=(4, 40, 30),
                                                                       time_precision_ms=3)
refined_mesh = original_mesh.refine(longitudinal_factor=2, circumferential_factor=2, radial_factor=3)
slice_dict = dict(slice_thickness=Quantity(3, "mm"), slice_normal=np.array((0., 0., 1.)),
                  slice_position=Quantity([0, 0, -3], "cm"), reference_time = Quantity(105, "ms"))

temp = refined_mesh.select_slice(**slice_dict)
slice_mesh = cmrsim.datasets.MeshDataset(temp, refined_mesh.timing)

t, r = slice_mesh.get_trajectories(start=Quantity(105, "ms"), end=Quantity(300, "ms"))
np.min(t), np.max(t), r.shape, r.units

(105.045 <Unit('millisecond')>,
 298.975 <Unit('millisecond')>,
 (2732, 49, 3),
 <Unit('meter')>)

Visualize the reference trajectories

plotter = pyvista.Plotter(off_screen=True, window_size=(600, 600))
local_functions.animate_trajectories(plotter, np.swapaxes(r.m_as("m"), 0, 1), t.m_as("ms"), filename="particlesup.gif",
                                     text_kwargs=dict(color="white", font_size=16),
                                     mesh_kwargs=dict(show_scalar_bar=False, color="gray", render_points_as_spheres=True),
                                     visualize_trajectories=True, trajectory_kwargs=dict(cmap="hot", scalar_bar_args={"title":"Displacement per step [m]"}))

plotter = pyvista.Plotter(off_screen=True, window_size=(600, 600))
local_functions.animate_trajectories(plotter,
                                     np.swapaxes(r.m_as("m"), 0, 1)[:, 0:10],
                                     t.m_as("ms"), filename="single_particle.gif",
                                     text_kwargs=dict(color="white", font_size=16),
                                     mesh_kwargs=dict(show_scalar_bar=False, color="gray", render_points_as_spheres=True),
                                     visualize_trajectories=True,
                                     trajectory_kwargs=dict(cmap="hot", scalar_bar_args={"title":"Displacement per step [m]"}))

gif1 = imageio.v3.imread(f"particlesup.gif")
gif2 = imageio.v3.imread(f"single_particle.gif")
combined_gif = imageio.get_writer(f'particles.gif')
for f1, f2 in tqdm(zip(gif1, gif2), total=gif1.shape[0]):
    combined_gif.append_data(np.concatenate([f1, f2], axis=1))
combined_gif.close()
clear_output()
b64 = base64.b64encode(open("particles.gif",'rb').read()).decode('ascii')
display(HTML(f'<img src="data:image/gif;base64,{b64}" />'))

Fit taylor expansion and evaluate the trajectories

trajectory_module = cmrsim.trajectory.TaylorTrajectoryN(order=8, time_grid=t.m_as("ms"), particle_trajectories=r.m_as("m"))

Evaluate positions at the reference time-points to check validity. Also compute the velocity estimation and check if it is valid

taylor_positions, _ = trajectory_module(tf.constant(t.m_as("ms"), dtype=tf.float32))
taylor_positions_p, _ = trajectory_module(tf.constant(t.m_as("ms")+0.1, dtype=tf.float32))

taylor_velocities = (taylor_positions - taylor_positions_p).numpy() / 0.1
reference_velocities = np.gradient(r.m_as("m"), t.m_as("ms"), axis=1)
taylor_positions.numpy().shape, r.shape, taylor_velocities.shape, reference_velocities.shape

((2732, 49, 3), (2732, 49, 3), (2732, 49, 3), (2732, 49, 3))

Visualize taylor trajectories

plotter = pyvista.Plotter(off_screen=True, window_size=(600, 600))
local_functions.animate_trajectories(plotter,
                                     np.swapaxes(taylor_positions, 0, 1),
                                     t.m_as("ms"),
                                     filename="particlestaylor.gif",
                                     scalars=np.swapaxes(taylor_positions-r.m_as("m"), 0, 1)*1000,
                                     text_kwargs=dict(color="white", font_size=16),
                                     mesh_kwargs=dict(show_scalar_bar=True, cmap="hot", render_points_as_spheres=True, scalar_bar_args={"title":"Difference to reference [mm]"}),
                                     visualize_trajectories=False)

plotter = pyvista.Plotter(off_screen=True, window_size=(600, 600))
local_functions.animate_trajectories(plotter,
                                     np.swapaxes(taylor_positions, 0, 1) / 1000,
                                     t.m_as("ms"),
                                     filename="particlestaylorv.gif",
                                     scalars=np.swapaxes(taylor_velocities-reference_velocities, 0, 1),
                                     text_kwargs=dict(color="white", font_size=16),
                                     mesh_kwargs=dict(show_scalar_bar=True, cmap="hot", render_points_as_spheres=True, scalar_bar_args={"title":"Difference to reference [mm/ms]"}),
                                     visualize_trajectories=False)

gif1 = imageio.v3.imread(f"particlestaylor.gif")
gif2 = imageio.v3.imread(f"particlestaylorv.gif")
combined_gif = imageio.get_writer(f'particlestay.gif')
for f1, f2 in tqdm(zip(gif1, gif2), total=gif1.shape[0]):
    combined_gif.append_data(np.concatenate([f1, f2], axis=1))
combined_gif.close()
clear_output()
b64 = base64.b64encode(open("particlestay.gif",'rb').read()).decode('ascii')
display(HTML(f'<img src="data:image/gif;base64,{b64}" />'))