Dynamic Mode Decomposition (DMD)

~1 year of 6-hourly data of global windspeeds from ERA5

from ipyparallel import Client
rc = Client(profile="default")
rc.ids

if len(rc.ids) == 0:
    print("No engines found")
else:
    print(f"{len(rc.ids)} engines found")

4 engines found

%%px
import heat as ht

[stderr:2] /home/hopp_fa/heat/heat/core/_config.py:104: UserWarning: Heat has CUDA GPU-support (PyTorch version 2.7.1+cu126 and `torch.cuda.is_available() = True`), but CUDA-awareness of MPI could not be detected. This may lead to performance degradation as direct MPI-communication between GPUs is not possible.
  warnings.warn(

[stderr:0] /home/hopp_fa/heat/heat/core/_config.py:104: UserWarning: Heat has CUDA GPU-support (PyTorch version 2.7.1+cu126 and `torch.cuda.is_available() = True`), but CUDA-awareness of MPI could not be detected. This may lead to performance degradation as direct MPI-communication between GPUs is not possible.
  warnings.warn(

[stderr:3] /home/hopp_fa/heat/heat/core/_config.py:104: UserWarning: Heat has CUDA GPU-support (PyTorch version 2.7.1+cu126 and `torch.cuda.is_available() = True`), but CUDA-awareness of MPI could not be detected. This may lead to performance degradation as direct MPI-communication between GPUs is not possible.
  warnings.warn(

[stderr:1] /home/hopp_fa/heat/heat/core/_config.py:104: UserWarning: Heat has CUDA GPU-support (PyTorch version 2.7.1+cu126 and `torch.cuda.is_available() = True`), but CUDA-awareness of MPI could not be detected. This may lead to performance degradation as direct MPI-communication between GPUs is not possible.
  warnings.warn(

Preprocessing

The following steps have already been described in the notebook on loading and preprocessing:

%%px
u = ht.load("/p/scratch/training2546/datasets/10m_wind_2024_every6hours.nc", "u10", device="gpu", split=0)
v = ht.load("/p/scratch/training2546/datasets/10m_wind_2024_every6hours.nc", "v10", device="gpu", split=0)
n_timesteps, n_lon, n_lat = u.shape
w = ht.stack([u,v],axis=-1)

In order to apply Dynamic Mode Decomposition, we need to reshape the data such that there is a single feature-axis; consequently, the two components of windspeed at every location on earth need to be reshaped into a single data vector (for every time step). Moreover, DMD requires the time axis to be the second one and the feature axis the first one. This is done as follows:

%%px
w = ht.reshape(w, (n_timesteps, n_lon * n_lat * 2)).T
w.shape, w.split

Out[1:3]: ((2076480, 1464), 1)

Out[0:3]: ((2076480, 1464), 1)

Out[2:3]: ((2076480, 1464), 1)

Out[3:3]: ((2076480, 1464), 1)

In order to be able to execute the computations within DMD efficiently, splitting along the feature axis is desirable. Changing the split from one axis to another can be done as follows:

%%px
w.resplit_(0)
w.split

Out[3:4]: 0

Out[2:4]: 0

Out[1:4]: 0

Out[0:4]: 0

Please note that, in general, a resplit operation is expensive and should be avoided if not necessary. However, in the present case it is necessary.

Explore low-rank structure

To check whether application of DMD is reasonable, we first check whether there is some kind of low-rank structure in the data.

%%px
from heat.decomposition import PCA
pca = PCA(svd_solver="full")
pca.fit(w)

Out[0:5]: 
PCA({
    "n_components": null,
    "copy": true,
    "whiten": false,
    "svd_solver": "full",
    "tol": null,
    "iterated_power": 0,
    "n_oversamples": 10,
    "power_iteration_normalizer": "qr",
    "random_state": null
})

Out[1:5]: 
PCA({
    "n_components": null,
    "copy": true,
    "whiten": false,
    "svd_solver": "full",
    "tol": null,
    "iterated_power": 0,
    "n_oversamples": 10,
    "power_iteration_normalizer": "qr",
    "random_state": null
})

Out[3:5]: 
PCA({
    "n_components": null,
    "copy": true,
    "whiten": false,
    "svd_solver": "full",
    "tol": null,
    "iterated_power": 0,
    "n_oversamples": 10,
    "power_iteration_normalizer": "qr",
    "random_state": null
})

Out[2:5]: 
PCA({
    "n_components": null,
    "copy": true,
    "whiten": false,
    "svd_solver": "full",
    "tol": null,
    "iterated_power": 0,
    "n_oversamples": 10,
    "power_iteration_normalizer": "qr",
    "random_state": null
})

%%px
ht.norm(pca.singular_values_[1000:])/ht.norm(pca.singular_values_)

Out[0:6]: DNDarray(0.1008, dtype=ht.float32, device=gpu:0, split=None)

Dynamic Mode Decomposition (DMD)

Although the data are not really low-rank, we proceed with computing the DMD for demonstration purposes.

%%px
del pca, u, v # clean up before expensive computation
from heat.decomposition import DMD
dmd = DMD(svd_solver="full", svd_rank=1000)
dmd.fit(w)

Once DMD is computed, you can check the eigenvalues in order to get an impression of stability/instability of the dynamical system.

%%px
import matplotlib.pyplot as plt
%matplotlib inline

dmd.rom_eigenvalues_
re,im = dmd.rom_eigenvalues_.numpy().real, dmd.rom_eigenvalues_.numpy().imag
plt.figure()
plt.scatter(re,im)
plt.show()

[stderr:0] /home/hopp_fa/.local/lib/python3.10/site-packages/matplotlib/projections/__init__.py:63: UserWarning: Unable to import Axes3D. This may be due to multiple versions of Matplotlib being installed (e.g. as a system package and as a pip package). As a result, the 3D projection is not available.
  warnings.warn("Unable to import Axes3D. This may be due to multiple versions of "

[stderr:3] /home/hopp_fa/.local/lib/python3.10/site-packages/matplotlib/projections/__init__.py:63: UserWarning: Unable to import Axes3D. This may be due to multiple versions of Matplotlib being installed (e.g. as a system package and as a pip package). As a result, the 3D projection is not available.
  warnings.warn("Unable to import Axes3D. This may be due to multiple versions of "

[stderr:2] /home/hopp_fa/.local/lib/python3.10/site-packages/matplotlib/projections/__init__.py:63: UserWarning: Unable to import Axes3D. This may be due to multiple versions of Matplotlib being installed (e.g. as a system package and as a pip package). As a result, the 3D projection is not available.
  warnings.warn("Unable to import Axes3D. This may be due to multiple versions of "

[stderr:1] /home/hopp_fa/.local/lib/python3.10/site-packages/matplotlib/projections/__init__.py:63: UserWarning: Unable to import Axes3D. This may be due to multiple versions of Matplotlib being installed (e.g. as a system package and as a pip package). As a result, the 3D projection is not available.
  warnings.warn("Unable to import Axes3D. This may be due to multiple versions of "

[output:2]

[output:1]

[output:3]

[output:0]

Moreover, you can plot the so-called DMD-modes. Of course, for plotting the reshapes done above must be undone first.

%%px
first_mode_u_re = ht.real(dmd.dmdmodes_[:,0].reshape(n_lon,n_lat,2)[:,:,0])

plt.figure()
plt.imshow(first_mode_u_re.numpy())
plt.show()

[output:0]

[output:3]

[output:1]

[output:2]