Merge pull request #2 from mfinzi/reorganized

Reorganization of file structure and imports in the package.

README.md

@@ -15,16 +15,16 @@ Use at your own caution. But if you notice things behaving unexpectedly or get f
 
 --------------------------------------------------------------------------------
 
-Our type system is centered on it making it easy to combine representations using ρᵤ⊗ρᵥ, ρᵤ⊕ρᵥ, ρ*. For any given matrix group and representation formed in our type system, you can get the equivariant basis with [`rep.symmetric_basis()`](https://emlp.readthedocs.io/en/latest/package/emlp.solver.representation.html#emlp.solver.representation.symmetric_basis) or a matrix which projects to that subspace with [`rep.symmetric_projector()`](https://emlp.readthedocs.io/en/latest/package/emlp.solver.representation.html#emlp.solver.representation.symmetric_projector). For example:
+Our type system is centered on it making it easy to combine representations using ρᵤ⊗ρᵥ, ρᵤ⊕ρᵥ, ρ*. For any given matrix group and representation formed in our type system, you can get the equivariant basis with [`rep.equivariant_basis()`](https://emlp.readthedocs.io/en/latest/package/emlp.solver.representation.html#emlp.solver.representation.equivariant_basis) or a matrix which projects to that subspace with [`rep.equivariant_projector()`](https://emlp.readthedocs.io/en/latest/package/emlp.solver.representation.html#emlp.solver.representation.equivariant_projector). For example:
 
 ```python
-from emlp.solver.representation import V
-from emlp.solver.groups import *
+from emlp.reps import V
+from emlp.groups import *
 
 W=V(O13())
 repin = (W+2*W**2)*(W.T+1*W).T + W.T
 repout = 3*W**0 + W + W*W.T
-Q = (repin>>repout).symmetric_basis()
+Q = (repin>>repout).equivariant_basis()
 ```
 
 is code that will run and produce the basis for linear maps from repin to repout that are equivariant to the Lorentz group O(1,3).
@@ -33,10 +33,10 @@ You can even mix and match representations from different groups. For example wi
 
 ```python
 rep = 2*V(Z(3))*V(S(4))+V(O(3))**2
-Q = (rep>>rep).symmetric_basis()
+Q = (rep>>rep).equivariant_basis()
 ```
 
-You can visualize these equivariant bases with [`vis(repin,repout)`](https://emlp.readthedocs.io/en/latest/package/emlp.solver.representation.html#emlp.solver.representation.vis), such as with the two examples above
+You can visualize these equivariant bases with [`vis(repin,repout)`](https://emlp.readthedocs.io/en/latest/package/emlp.reps.html#emlp.reps.vis), such as with the two examples above
 
 <img src="https://user-images.githubusercontent.com/12687085/111226517-a2192b80-85b7-11eb-8dba-c01399fb7105.png" width="350"/> <img src="https://user-images.githubusercontent.com/12687085/111226510-a0e7fe80-85b7-11eb-913b-09776cdaa92e.png" width="230"/>  
 <!-- ![basis B](https://user-images.githubusercontent.com/12687085/111226517-a2192b80-85b7-11eb-8dba-c01399fb7105.png "title2")
@@ -59,21 +59,35 @@ pip install -e .
 
 # Experimental Results from Paper
 
-Assuming you have installed the repo locally, you can run the experiments we described in the paper. The relevant scripts are
-[`emlp/experiments/neuralode.py`](https://github.com/mfinzi/equivariant-MLP/blob/master/emlp/experiments/neuralode.py) to train (equivariant) Neural ODEs, [`emlp/experiments/hnn.py`](https://github.com/mfinzi/equivariant-MLP/blob/master/emlp/experiments/hnn.py) to train (equivariant) Hamiltonian Neural Networks, [`emlp/experiments/inertia.py`](https://github.com/mfinzi/equivariant-MLP/blob/master/emlp/experiments/inertia.py) to train (E)MLP on the inertia task, [`emlp/experiments/o5_synthetic.py`](https://github.com/mfinzi/equivariant-MLP/blob/master/emlp/experiments/o5_synthetic.py) for the O(5)-Invariant synthetic task, and [`emlp/experiments/scattering.py`](https://github.com/mfinzi/equivariant-MLP/blob/master/emlp/experiments/scattering.py) to train (E)MLP on the Lorentz invariant scattering task. If in doubt, you can see options for command line arguments for each of these scripts by running `python script.py -h`.
+Assuming you have installed the repo locally, you can run the experiments we described in the paper. 
 
-For the three regression tasks, you can run [EMLP](https://emlp.readthedocs.io/en/latest/package/emlp.models.mlp.html#emlp.models.EMLP) or the baseline with `--network EMLP` or `--network MLP` to train the baseline. For example:
+To train the regression models on one of the `Inertia`, `O5Synthetic`, or `ParticleInteraction` datasets found in [`emlp.datasets.py`](https://github.com/mfinzi/equivariant-MLP/blob/master/emlp/datasets.py) you can run the script [`experiments/train_regression.py`](https://github.com/mfinzi/equivariant-MLP/blob/master/experiments/train_regression.py) with command line arguments specifying the dataset, network, and symmetry group. For example to train [`EMLP`](https://emlp.readthedocs.io/en/latest/package/emlp.nn.html#emlp.nn.EMLP) with [`SO(3)`](https://emlp.readthedocs.io/en/latest/package/emlp.groups.html#emlp.groups.SO) equivariance on the `Inertia` dataset, you can run
 
 ```
-python emlp/experiments/scattering.py --network EMLP
+python experiments/train_regression.py --dataset Inertia --network EMLP --group "SO(3)"
 ```
 
-For the dynamical system task, the Neural ODE and HNN models have special names. [`EMLPode`](https://emlp.readthedocs.io/en/latest/package/emlp.models.mlp.html#emlp.models.EMLPode) and [`MLPode`](https://emlp.readthedocs.io/en/latest/package/emlp.models.mlp.html#emlp.models.MLPode) for the Neural ODEs in `neuralode.py` and [`EMLPH`](https://emlp.readthedocs.io/en/latest/package/emlp.models.mlp.html#emlp.models.EMLPH) and [`MLPH`](https://emlp.readthedocs.io/en/latest/package/emlp.models.mlp.html#emlp.models.MLPH) for the HNNs in `hnn.py`. For example,
+or to train the MLP baseline you can run
+
 ```
-python emlp/experiments/neuralode.py --network EMLPode
+python experiments/train_regression.py --dataset Inertia --network MLP
 ```
+Other command line arguments such as `--aug=True` for data augmentation or `--ch=512` for number of hidden units and others are available, and you can browse the options and their defaults with `python experiments/train_regression.py -h`. If no group is specified, EMLP will automatically choose the one matched to the dataset, but you can also go crazy with any of the other groups implemented in [`groups.py`](https://github.com/mfinzi/equivariant-MLP/blob/master/emlp/groups.py) provided the dimensions match the data (e.g. you could do `--group=` [`"Z(3)"`](https://emlp.readthedocs.io/en/latest/package/emlp.groups.html#emlp.groups.Z) or [`"DkeR3(3)"`](https://emlp.readthedocs.io/en/latest/package/emlp.groups.html#emlp.groups.DkeR3) but not [`"Sp(2)"`](https://emlp.readthedocs.io/en/latest/package/emlp.groups.html#emlp.groups.Sp) or [`"SU(5)"`](https://emlp.readthedocs.io/en/latest/package/emlp.groups.html#emlp.groups.SU)).
+
+For the dynamical systems modeling experiments you can use the scripts
+[`experiments/neuralode.py`](https://github.com/mfinzi/equivariant-MLP/blob/master/experiments/neuralode.py) to train (equivariant) Neural ODEs and [`experiments/hnn.py`](https://github.com/mfinzi/equivariant-MLP/blob/master/experiments/hnn.py) to train (equivariant) Hamiltonian Neural Networks.
+
 
-You can also train EMLP using other equivariance groups provided the dimensions match the data. To do this, specify a group from [`groups.py`](https://github.com/mfinzi/equivariant-MLP/blob/master/emlp/solver/groups.py) in quotations, for example you could swap out [SO(3)](https://emlp.readthedocs.io/en/latest/package/emlp.solver.groups.html#emlp.solver.groups.SO) in `inertia.py` with the permutation group [S(3)](https://emlp.readthedocs.io/en/latest/package/emlp.solver.groups.html#emlp.solver.groups.S) using the command line argument `--group "S(3)"`. However, other groups will likely be less relevant to these specific problems.
+For the dynamical system task, the Neural ODE and HNN models have special names. [`EMLPode`](https://emlp.readthedocs.io/en/latest/package/emlp.nn.html#emlp.nn.EMLPode) and [`MLPode`](https://emlp.readthedocs.io/en/latest/package/emlp.nn.html#emlp.nn.MLPode) for the Neural ODEs in `neuralode.py` and [`EMLPH`](https://emlp.readthedocs.io/en/latest/package/emlp.nn.html#emlp.nn.EMLPH) and [`MLPH`](https://emlp.readthedocs.io/en/latest/package/emlp.nn.html#emlp.nn.MLPH) for the HNNs in `hnn.py`. For example,
+
+```
+python experiments/neuralode.py --network EMLPode --group="O2eR3()"
+```
+or 
+
+```
+python experiments/hnn.py --network EMLPH --group="DkeR3(6)"
+```
 
 <!-- # 
 <p align="center">

docs/CHANGELOG.md

@@ -14,5 +14,13 @@ Remember to align the itemized text with the first line of an item within a list
 * Bug Fixes:
   * Element ordering of mixed groups is now correctly maintained in the solution
   * Fixed edge case of {func}`lazy_direct_matmat` when concatenating matrices of size 0 
-    affecting {func}`emlp.solver.representation.Rep.symmetric_basis` but not 
-    {func}`emlp.solver.representation.Rep.symmetric_projector`
+    affecting {func}`emlp.reps.Rep.equivariant_basis` but not 
+    {func}`emlp.reps.Rep.equivariant_projector`
+* API Changes:
+  * `emlp.solver.representation` -> `emlp.reps`
+  * `emlp.solver.groups` -> `emlp.groups`
+  * `emlp.models.mlp` -> `emlp.nn`
+  * `rep.symmetric_basis()` -> `rep.equivariant_basis()`
+  * `rep.symmetric_projector()` -> `rep.equivariant_projector()`
+  * Tests and experiments separated from package and api
+

docs/conf.py

@@ -25,6 +25,10 @@
 copyright = '2021, Marc Finzi'
 author = 'Marc Finzi'
 
+# import emlp
+# release = emlp.__version__
+
+
 sys.path.append(os.path.abspath('sphinxext'))
 extensions = [
     'nbsphinx',
@@ -37,6 +41,7 @@
     'sphinx.ext.viewcode',
     'matplotlib.sphinxext.plot_directive',
     'sphinx_autodoc_typehints',
+    'sphinx.ext.autosectionlabel',
 ]
 autosummary_generate = True
 autodoc_default_options = {'autosummary': True}

docs/documentation.md

@@ -15,13 +15,10 @@ 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.
 
 ### Editing ipynb
 
-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,
+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` (for editing and running in Colab you will need to add 
 `!pip install git+https://github.com/mfinzi/equivariant-MLP.git`).

docs/index.rst

@@ -36,7 +36,7 @@ and representations an easy task, and one that does not require knowledge of ana
 .. toctree::
    :glob:
    :maxdepth: 1
-   :caption: Package Reference
+   :caption: API Reference
 
    package/emlp.solver.groups
    package/emlp.solver.representation