GMLDatasets

Dataset-backed demonstrations for GeometricMachineLearning and GeometricOptimizers.

Both of those packages are libraries for scientific machine learning and neither should depend on an image-dataset package to document itself. This one does: it holds the MLDatasets glue, the MNIST and Fashion-MNIST demonstrations, the numerical experiment of [1] that shows manifold optimization making a vision transformer trainable at all, and a symplectically integrated pendulum data set for the autoencoders.

The data loaders

mnist_data_loader and fashion_mnist_data_loader hand back a GeometricMachineLearning.DataLoader in the time series format a transformer consumes — each $28\times28$ image cut into 16 patches of $7\times7$, each patch flattened into a column, the labels one-hot encoded:

using GMLDatasets

dl = mnist_data_loader(:train; patch_length = 7, suppress_info = true)
(dl.input_dim, dl.input_time_steps, dl.n_params, dl.output_dim)
(49, 16, 60000, 10)

The pieces are available on their own too: mnist and fashion_mnist return the raw images and labels, split_and_flatten cuts an image into flattened patches and onehotbatch encodes the labels.

The pendulum

pendulum integrates a grid of pendulum trajectories and angular_to_euclidean lifts them into $\mathbb{R}^4$, where two dimensions of data sit on a curved submanifold of four — which is what a symplectic autoencoder is for. There is no data loader of its own here; the array is already in the shape DataLoader reads:

solution = pendulum()
dl = DataLoader(angular_to_euclidean(solution); autoencoder = true, suppress_info = true)
(dl.input_dim, dl.input_time_steps, dl.n_params)
(4, 101, 100)

Where to go next

  • The Pendulum Data Set is the pendulum in full: the sign convention, the lift into $\mathbb{R}^4$, and the energy check that says the integrator did its job.
  • MNIST Tutorial trains a vision transformer with weights on the Stiefel manifold, using GeometricMachineLearning, and compares four optimizers on it.
  • The Numerical Experiment on Homogeneous Spaces carries the figures for a 500-epoch run of that experiment against GeometricOptimizers — training loss, test accuracy and drift off the manifold.
  • Running the Experiments is the operational side: what each script in scripts/ does, what the four configurations are for, and the device-memory handling the GPU runs need.

Library

GMLDatasets.GMLDatasetsModule
GMLDatasets

Dataset-backed demonstrations for GeometricMachineLearning and GeometricOptimizers.

Both of those packages are libraries for scientific machine learning and neither should depend on an image-dataset package to document itself. Everything that does lives here: the MLDatasets glue, the MNIST and Fashion-MNIST demonstrations and the numerical experiment from [1].

The entry points are mnist_data_loader and fashion_mnist_data_loader, which hand back a GeometricMachineLearning.DataLoader in the time series format a transformer wants:

using GMLDatasets

dl      = mnist_data_loader(:train; patch_length = 7)
dl_test = mnist_data_loader(:test;  patch_length = 7)

Not everything here is an image. pendulum integrates a grid of mathematical pendula with GeometricProblems and GeometricIntegrators, and angular_to_euclidean lifts them into four dimensions, where they sit on a two-dimensional submanifold — a data set with a geometry worth recovering, for the symplectic autoencoders:

dl = DataLoader(angular_to_euclidean(pendulum()); autoencoder = true)
source
GeometricMachineLearning.DataLoaderMethod
DataLoader(images::AbstractArray{T, 3}, labels::AbstractVector)

Make an instance of GeometricMachineLearning.DataLoader for an image classification problem.

images is a tensor whose first two axes are the image axes and whose third axis indexes the images; labels is a vector of integer labels, one per image. Each image is cut into square patches by split_and_flatten and each patch is flattened into a column, which casts the data into the time series format a transformer consumes. The labels are encoded by onehotbatch.

Arguments

There are two keyword arguments:

  • patch_length = 7. This is the length of the patch in the $x$ and the $y$ direction;
  • suppress_info = false.

MNIST images are of size $28\times28$. For patch_length = 7 such an image is therefore split into 16 patches of $7\times7$, i.e. it becomes a matrix in $\mathbb{R}^{49\times16}$ [1].

Implementation

This method lives in GMLDatasets rather than in GeometricMachineLearning because it is the piece that is specific to labelled image data. It is also the only way to build a DataLoader whose input is Float32 and whose output is Int: the generic tensor–tensor constructor in GeometricMachineLearning requires a single shared element type, whereas a one-hot target is integer-valued whatever the images are.

source
GMLDatasets.angular_to_euclideanMethod
angular_to_euclidean(θ, pθ; l = 1)
angular_to_euclidean(solution)

Lift the canonical pendulum coordinates into the Euclidean coordinates of the bob in the plane:

\[q = (\ell\sin\theta,\; \ell\cos\theta), \qquad p = (p_\theta\cos\theta/\ell,\; -p_\theta\sin\theta/\ell).\]

The lift is a symplectomorphism onto its image, so the Euclidean Hamiltonian pendulum_energy agrees with GeometricProblems.Pendulum.hamiltonian and a symplectic integrator stays symplectic under it. Its image is a two-dimensional submanifold of $\mathbb{R}^4$: $q$ lies on the circle of radius $\ell$ and $p$ is tangent to that circle, i.e. $\|q\| = \ell$ and $q\cdot{}p = 0$ hold exactly.

Given two vectors of $n$ samples this returns the pair of $2\times{}n$ matrices $q$ and $p$.

Given a GeometricSolutions.GeometricSolution or EnsembleSolution — what pendulum hands back — it returns the data set instead: one $4\times{}n_t\times{}n$ array whose rows are $(q_1, q_2, p_1, p_2)$, one column per time step and one slice per trajectory. That is the layout every symplectic architecture in GeometricMachineLearning assumes (the first half of the rows is $q$, the second half is $p$) and the one DataLoader reads off a tensor, so nothing further is needed:

dl = DataLoader(angular_to_euclidean(pendulum()); autoencoder = true)

$\ell$ is taken from the problem each solution was integrated from rather than assumed to be one.

euclidean_to_angular is the inverse, up to the $2\pi$-periodicity of $\theta$.

source
GMLDatasets.euclidean_to_angularMethod
euclidean_to_angular(q, p; l = 1)

Project the Euclidean coordinates of the bob back onto the canonical $(\theta, p_\theta)$.

q and p are $2\times{}n$ matrices as produced by angular_to_euclidean, which this inverts. The angle comes back wrapped into $(-\pi, \pi]$, so the round trip is the identity only for angles that were in that interval to begin with.

source
GMLDatasets.fashion_mnistFunction
fashion_mnist(split = :train)

Return the Fashion-MNIST images and labels for split, which is either :train or :test.

Fashion-MNIST is drop-in compatible with MNIST — $28\times28$ grayscale images in ten classes — so everything that works on one works on the other. See mnist.

source
GMLDatasets.mnistFunction
mnist(split = :train)

Return the MNIST images and labels for split, which is either :train or :test.

This is MLDatasets.MNIST(split = split)[:] and nothing more — the images come back as a $28\times28\times{}n$ tensor of Float32 and the labels as a vector of Int in $0, \ldots, 9$. Use mnist_data_loader to get them in the form a transformer consumes.

The data set is downloaded on first use. Set ENV["DATADEPS_ALWAYS_ACCEPT"] = true to skip the download prompt in a non-interactive session.

source
GMLDatasets.mnist_data_loaderFunction
mnist_data_loader(split = :train; patch_length = 7, transform = identity, suppress_info = false)

Load MNIST and wrap it in a GeometricMachineLearning.DataLoader.

This composes mnist with DataLoader, i.e. it cuts each image into patches of patch_length $\times$ patch_length, flattens each patch into a column and one-hot encodes the labels.

transform is applied to the images and to the labels before the DataLoader is built, which is where the data are moved onto an accelerator:

using CUDA

dl = mnist_data_loader(:train; transform = cu)

The transform is applied before rather than after the DataLoader is constructed so that split_and_flatten and onehotbatch run their kernels on the device.

source
GMLDatasets.onehotbatchMethod
onehotbatch(target)

Performs a one-hot-batch encoding of a vector of integers: $input\in\{0,1,\ldots,9\}^\ell$.

The output is a tensor of shape $10\times1\times\ell$.

If the input is $0$, this function produces:

\[0 \mapsto \begin{bmatrix} 1 & 0 & \ldots & 0 \end{bmatrix}^T.\]

In more abstract terms: $i \mapsto e_i$.

Examples

using GMLDatasets

target = [0]
onehotbatch(target)

# output

10×1×1 Array{Int64, 3}:
[:, :, 1] =
 1
 0
 0
 0
 0
 0
 0
 0
 0
 0
source
GMLDatasets.pendulumMethod
pendulum(; qmin, qmax, pmin, pmax, qsamples, psamples, parameters, timespan, timestep, integrator)

Integrate an ensemble of mathematical pendula and return the GeometricSolutions.EnsembleSolution.

This is GeometricProblems.Pendulum.hodeensemble composed with GeometricIntegrators.integrate and nothing more. The Hamiltonian is the one GeometricProblems defines,

\[H(\theta, p_\theta) = \frac{p_\theta^2}{2m\ell^2} + mg\ell\cos(\theta),\]

so the potential is at its minimum at $\theta = \pi$: the pendulum hangs down at $\theta = \pi$ and stands upright at $\theta = 0$. Trajectories with $H < mg\ell$ librate about $\theta = \pi$ and trajectories with $H > mg\ell$ rotate.

The initial conditions are a Cartesian grid: qsamples angles spread evenly over [qmin, qmax] times psamples momenta spread evenly over [pmin, pmax], each given as a one-element vector because the pendulum has one degree of freedom. The defaults are the grid GeometricProblems itself uses — a hundred trajectories covering both libration and rotation — over a longer timespan than its default, so that there is enough of each trajectory to learn from.

parameters is the $(\ell, m, g)$ named tuple, integrator any GeometricIntegrators method. The default is Gauss collocation with two stages, which is symplectic, so the energy of each trajectory oscillates within a bounded band around its initial value over the whole run rather than drifting away from it. On the default grid that band is about $2\cdot{}10^{-6}$ wide.

The canonical coordinates the ensemble carries are two-dimensional. Use angular_to_euclidean to lift them into $\mathbb{R}^4$, where they trace out a two-dimensional submanifold — which is what makes them a worthwhile test case for a GeometricMachineLearning.SymplecticAutoencoder:

using GeometricMachineLearning

dl = DataLoader(angular_to_euclidean(pendulum()); autoencoder = true)

See also pendulum_energy.

source
GMLDatasets.pendulum_energyFunction
pendulum_energy(θ, pθ, parameters = GeometricProblems.Pendulum.default_parameters())
pendulum_energy(q, p, parameters = ...)
pendulum_energy(data, parameters = ...)
pendulum_energy(solution)

Evaluate the pendulum Hamiltonian, in canonical or in Euclidean coordinates.

Given two vectors this is GeometricProblems.Pendulum.hamiltonian broadcast over them. Given the Euclidean coordinates of angular_to_euclidean — either as two arrays whose first axis is the two Euclidean components, or as the single four-row array that function returns for a solution — it is the same Hamiltonian written in those coordinates,

\[H(q, p) = \frac{\|p\|^2}{2m} + mgq_2,\]

which is what the lift being a symplectomorphism buys: $\|p\|^2 = p_\theta^2/\ell^2$ because $p$ is tangent to the circle, and $q_2 = \ell\cos\theta$. Given a solution of pendulum the parameters are read off the problem it was integrated from.

The result keeps every axis but the first, so a $2\times{}n_t\times{}n$ tensor gives an $n_t\times{}n$ matrix of energies — one column per trajectory, which is how a symplectic integrator is checked for drift.

source
GMLDatasets.split_and_flattenMethod
split_and_flatten(input::AbstractArray)::AbstractArray

Perform a preprocessing of an image into flattened patches.

This rearranges the input data so that it can easily be processed with a transformer.

Examples

Consider a matrix of size $6\times6$ which we want to divide into patches of size $3\times3$.

using GMLDatasets

input = [ 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]

split_and_flatten(input; patch_length = 3, number_of_patches = 4)

# output

9×4 Matrix{Int64}:
  1  19   4  22
  7  25  10  28
 13  31  16  34
  2  20   5  23
  8  26  11  29
 14  32  17  35
  3  21   6  24
  9  27  12  30
 15  33  18  36

Here we see that split_and_flatten:

  1. splits the original matrix into four $3\times3$ matrices and then
  2. flattens each matrix into a column vector of size $9.$

After this all the vectors are put together again to yield a $9\times4$ matrix.

Arguments

The optional keyword arguments are:

  • patch_length: by default this is 7.
  • number_of_patches: by default this is 16.

The sizes of the first and second axis of the output of split_and_flatten are

  1. $\mathtt{path\_length}^2$ and
  2. number_of_patches.
source