documentation and testing updates

docs/CHANGELOG.md

@@ -3,15 +3,11 @@
 <!---
 This is a comment.
 Remember to align the itemized text with the first line of an item within a list.
-
-PLEASE REMEMBER TO CHANGE THE '..master' WITH AN ACTUAL TAG in GITHUB LINK.
 -->
 
-## jax 0.2.10 (Unreleased)
+## EMLP 0.8.0 (Unreleased)
 
-* Mockup
-* a
-* New features:
+<!-- * New features:
   * {func}`jax.scipy.stats.chi2` is now available as a distribution with logpdf and pdf methods.
   * {func}`jax.scipy.stats.betabinom` is now available as a distribution with logpmf and pmf methods.
-  *
+  * -->

docs/documentation.md

@@ -1,6 +1,6 @@
 # Update documentation
 
-To rebuild the documentation, install several packages:
+To rebuild the documentation, you need to install the requirement packages:
 ```
 pip install -r docs/requirements.txt
 ```
@@ -16,7 +16,7 @@ sphinx-build -b html -D jupyter_execute_notebooks=off docs docs/build/html
 You can then see the generated documentation in `docs/build/html/index.html`.
 
 ## Update notebooks
-
+We follow the approach of [Jax](https://jax.readthedocs.io/) in how the documentation is setup and how to contribute.
 We use [jupytext](https://jupytext.readthedocs.io/) to maintain two synced copies of the notebooks
 in `docs/notebooks`: one in `ipynb` format, and one in `md` format. The advantage of the former
 is that it can be opened and executed directly in Colab; the advantage of the latter is that
@@ -27,7 +27,8 @@ it makes it much easier to track diffs within version control.
 For making large changes that substantially modify code and outputs, it is easiest to
 edit the notebooks in Jupyter or in Colab. To edit notebooks in the Colab interface,
 open <http://colab.research.google.com> and `Upload` from your local repo.
-Update it as needed, `Run all cells` then `Download ipynb`.
+Update it as needed, `Run all cells` then `Download ipynb` (for editing and running in Colab you will need to add 
+`!pip install git+https://github.com/mfinzi/equivariant-MLP.git`).
 You may want to test that it executes properly, using `sphinx-build` as explained above.
 
 ### Editing md
@@ -42,14 +43,4 @@ using [jupytext](https://jupytext.readthedocs.io/) by running:
 
 ```
 $ jupytext --sync docs/notebooks/*
-```
-
-Alternatively, you can run this command via the [pre-commit](https://pre-commit.com/)
-framework by executing the folloing in the main JAX directory:
-
-```
-$ pre-commit run --all
-```
-
-See the pre-commit framework documentation for information on how to set your local git
-environment to execute this automatically.
+```

docs/index.rst

@@ -12,19 +12,24 @@ A type system for the automated construction of equivariant layers.
 
 .. toctree::
    :maxdepth: 1
-   :caption: Tutorials
+   :caption: Getting Started
 
    notebooks/quickstart.md
-   notebooks/model.md
+   notebooks/building_a_model.md
    notebooks/new_groups.md
 
 .. toctree::
    :maxdepth: 1
-   :caption: Advanced Tutorials
+   :caption: Advanced Features
 
    notebooks/mixed_tensors.md
    notebooks/new_representations.md
 
+.. toctree::
+   :maxdepth: 1
+   :caption: Examples
+
+
 .. toctree::
    :maxdepth: 1
    :caption: Notes
@@ -34,7 +39,7 @@ A type system for the automated construction of equivariant layers.
 .. toctree::
    :maxdepth: 2
    :caption: Developer documentation
-
+   testing.md
    documentation.md
 
 

docs/notebooks/quickstart.md

@@ -12,24 +12,30 @@ kernelspec:
   name: python3
 ---
 
-# Getting Familiar with the Type System
+# Getting Started with Equivariant Representations
 
 ```{code-cell} ipython3
-from emlp.solver.representation import V,sparsify_basis
-from emlp.solver.groups import *
-import jax.numpy as jnp
-import logging
-logging.getLogger().setLevel(logging.INFO)
+
 ```
 
-EMLP computes the symmetric subspace for a linear representation $\rho$ and a matrix group $G$, solving the constraint to find an element $v\in V$ that satisfies$$\forall g\in G: \ \ \rho(g)v=v$$
+EMLP computes the symmetric subspace for a linear representation $\rho$ and a matrix group $G$, solving the constraint to find an element $v\in V$ that satisfies $\forall g\in G: \ \ \rho(g)v=v$
 
-For example, we can find invariant vectors of the cyclic translation group $\mathbb{Z}_n$ which is just the constant $\vec{1}$ scaled to have unit norm. 
+For example, we can find invariant vectors of the cyclic translation group $\mathbb{Z}_n$ which is just the constant $\vec{1}$ scaled to have unit norm.
+
+```{code-cell} ipython3
+from emlp.solver.representation import V,sparsify_basis
+from emlp.solver.groups import Z,S,SO,O,O13,RubiksCube
+import jax.numpy as jnp
+```
 
 ```{code-cell} ipython3
 V(Z(5)).symmetric_basis()
 ```
 
+## The Representation Type System
+
++++
+
 Each implemented group comes with a faithful 'base' representation $V$. Because faithful representations are one-to-one, we can build any representation by transforming this base representation.
 
 We provide several operators to transform and construct representations in different ways built and later go on to show how to do this more generally. In our type system, representations can be combined with the direct sum $\rho_a \oplus\rho_b$ operator, the tensor product $\rho_a\otimes\rho_b$, the dual $\rho^*$. We implement these with the python operators `+`, `*`, and `.T`.
@@ -44,7 +50,7 @@ We can combine and use these operators interchangeably:
 (V+V.T)*(V*V.T+V)
 ```
 
-We use the shorthand $cV$ can be used for $V\oplus V\oplus...\oplus V$ and $V^c = V\otimes V\otimes...\otimes V$, not that this different from the typical notation with cartesian products of sets.
+We use the shorthand $cV$ to mean $V\oplus V\oplus...\oplus V$ and $V^c = V\otimes V\otimes...\otimes V$. Note that this differs from the common notation where $V^c$ denotes the cartesian products the set (like with $\mathbb{R}^c$) which would be the same as $cV$ in this notation. Being more formal we could distinguish the two by denoting $V^{\otimes c}=V\otimes V\otimes...\otimes V$ but to stay consistent with the python interface, we will not.
 
 ```{code-cell} ipython3
 5*V*2
@@ -67,14 +73,20 @@ Although for groups like the Lorentz group $SO(1,3)$ with non orthogonal represe
 
 ```{code-cell} ipython3
 V(SO(3)).T+V(SO(3))
+```
 
+```{code-cell} ipython3
 V(SO13()).T+V(SO13())
 ```
 
 Linear maps from $V_1\rightarrow V_2$ have the type $V_2\otimes V_1^*$. The `V>>W` is shorthand for `W*V.T` and produces linear maps from `V` to `W`.
 
 Imposing (cyclic) Translation Equivariance $G=\mathbb{Z}_n$ on linear maps $V\rightarrow V$ yields circular convolutions (circulant matrices) which can be expressed as a linear combination of $n$ basis elements of size $n\times n$.
 
++++
+
+## Exploring and Visualizing Equivariant Bases
+
 ```{code-cell} ipython3
 G = Z(6)
 repin = V(G)
@@ -142,7 +154,44 @@ vis(repin,repout)
 print((repin>>repout).symmetric_basis().shape)
 ```
 
-## Composite Representations
+How about the continuous $2$D rotation group $SO(3)$? It's well known that the only equivariant object for the vector space $V^{\otimes 3}$ is the Levi-Civita symbol $\epsilon_{ijk}$. Since the values are both $0$, positive, and negative (leading to more than `Q.shape[-1]` clusters) we disable the clustering.
+
+```{code-cell} ipython3
+W = V(SO(3))
+repin = W**2
+repout = W
+Q = (repin>>repout).symmetric_basis()
+print(f"Basis matrix of shape {Q.shape}")
+vis(repin,repout,cluster=False)
+```
+
+```{code-cell} ipython3
+print(sparsify_basis(Q).reshape(3,3,3))
+```
+
+```{code-cell} ipython3
+from emlp.solver.representation import T
+```
+
+### High Dimensional Representations
+
++++
+
+We can also solve for very high dimensional representations which we automatically switch to using the automated iterative Krylov subspace method
+
+```{code-cell} ipython3
+vis(W**3,W**3)
+```
+
+```{code-cell} ipython3
+vis(W**5,W**3)
+```
+
+```{code-cell} ipython3
+vis(V(RubiksCube()),V(RubiksCube()))
+```
+
+## Composite Representations and Lazy Matrices
 
 +++
 
@@ -157,25 +206,24 @@ vis(repin,repout)
 print((repin>>repout).symmetric_basis().shape)
 ```
 
-We can compute the bases for representations that have many copies or multiplicity of a given representation type, such as for the many channels in a neural network. The `rep.symmetric_basis()` and `rep.symmetric_projector()` can return lazy matrices $Q$ and $P=QQ^T$ when the representations are composite (or when the representation is specified lazily). [implementation change is making this much slower than normal, using smaller values]
+Representations that have many copies or multiplicity of a given representation type, such as for the many channels in a neural network, are simply examples of the $\otimes$ operator (`+` in python). The `rep.symmetric_basis()` and `rep.symmetric_projector()` can return lazy matrices $Q$ and $P=QQ^T$ when the representations are composite (or when the representation is specified lazily). [implementation change is making this much slower than normal, using smaller values]
 
 +++
 
 For example with a more realistically sized layer with 100 global constants, 100 set feature channels, and 20 edge feature channels ($100V^0+100V^1+20V^2$) we have
 
 ```{code-cell} ipython3
-
 W = V(S(6))
 repin = 100*W**0 + 100*W+20*W**2
 repout = repin
 rep_map = repin>>repout
 print(f"{rep_map}, of size {rep_map.size()}")
 
-# Q = rep_map.symmetric_basis()
-# print(Q.shape)
+Q = rep_map.symmetric_basis()
+print(f"Basis matrix of shape {Q.shape}")
 ```
 
-These Lazy matrices are modeled after https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.LinearOperator.html. Unfortunately the larger matrices are harder to visualize, maybe someone can figure out how to do this better?
+These Lazy matrices are modeled after https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.LinearOperator.html. Unfortunately the larger matrices are harder to visualize, let us know if you have a good idea for how!
 
 ```{code-cell} ipython3
 P =rep_map.symmetric_projector()
@@ -184,39 +232,3 @@ v = P@v
 plt.imshow(v.reshape(repout.size(),repin.size()))
 plt.axis('off')
 ```
-
-But more fun than continuous groups are discrete groups!
-How about the $2$D rotation group $SO(3)$. It's well known that the only equivariant object for the vector space $V^{\otimes 3}$ is the Levi-Civita symbol $\epsilon_{ijk}$. Since the values are both $0$, positive, and negative (leading to more than `Q.shape[-1]` clusters) we disable the clustering.
-
-```{code-cell} ipython3
-W = V(SO(3))
-repin = W**2
-repout = W
-Q = (repin>>repout).symmetric_basis()
-print(Q.shape)
-vis(repin,repout,cluster=False)
-
-print(sparsify_basis(Q).reshape(3,3,3))
-```
-
-```{code-cell} ipython3
-from emlp.solver.representation import T
-```
-
-## High Dimensional Representations
-
-+++
-
-We can also solve for very high dimensional representations which we automatically switch to using the automated iterative Krylov subspace method
-
-```{code-cell} ipython3
-vis(W**3,W**3)
-```
-
-```{code-cell} ipython3
-vis(W**5,W**3)
-```
-
-```{code-cell} ipython3
-vis(V(RubiksCube()),V(RubiksCube()))
-```

docs/testing.md

@@ -0,0 +1,15 @@
+# Running the Tests
+
+We have a number of tests checking the equivariance of representations constructed in
+ different ways (`.T`, `*`, `+`) for the groups that have been implemented (`Z(n)`,`S(n)`,`D(k)`,`SO(n)`, `O(n)`,`Sp(n)`,`SO13()`,`O13()`,`SU(n)`).
+We use pytest and some of the tests are automatically generated. Because there is a large amount of tests and it can take quite some time to run them all,
+you can run a subset using pytests built in features to filter by the matches on the name of the testcase using the `-k` argument.
+
+For example to run `test_prod` with all the groups you can run
+```python emlp/tests/equivariant_subspaces_tests.py -k "test_prod"```
+
+To run the test case for a specific group could use the filter `-k "test_prod_O(3)"` and to run all tests with that group
+you could run
+```python emlp/tests/equivariant_subspaces_tests.py -k "O(3)"```
+
+The usual pytest command line arguments apply (like `-v` for verbose) and we add an additional `--log` argument for the log level.