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 ReducedComplexityModeling.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.GMLDatasets — Module
GMLDatasetsDataset-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 ReducedComplexityModeling.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)ReducedComplexityModeling.DataLoader — Method
DataLoader(images::AbstractArray{T, 3}, labels::AbstractVector)Make an instance of ReducedComplexityModeling.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 ReducedComplexityModeling because it is the piece that is specific to labelled image data: it cuts the images into patches and one-hot encodes the labels, and it builds a :TimeSeries loader. The generic tensor–tensor constructor of ReducedComplexityModeling takes its two tensors as they are and builds a :RegularData loader.
GMLDatasets.angular_to_euclidean — Method
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$.
GMLDatasets.euclidean_to_angular — Method
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.
GMLDatasets.fashion_mnist — Function
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.
GMLDatasets.fashion_mnist_data_loader — Function
fashion_mnist_data_loader(split = :train; patch_length = 7, transform = identity, suppress_info = false)Load Fashion-MNIST and wrap it in a ReducedComplexityModeling.DataLoader.
Like mnist_data_loader, but for fashion_mnist.
GMLDatasets.linear_wave — Method
linear_wave(; N, μ, timespan, n_time_steps, integrator)Integrate the discretized linear wave equation for each wave speed in μ and return the GeometricSolutions.EnsembleSolution.
This is GeometricProblems.LinearWave.hodeensemble composed with GeometricIntegrators.integrate and nothing more. The problem is the one GeometricProblems defines: the wave equation $\partial_{tt}u = \mu^2\partial_{\xi\xi}u$ on $\xi \in (-1/2, 1/2)$ with homogeneous Dirichlet boundary conditions, discretized on the N interior points $\xi_i = i\Delta\xi - 1/2$, $\Delta\xi = 1/(N+1)$, with the Hamiltonian
\[H(q, p; \mu) = \frac{1}{2}p^Tp - \frac{\mu^2}{2}q^TD_{\xi\xi}q ,\]
where $D_{\xi\xi}$ is the central-difference matrix. The state has $2N$ components.
Every wave speed gets its own initial condition, GeometricProblems.LinearWave.initial_condition(μ, N): a cubic-spline pulse that starts at $\xi = -\mu/2$ with width $\propto 1/\mu$ and initial momentum $p = -\mu\,\partial_\xi q$, so that it travels to the right at speed $\mu$. Over the default timespan it does not reach the boundary for any $\mu \in [5/12, 2/3]$.
The defaults are the training set of the symplectic-autoencoder experiment: twenty equally spaced values $\mu \in [5/12, 2/3]$, $N = 128$, n_time_steps = 200 snapshots on $t \in [0, 1]$ (so the time step is $1/199$), integrated with the implicit midpoint rule.
Use linear_wave_snapshots to turn the solution into the data tensor a ReducedComplexityModeling.DataLoader takes:
using GeometricMachineLearning
dl = DataLoader(linear_wave_snapshots(linear_wave()); autoencoder = true)GMLDatasets.linear_wave_parameters — Method
linear_wave_parameters(solution)The wave speeds $\mu$ of a linear_wave solution, one per trajectory and in the order of the slices of linear_wave_snapshots.
GMLDatasets.linear_wave_snapshots — Method
linear_wave_snapshots(solution)The snapshots of a linear_wave solution as one $2N × n_t × n_μ$ array: the rows are $(q_1, …, q_N, p_1, …, p_N)$, one column per time step (the initial condition first) and one slice per wave speed, in the order of μ.
That is the layout every symplectic architecture in GeometricMachineLearning assumes (the first half of the rows is $q$, the second half $p$) and the one DataLoader reads off a tensor.
Given a single GeometricSolution it returns the $2N × n_t$ matrix of that trajectory.
GMLDatasets.mnist — Function
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.
GMLDatasets.mnist_data_loader — Function
mnist_data_loader(split = :train; patch_length = 7, transform = identity, suppress_info = false)Load MNIST and wrap it in a ReducedComplexityModeling.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.
GMLDatasets.onehotbatch — Method
onehotbatch(target)
onehotbatch(S, 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$. By default its element type is the label type; pass a numeric type S to choose a different element type. In either form, the output uses the same compute backend as target.
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
0GMLDatasets.pendulum — Method
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.
momentum_fractions replaces that grid with a separatrix-relative one, which is what a data set straddling the separatrix wants. Passing a vector of numbers f takes the same qsamples angles and gives each of them the momenta $p_\theta = f\,p_\mathrm{sep}(\theta)$, where separatrix_momentum is the momentum that puts the pendulum exactly on the separatrix at that angle. The energy of such an initial condition is
\[H = mg\ell\left[f^2(1 - \cos\theta) + \cos\theta\right] ,\]
so $|f| < 1$ librates, $|f| = 1$ is the separatrix itself and $|f| > 1$ rotates, whatever the angle, and the sign of $f$ picks the direction of rotation. A Cartesian $(\theta, p_\theta)$ grid cannot express that: a fixed $p_\theta$ crosses the separatrix as $\theta$ varies, so a rectangle in $(\theta, p_\theta)$ always mixes the two regimes and never controls how close to the separatrix the data comes. With fractions, how densely the crossing is covered is chosen directly — fractions just above one in modulus fill the band of rotating orbits nearest the separatrix, which is the region a single Cartesian grid leaves empty.
pmin, pmax and psamples are unused when momentum_fractions is given. The fraction is undefined where $p_\mathrm{sep}$ vanishes, that is at the unstable equilibrium $\theta \equiv 0$ modulo $2\pi$, and an angle grid that contains one of those points is rejected rather than silently collapsed onto the fixed point.
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 separatrix_momentum and pendulum_energy.
GMLDatasets.pendulum_energy — Function
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.
GMLDatasets.separatrix_momentum — Function
separatrix_momentum(θ, parameters = GeometricProblems.Pendulum.default_parameters())The momentum that puts the pendulum exactly on the separatrix at angle $\theta$:
\[p_\mathrm{sep}(\theta) = m\ell\sqrt{2g\ell\,(1 - \cos\theta)} ,\]
the positive root of $H(\theta, p_\theta) = mg\ell$. It vanishes at the unstable equilibrium $\theta \equiv 0$ modulo $2\pi$, where the separatrix meets itself, and is largest at $\theta = \pi$.
This is the scale pendulum's momentum_fractions measures against.
GMLDatasets.split_and_flatten — Method
split_and_flatten(input::AbstractArray)::AbstractArrayPerform 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 36Here we see that split_and_flatten:
- splits the original matrix into four $3\times3$ matrices and then
- 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
- $\mathtt{path\_length}^2$ and
number_of_patches.