├── .conda.yml ├── .gitignore ├── .readthedocs.yml ├── .travis.yml ├── LICENSE ├── README.md ├── doc ├── Makefile ├── make.bat ├── requirements.txt └── source │ ├── _static │ └── css │ │ └── custom.css │ ├── conf.py │ ├── index.rst │ ├── modules.rst │ ├── neurophox-logo.png │ ├── neurophox.components.rst │ ├── neurophox.ml.rst │ ├── neurophox.numpy.rst │ ├── neurophox.rst │ ├── neurophox.tensorflow.rst │ └── neurophox.torch.rst ├── neurophox ├── __init__.py ├── components │ ├── __init__.py │ ├── mzi.py │ └── transfermatrix.py ├── config.py ├── decompositions.py ├── helpers.py ├── initializers.py ├── meshmodel.py ├── ml │ ├── __init__.py │ ├── linear.py │ └── nonlinearities.py ├── numpy │ ├── __init__.py │ ├── generic.py │ └── layers.py ├── tensorflow │ ├── __init__.py │ ├── generic.py │ └── layers.py └── torch │ ├── __init__.py │ ├── generic.py │ └── layers.py ├── requirements.txt ├── setup.py └── tests ├── __init__.py ├── test_components.py ├── test_decompositions.py ├── test_fieldpropagation.py ├── test_layers.py ├── test_phasecontrolfn.py ├── test_svd.py └── test_transformers.py /.conda.yml: -------------------------------------------------------------------------------- 1 | name: neurophox 2 | channels: 3 | - pytorch-nightly 4 | - conda-forge 5 | dependencies: 6 | - python=3.7 # or 2.7 7 | - numpy 8 | - scipy=1.4.1 9 | - tensorflow=2.2.* 10 | - pytorch=1.7.* 11 | - sphinx-autodoc-typehints -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.py? 2 | *~ 3 | *.swp 4 | .cache 5 | __pycache__ 6 | *.egg-info 7 | .env 8 | .env* 9 | .idea 10 | .DS_Store 11 | doc/build/* -------------------------------------------------------------------------------- /.readthedocs.yml: -------------------------------------------------------------------------------- 1 | version: 2 2 | 3 | conda: 4 | environment: .conda.yml 5 | 6 | python: 7 | version: 3.7 8 | install: 9 | - requirements: doc/requirements.txt 10 | - method: setuptools 11 | path: . 12 | 13 | sphinx: 14 | builder: html 15 | configuration: doc/source/conf.py -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | sudo: false 2 | language: python 3 | python: 3.7 4 | 5 | env: 6 | - PYTORCH_CHANNEL=pytorch-nightly 7 | 8 | before_install: &before_install 9 | - sudo apt-get update 10 | - wget https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh; 11 | - bash miniconda.sh -b -p $HOME/miniconda 12 | - export PATH="$HOME/miniconda/bin:$PATH" 13 | - hash -r 14 | - conda config --set always_yes yes --set changeps1 no 15 | - conda update -q conda 16 | # Useful for debugging any issues with conda 17 | - conda info -a 18 | - conda create -q -n test-environment pytorch cpuonly python=$TRAVIS_PYTHON_VERSION -c $PYTORCH_CHANNEL 19 | - source activate test-environment 20 | - conda install numpy scipy 21 | - conda install tensorflow 22 | 23 | install: 24 | pip install codecov && pip install pytest 25 | 26 | script: 27 | coverage run -m unittest discover 28 | after_success: 29 | - codecov 30 | deploy: 31 | provider: pypi 32 | user: sunilpai 33 | password: 34 | secure: QdY7eX8n/C4Nb7Lm/CvwO4vNltTaeVS0XuJV3kA4fvuFNCtXkfi7Bg85lkpHX/IPhe0A8SR3IT4UdwpSYR7iKPw0iM3/+eboaGuQREQ59Vi0+X85oY4oYnebJFMZ61FhHFwBmVi1pQ6IgNqNqLmG+y1ni3Op2ZDpECabjIjnEhr5SRD0j98DIT9K29QKdbBTokLf0Xb7jMXtYn+dU8401eFv+NrCXFUm6xQ0C7MPPaWTi0QxWHBkoc8acA9t20QlJoj63EufJOFA8nKSxzp3Phmzstobj/82vshl9glDb1tHKKyw/Y+zL3JVaNneYKcoR5CGThfqhsi8MklGZRGWB2/LiLlHA52u8UGfFBFk4IWhTbM6LQCYYCofzbQBsDjDvbP3W/20oGz2Oay8vD4pXEUjjjb772tdNrCeT+QFOYNAF07P1UcgHceVw7fGO0w2grniIqqlNmCcqfSBOFZtQLp7fhE/FtLxlrNZ39OLFG8LU39+f+X7U3bc1hnqfgWMlFQK0ofwq2GSjk7AZpMuobLlJXNkXUYpMPaejhIBTzVLz3AsDeU5+nC0Qji+LM6IIWnLzMvhLm2zm3cI9l9EUpi9zqYmSaKu3W8OPLgTUZ7u9cyTVarMCI7u3N88CUFsxlBx4hch6K4jnolWUH7JakkYOf2N/Ix6Uiqo/lgBAao= 35 | on: 36 | branch: master 37 | tags: true 38 | python: 3.7 -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 Solgaard Lab 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 |
Logo
2 | 3 | # 4 | ![Build Status](https://img.shields.io/travis/solgaardlab/neurophox/master.svg?style=for-the-badge) 5 | ![Docs](https://readthedocs.org/projects/neurophox/badge/?style=for-the-badge) 6 | ![PiPy](https://img.shields.io/pypi/v/neurophox.svg?style=for-the-badge) 7 | ![CodeCov](https://img.shields.io/codecov/c/github/solgaardlab/neurophox/master.svg?style=for-the-badge) 8 | 9 | 10 | The [Neurophox module](https://neurophox.readthedocs.io) is an open source machine learning and photonic simulation framework based on unitary mesh networks presented in [arxiv/1909.06179](https://arxiv.org/pdf/1909.06179.pdf), [arxiv/1808.00458](https://arxiv.org/pdf/1808.00458.pdf), and [arxiv/1903.04579](https://arxiv.org/pdf/1903.04579.pdf). 11 | 12 | ![neurophox](https://user-images.githubusercontent.com/7623867/57964658-87a79580-78ed-11e9-8f1e-c4af30c32e65.gif) 13 | 14 | 15 | ![neurophox](https://user-images.githubusercontent.com/7623867/57976056-57fb9a80-798c-11e9-9aef-8d1f07af7ca7.gif) 16 | 17 | 18 | ## Motivation 19 | 20 | Neurophox provides a robust and efficient framework for simulating **optical neural networks** (ONNs) that promise fast and energy-efficient machine learning. Scalable ONNs are made possible by integrated [**reconfigurable nanophotonic processors**](https://www.osapublishing.org/optica/abstract.cfm?uri=optica-5-12-1623), feedforward networks of 2 x 2 nodes that compute matrix multiplication by simply allowing light to flow through them (thus, a time- and energy-efficient way to compute). 21 | 22 | Reconfigurable photonic networks are currently being [developed to scale](https://medium.com/lightmatter/matrix-processing-with-nanophotonics-998e294dabc1) in the optical domain using Mach-Zehnder interferometers (MZIs). Such matrix multiplier processors behave differently from conventional matrix multipliers: 23 | 1. They act as **unitary operators rather than general linear operators**, preserving the norm of the data flowing through the network (i.e., intensity of the light is preserved due to energy conservation). 24 | 2. The matrix elements are not directly trained during backpropagation. Instead, the node parameters are trained, and the **matrix elements are nonlinear functions of those node parameters**. 25 | 26 | The goal of Neurophox is to provide utilities for calibrating and training such processors for machine learning and other signal processing tasks. These simulated models might furthermore be useful for a relatively unexplored regime of machine learning: "unitary machine learning" (see e.g., [arxiv/1612.05231](https://arxiv.org/pdf/1612.05231.pdf)). 27 | 28 | ## Dependencies and requirements 29 | 30 | Some important requirements for Neurophox are: 31 | 1. Python >=3.6 32 | 2. Tensorflow >=2.0.0a 33 | 3. PyTorch 34 | 35 | The dependencies for Neurophox (specified in `requirements.txt`) are: 36 | ```text 37 | numpy>=1.16 38 | scipy 39 | matplotlib 40 | tensorflow>=2.0.0a 41 | ``` 42 | 43 | Optionally, the user may install `torch>=1.7` to run the Pytorch layers. 44 | NOTE: PyTorch layers are currently nonfunctional as we migrate to complex PyTorch, but will be up and running once this 45 | [issue](https://github.com/pytorch/pytorch/issues/53645) is resolved. 46 | 47 | ## Getting started 48 | 49 | ### Installation 50 | 51 | There are currently two options to install Neurophox: 52 | 1. Installation via `pip`: 53 | ```bash 54 | pip install neurophox 55 | ``` 56 | 2. Installation from source via `pip`: 57 | ```bash 58 | git clone https://github.com/solgaardlab/neurophox 59 | pip install -e neurophox 60 | pip install -r neurophox/requirements.txt 61 | ``` 62 | 63 | If installing other dependencies manually, ensure you install [PyTorch](https://pytorch.org/) (since PyTorch mesh layers are currently in development) and [Tensorflow 2.0](https://www.tensorflow.org/versions/r2.0/api_docs/python/tf). 64 | 65 | #### Using the GPU 66 | 67 | If using a GPU, we recommend using a `conda` environement to install GPU dependencies using CUDA 10.0 with the following commands: 68 | ```bash 69 | conda install pytorch torchvision cudatoolkit=10.2 -c pytorch-nightly 70 | pip install tensorflow-gpu==2.2.0 71 | ``` 72 | 73 | ### Features 74 | 75 | Using Neurophox, we can simulate light physically propagating through such networks (using `layer.propagate`), and observe its equivalence to a _matrix operation_ (using `layer.matrix`). 76 | 77 | We demonstrate the **R**ectangular **M**esh (RM layer) using either `numpy` or `tensorflow`: 78 | 79 | ```python 80 | import numpy as np 81 | from neurophox.numpy import RMNumpy 82 | from neurophox.tensorflow import RM 83 | 84 | N = 16 85 | 86 | tf_layer = RM(N) 87 | np_layer = RMNumpy(N, phases=tf_layer.phases) 88 | 89 | np.allclose(tf_layer.matrix, np_layer.matrix) # True 90 | np.allclose(tf_layer(np.eye(N)), np_layer.matrix) # True 91 | ``` 92 | 93 | ### Visualizations 94 | 95 | We provide code to visualize network phase shift parameters and light propagation in [`neurophox-notebooks`](https://github.com/solgaardlab/neurophox-notebooks) for rectangular and triangular meshes. 96 | 97 | Rectangular mesh: 98 | ![neurophox](https://user-images.githubusercontent.com/7623867/57964850-aeb39680-78f0-11e9-8785-e6e46c705b34.png) 99 | Triangular mesh: 100 | ![neurophox](https://user-images.githubusercontent.com/7623867/57964852-aeb39680-78f0-11e9-8a5c-d08e9f6dce89.png) 101 | 102 | For the **phase shift settings** above, we visualize the **propagation of light** (scaled by light intensity), equivalently representing data "flowing" through the mesh. 103 | 104 | Rectangular mesh: 105 | ![neurophox](https://user-images.githubusercontent.com/7623867/57964851-aeb39680-78f0-11e9-9ff3-41e8cebd25a6.png) 106 | Triangular mesh: 107 | ![neurophox](https://user-images.githubusercontent.com/7623867/57964853-aeb39680-78f0-11e9-8cd4-1364d2cec339.png) 108 | 109 | ### Small machine learning example 110 | 111 | It is possible to compose Neurophox Tensorflow layers into unitary neural networks using `tf.keras.Sequential` 112 | to solve machine learning problems. Here we use absolute value nonlinearities and categorical cross entropy. 113 | 114 | ```python 115 | import tensorflow as tf 116 | from neurophox.tensorflow import RM 117 | from neurophox.ml.nonlinearities import cnorm, cnormsq 118 | 119 | ring_model = tf.keras.Sequential([ 120 | RM(3, activation=tf.keras.layers.Activation(cnorm)), 121 | RM(3, activation=tf.keras.layers.Activation(cnorm)), 122 | RM(3, activation=tf.keras.layers.Activation(cnorm)), 123 | RM(3, activation=tf.keras.layers.Activation(cnorm)), 124 | tf.keras.layers.Activation(cnormsq), 125 | tf.keras.layers.Lambda(lambda x: tf.math.real(x[:, :2])), # get first 2 output ports (we already know it is real from the activation), 126 | tf.keras.layers.Activation('softmax') 127 | ]) 128 | 129 | ring_model.compile( 130 | loss='categorical_crossentropy', 131 | optimizer=tf.keras.optimizers.Adam(lr=0.0025) 132 | ) 133 | ``` 134 | 135 | Below is a visualization for many planar classification problems: 136 | 137 | ![neurophox](https://user-images.githubusercontent.com/7623867/58128090-b2cf0500-7bcb-11e9-8986-25450bfd68a9.png) 138 | ![neurophox](https://user-images.githubusercontent.com/7623867/58132218-95069d80-7bd5-11e9-9d08-20e1de5c3727.png) 139 | The code to generate the above example is provided in [`neurophox-notebooks`](https://github.com/solgaardlab/neurophox-notebooks). 140 | 141 | 142 | ## Authors and citing this repository 143 | Neurophox was written by Sunil Pai (email: sunilpai@stanford.edu). 144 | 145 | If you find this repository useful, please cite at least one of the following papers depending on your application: 146 | 1. Definition and calibration of feedforward photonic networks: 147 | ```text 148 | @article{pai2019parallel, 149 | title={Parallel fault-tolerant programming of an arbitrary feedforward photonic network}, 150 | author={Pai, Sunil and Williamson, Ian AD and Hughes, Tyler W and Minkov, Momchil and Solgaard, Olav and Fan, Shanhui and Miller, David AB}, 151 | journal={arXiv preprint arXiv:1909.06179}, 152 | year={2019} 153 | } 154 | ``` 155 | 2. Optimization of unitary mesh networks: 156 | ```text 157 | @article{pai_matrix_2019, 158 | author = {Pai, Sunil and Bartlett, Ben and Solgaard, Olav and Miller, David A. B.}, 159 | doi = {10.1103/PhysRevApplied.11.064044}, 160 | journal = {Physical Review Applied}, 161 | month = jun, 162 | number = {6}, 163 | pages = {064044}, 164 | title = {Matrix Optimization on Universal Unitary Photonic Devices}, 165 | volume = {11}, 166 | year = {2019} 167 | } 168 | ``` 169 | 3. Optical neural network nonlinearities: 170 | ```text 171 | @article{williamson_reprogrammable_2020, 172 | author = {Williamson, I. A. D. and Hughes, T. W. and Minkov, M. and Bartlett, B. and Pai, S. and Fan, S.}, 173 | doi = {10.1109/JSTQE.2019.2930455}, 174 | issn = {1077-260X}, 175 | journal = {IEEE Journal of Selected Topics in Quantum Electronics}, 176 | month = jan, 177 | number = {1}, 178 | pages = {1-12}, 179 | title = {Reprogrammable Electro-Optic Nonlinear Activation Functions for Optical Neural Networks}, 180 | volume = {26}, 181 | year = {2020} 182 | } 183 | ``` 184 | 185 | ## Future Work and Contributions 186 | 187 | Neurophox is under development and is not yet stable. 188 | 189 | If you find a bug, have a question, or would like to recommend a feature, please submit an issue on Github. 190 | 191 | We welcome pull requests and contributions from the broader community. If you would like to contribute, please submit a pull request and title your branch `bug/bug-fix-title` or `feature/feature-title`. 192 | 193 | 194 | -------------------------------------------------------------------------------- /doc/Makefile: -------------------------------------------------------------------------------- 1 | # Minimal makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line. 5 | SPHINXOPTS = 6 | SPHINXBUILD = sphinx-build 7 | SOURCEDIR = source 8 | BUILDDIR = build 9 | 10 | # Put it first so that "make" without argument is like "make help". 11 | help: 12 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 13 | 14 | .PHONY: help Makefile 15 | 16 | # Catch-all target: route all unknown targets to Sphinx using the new 17 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). 18 | %: Makefile 19 | @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) -------------------------------------------------------------------------------- /doc/make.bat: -------------------------------------------------------------------------------- 1 | @ECHO OFF 2 | 3 | pushd %~dp0 4 | 5 | REM Command file for Sphinx documentation 6 | 7 | if "%SPHINXBUILD%" == "" ( 8 | set SPHINXBUILD=sphinx-build 9 | ) 10 | set SOURCEDIR=source 11 | set BUILDDIR=build 12 | 13 | if "%1" == "" goto help 14 | 15 | %SPHINXBUILD% >NUL 2>NUL 16 | if errorlevel 9009 ( 17 | echo. 18 | echo.The 'sphinx-build' command was not found. Make sure you have Sphinx 19 | echo.installed, then set the SPHINXBUILD environment variable to point 20 | echo.to the full path of the 'sphinx-build' executable. Alternatively you 21 | echo.may add the Sphinx directory to PATH. 22 | echo. 23 | echo.If you don't have Sphinx installed, grab it from 24 | echo.http://sphinx-doc.org/ 25 | exit /b 1 26 | ) 27 | 28 | %SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% 29 | goto end 30 | 31 | :help 32 | %SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% 33 | 34 | :end 35 | popd 36 | -------------------------------------------------------------------------------- /doc/requirements.txt: -------------------------------------------------------------------------------- 1 | sphinx-autodoc-typehints 2 | sphinx-autodoc-annotation -------------------------------------------------------------------------------- /doc/source/conf.py: -------------------------------------------------------------------------------- 1 | # Configuration file for the Sphinx documentation builder. 2 | # 3 | # This file only contains a selection of the most common options. For a full 4 | # list see the documentation: 5 | # http://www.sphinx-doc.org/en/master/config 6 | 7 | # -- Path setup -------------------------------------------------------------- 8 | 9 | # If extensions (or modules to document with autodoc) are in another directory, 10 | # add these directories to sys.path here. If the directory is relative to the 11 | # documentation root, use os.path.abspath to make it absolute, like shown here. 12 | # 13 | import os 14 | import sys 15 | sys.path.insert(0, os.path.abspath('../../')) 16 | 17 | 18 | # -- Project information ----------------------------------------------------- 19 | 20 | project = 'Neurophox' 21 | copyright = '2019, Sunil Pai' 22 | author = 'Sunil Pai' 23 | 24 | 25 | # -- General configuration --------------------------------------------------- 26 | master_doc = 'index' 27 | 28 | # Add any Sphinx extension module names here, as strings. They can be 29 | # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom 30 | # ones. 31 | extensions = [ 32 | 'sphinx.ext.napoleon', 33 | 'sphinx.ext.mathjax', 34 | 'sphinx.ext.autodoc', 35 | 'sphinx_autodoc_typehints', 36 | 'sphinx.ext.viewcode', 37 | 'sphinx.ext.inheritance_diagram' 38 | ] 39 | 40 | mathjax_path = "https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-MML-AM_CHTML" 41 | 42 | # mathjax_config = { 43 | # 'extensions': ['tex2jax.js'], 44 | # 'jax': ['input/TeX', 'output/CommonHTML'], 45 | # 'CommonHTML': {} 46 | # } 47 | 48 | # Add any paths that contain templates here, relative to this directory. 49 | templates_path = ['_templates'] 50 | 51 | # List of patterns, relative to source directory, that match files and 52 | # directories to ignore when looking for source files. 53 | # This pattern also affects html_static_path and html_extra_path. 54 | exclude_patterns = [] 55 | 56 | 57 | # -- Options for HTML output ------------------------------------------------- 58 | 59 | # The theme to use for HTML and HTML Help pages. See the documentation for 60 | # a list of builtin themes. 61 | # 62 | html_theme_path = ["_themes"] 63 | html_theme = 'sphinx_rtd_theme' 64 | html_theme_options = { 65 | "logo_only": True 66 | } 67 | 68 | # Add any paths that contain custom static files (such as style sheets) here, 69 | # relative to this directory. They are copied after the builtin static files, 70 | # so a file named "default.css" will overwrite the builtin "default.css". 71 | html_static_path = ['_static'] 72 | html_logo = "neurophox-logo.png" 73 | 74 | 75 | def setup(app): 76 | app.add_stylesheet('css/custom.css') 77 | -------------------------------------------------------------------------------- /doc/source/index.rst: -------------------------------------------------------------------------------- 1 | .. neurophox documentation master file, created by 2 | sphinx-quickstart on Sun May 19 10:27:10 2019. 3 | You can adapt this file completely to your liking, but it should at least 4 | contain the root `toctree` directive. 5 | 6 | Neurophox 7 | ========= 8 | 9 | Neurophox is an open source machine learning and photonic simulation framework based on unitary mesh networks presented in `arxiv/1808.00458 `_ and `arxiv/1903.04579 `_. 10 | 11 | `Note`: ``neurophox`` is under active development and is not yet stable, so expect to see the documentation change frequently. 12 | 13 | .. image:: https://user-images.githubusercontent.com/7623867/57964658-87a79580-78ed-11e9-8f1e-c4af30c32e65.gif 14 | :target: https://user-images.githubusercontent.com/7623867/57964658-87a79580-78ed-11e9-8f1e-c4af30c32e65.gif 15 | :alt: neurophox 16 | :width: 100% 17 | 18 | 19 | 20 | .. image:: https://user-images.githubusercontent.com/7623867/57976056-57fb9a80-798c-11e9-9aef-8d1f07af7ca7.gif 21 | :target: https://user-images.githubusercontent.com/7623867/57976056-57fb9a80-798c-11e9-9aef-8d1f07af7ca7.gif 22 | :alt: neurophox 23 | :width: 100% 24 | 25 | 26 | Motivation 27 | ---------- 28 | 29 | Orthogonal and unitary neural networks have interesting properties and have been studied for synthetic natural language processing tasks (see `unitary mesh-based RNN `_\ , `unitary evolution RNN `_\ , and `orthogonal evolution RNN `_\ ). Furthermore, new energy-efficient photonic technologies are being built to realize unitary mesh-based neural networks using light as the computing medium as opposed to conventional analog electronics. 30 | 31 | Introduction 32 | ------------ 33 | 34 | Neurophox provides a general framework for mesh network layers in orthogonal and unitary neural networks. We use an efficient definition for any feedforward mesh architecture, ``neurophox.meshmodel.MeshModel``\ , to develop mesh layer architectures in Numpy (\ ``neurophox.numpy.layers``\ ), `Tensorflow 2.0 `_ (\ ``neurophox.tensorflow.layers``\ ), and (soon) `PyTorch `_. 35 | 36 | Scattering matrix models used in unitary mesh networks for photonics simulations are provided in ``neurophox.components``. The models for all layers are fully defined in ``neurophox.meshmodel``\ , which provides a general framework for efficient implementation of any unitary mesh network. 37 | 38 | 39 | Dependencies and requirements 40 | ----------------------------- 41 | 42 | Some important requirements for ``neurophox`` are: 43 | 44 | 45 | #. Python >=3.6 46 | #. Tensorflow >=2.0 47 | #. PyTorch >=1.3 48 | 49 | The dependencies for ``neurophox`` (specified in ``requirements.txt``\ ) are: 50 | 51 | .. code-block:: text 52 | 53 | numpy>=1.16 54 | scipy 55 | matplotlib 56 | tensorflow>=2.0.0a 57 | 58 | The user may also optionally install ``torch>=1.3`` to run the ``neurophox.torch`` module. 59 | 60 | Getting started 61 | --------------- 62 | 63 | Installation 64 | ^^^^^^^^^^^^ 65 | 66 | There are currently two options to install ``neurophox``\ : 67 | 68 | 69 | #. Installation via ``pip``\ : 70 | 71 | .. code-block:: bash 72 | 73 | pip install neurophox 74 | 75 | #. Installation from source via ``pip``\ : 76 | 77 | .. code-block:: bash 78 | 79 | git clone https://github.com/solgaardlab/neurophox 80 | pip install -e neurophox 81 | pip install -r requirements.txt 82 | 83 | If installing other dependencies manually, ensure you install `PyTorch `_ (since PyTorch mesh layers are currently in development) and `Tensorflow 2.0 `_. 84 | 85 | Using the GPU 86 | ~~~~~~~~~~~~~ 87 | 88 | If using a GPU, we recommend using a ``conda`` environement to install GPU dependencies using CUDA 10.0 with the following commands: 89 | 90 | .. code-block:: bash 91 | 92 | conda install pytorch torchvision cudatoolkit=10.0 -c pytorch 93 | pip install tensorflow-gpu==2.0.0-alpha0 94 | 95 | Imports 96 | ^^^^^^^ 97 | 98 | .. code-block:: python 99 | 100 | import numpy as np 101 | from neurophox.numpy import RMNumpy 102 | from neurophox.tensorflow import RM 103 | 104 | N = 16 105 | 106 | tf_layer = RM(N) 107 | np_layer = RMNumpy(N, phases=tf_layer.phases) 108 | 109 | np.allclose(tf_layer.matrix, np_layer.matrix) # True 110 | np.allclose(tf_layer(np.eye(N)), np_layer.matrix) # True 111 | 112 | 113 | Inspection 114 | ^^^^^^^^^^ 115 | 116 | We can inspect the parameters for each layer using ``neurophox.control.MeshPhases`` which can be accessed via ``tf_layer.phases`` and ``np_layer.phases``. 117 | 118 | We can inspect the matrix elements implemented by each layer as follows via ``tf_layer.matrix`` and ``np_layer.matrix``. 119 | 120 | 121 | 122 | Bottleneck 123 | ^^^^^^^^^^ 124 | The bottleneck in time for a ``neurophox`` mesh layer is `not` dominated by the number of inputs/outputs (referred to as ``units`` or :math:`N`), but rather the number of "vertical layers" in the mesh layer (referred to as ``num_layers`` or :math:`L`). This is to say that fairly large unitary matrices can be implemented by ``neurophox`` as long as the ``num_layers`` parameter is kept small (e.g., empirically, :math:`\log N` layers seems to be enough for butterfly mesh architectures used in `unitary mesh-based RNNs `_\ ). 125 | 126 | 127 | Visualizations 128 | -------------- 129 | 130 | 131 | Phase shift settings visualization 132 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 133 | 134 | The phase shift patterns used to generate the above propagation patterns can also be visualized by plotting ``np_layer.phases``\ : 135 | 136 | Rectangular mesh: 137 | 138 | .. image:: https://user-images.githubusercontent.com/7623867/57964850-aeb39680-78f0-11e9-8785-e6e46c705b34.png 139 | :target: https://user-images.githubusercontent.com/7623867/57964850-aeb39680-78f0-11e9-8785-e6e46c705b34.png 140 | :alt: neurophox 141 | :width: 100% 142 | 143 | Triangular mesh: 144 | 145 | .. image:: https://user-images.githubusercontent.com/7623867/57964852-aeb39680-78f0-11e9-8a5c-d08e9f6dce89.png 146 | :target: https://user-images.githubusercontent.com/7623867/57964852-aeb39680-78f0-11e9-8a5c-d08e9f6dce89.png 147 | :alt: neurophox 148 | :width: 100% 149 | 150 | 151 | Light propagation visualization 152 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 153 | 154 | For the phase shift settings above, we can visualize the propagation of light (field magnitude), as the data "flows" through the mesh. 155 | 156 | Rectangular mesh: 157 | 158 | .. image:: https://user-images.githubusercontent.com/7623867/57964851-aeb39680-78f0-11e9-9ff3-41e8cebd25a6.png 159 | :target: https://user-images.githubusercontent.com/7623867/57964851-aeb39680-78f0-11e9-9ff3-41e8cebd25a6.png 160 | :alt: neurophox 161 | :width: 100% 162 | 163 | Triangular mesh: 164 | 165 | .. image:: https://user-images.githubusercontent.com/7623867/57964853-aeb39680-78f0-11e9-8cd4-1364d2cec339.png 166 | :target: https://user-images.githubusercontent.com/7623867/57964853-aeb39680-78f0-11e9-8cd4-1364d2cec339.png 167 | :alt: neurophox 168 | :width: 100% 169 | 170 | 171 | The code to generate these visualization examples are provided in `neurophox-notebooks `_\. 172 | 173 | Small machine learning example 174 | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 175 | 176 | It is simple to compose ``neurophox`` Tensorflow layers into unitary neural networks using ``Sequential`` to solve machine learning problems. 177 | 178 | 179 | .. code-block:: python 180 | 181 | import tensorflow as tf 182 | from neurophox.tensorflow import RM 183 | from neurophox.ml.nonlinearities import cnorm, cnormsq 184 | 185 | ring_model = tf.keras.Sequential([ 186 | RM(3, activation=tf.keras.layers.Activation(cnorm)), 187 | RM(3, activation=tf.keras.layers.Activation(cnorm)), 188 | RM(3, activation=tf.keras.layers.Activation(cnorm)), 189 | RM(3, activation=tf.keras.layers.Activation(cnorm)), 190 | tf.keras.layers.Activation(cnormsq), 191 | tf.keras.layers.Lambda(lambda x: tf.math.real(x[:, :2])), # get first 2 output ports, 192 | tf.keras.layers.Activation('softmax') 193 | ]) 194 | 195 | ring_model.compile( 196 | loss='categorical_crossentropy', 197 | optimizer=tf.keras.optimizers.Adam(lr=0.0025) 198 | ) 199 | 200 | Below is a visualization for many planar classification problems (including the ring model defined above): 201 | 202 | .. image:: https://user-images.githubusercontent.com/7623867/58128090-b2cf0500-7bcb-11e9-8986-25450bfd68a9.png 203 | :target: https://user-images.githubusercontent.com/7623867/58128090-b2cf0500-7bcb-11e9-8986-25450bfd68a9.png 204 | :alt: neurophox 205 | :width: 100% 206 | 207 | .. image:: https://user-images.githubusercontent.com/7623867/58132218-95069d80-7bd5-11e9-9d08-20e1de5c3727.png 208 | :target: https://user-images.githubusercontent.com/7623867/58132218-95069d80-7bd5-11e9-9d08-20e1de5c3727.png 209 | :alt: neurophox 210 | :width: 100% 211 | 212 | 213 | 214 | The code to generate the above example is provided in `neurophox-notebooks `_\. 215 | 216 | 217 | Authors and citing this repository 218 | ---------------------------------- 219 | 220 | ``neurophox`` was written by Sunil Pai (email: sunilpai@stanford.edu). 221 | 222 | If you find this repository useful, please cite at least one of the following papers depending on your application: 223 | 224 | #. Calibration of optical neural networks 225 | 226 | .. code-block:: text 227 | 228 | @article{pai2019parallel, 229 | title={Parallel fault-tolerant programming of an arbitrary feedforward photonic network}, 230 | author={Pai, Sunil and Williamson, Ian AD and Hughes, Tyler W and Minkov, Momchil and Solgaard, Olav and Fan, Shanhui and Miller, David AB}, 231 | journal={arXiv preprint arXiv:1909.06179}, 232 | year={2019} 233 | } 234 | 235 | 236 | #. Optimization of unitary mesh networks: 237 | 238 | .. code-block:: text 239 | 240 | @article{pai_matrix_2019, 241 | author = {Pai, Sunil and Bartlett, Ben and Solgaard, Olav and Miller, David A. B.}, 242 | doi = {10.1103/PhysRevApplied.11.064044}, 243 | journal = {Physical Review Applied}, 244 | month = jun, 245 | number = {6}, 246 | pages = {064044}, 247 | title = {Matrix Optimization on Universal Unitary Photonic Devices}, 248 | volume = {11}, 249 | year = {2019} 250 | } 251 | 252 | #. Optical neural network nonlinearities: 253 | 254 | .. code-block:: text 255 | 256 | @article{williamson_reprogrammable_2020, 257 | author = {Williamson, I. A. D. and Hughes, T. W. and Minkov, M. and Bartlett, B. and Pai, S. and Fan, S.}, 258 | doi = {10.1109/JSTQE.2019.2930455}, 259 | issn = {1077-260X}, 260 | journal = {IEEE Journal of Selected Topics in Quantum Electronics}, 261 | month = jan, 262 | number = {1}, 263 | pages = {1-12}, 264 | title = {Reprogrammable Electro-Optic Nonlinear Activation Functions for Optical Neural Networks}, 265 | volume = {26}, 266 | year = {2020} 267 | } 268 | 269 | Contributions 270 | ------------- 271 | 272 | ``neurophox`` is under active development and is not yet stable. 273 | 274 | If you find a bug or would like to recommend a feature, please submit an issue on Github. 275 | 276 | We welcome pull requests and contributions from the broader community. If you would like to contribute to the codebase, please submit a pull request and title your branch ``bug/bug-fix-title`` or ``feature/feature-title``. 277 | 278 | Documentation 279 | ============= 280 | 281 | ``neurophox`` is under active development and is not yet stable, so expect to see the documentation change frequently. 282 | 283 | Summary 284 | ------- 285 | 286 | The layer APIs of ``neurophox`` are provided via the subpackages: 287 | 288 | * ``neurophox.tensorflow`` 289 | * ``neurophox.numpy`` 290 | * ``neurophox.torch`` 291 | 292 | Some machine learning experiment starter code with ``neurophox`` can be found in `neurophox-notebooks `_\ and/or ``neurophox.ml``. 293 | 294 | All other subpackages have code relevant to the inner workings of ``neurophox``. 295 | 296 | 297 | Indices and search 298 | ------------------ 299 | 300 | * :ref:`genindex` 301 | * :ref:`modindex` 302 | * :ref:`search` 303 | 304 | Contents 305 | -------- 306 | 307 | .. toctree:: 308 | :hidden: 309 | 310 | index 311 | 312 | .. toctree:: 313 | neurophox 314 | -------------------------------------------------------------------------------- /doc/source/modules.rst: -------------------------------------------------------------------------------- 1 | neurophox 2 | ========= 3 | 4 | .. toctree:: 5 | :maxdepth: 4 6 | 7 | neurophox 8 | -------------------------------------------------------------------------------- /doc/source/neurophox-logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/solgaardlab/neurophox/29a9b03a4efc241a85ecb92463e32cfd69776d23/doc/source/neurophox-logo.png -------------------------------------------------------------------------------- /doc/source/neurophox.components.rst: -------------------------------------------------------------------------------- 1 | neurophox.components package 2 | ============================ 3 | 4 | neurophox.components.mzi module 5 | ------------------------------- 6 | 7 | .. automodule:: neurophox.components.mzi 8 | :members: 9 | :undoc-members: 10 | :show-inheritance: 11 | 12 | neurophox.components.transfermatrix module 13 | ------------------------------------------ 14 | 15 | .. automodule:: neurophox.components.transfermatrix 16 | :members: 17 | :undoc-members: 18 | :show-inheritance: -------------------------------------------------------------------------------- /doc/source/neurophox.ml.rst: -------------------------------------------------------------------------------- 1 | neurophox.ml package 2 | ==================== 3 | 4 | neurophox.ml.linear module 5 | -------------------------- 6 | 7 | .. automodule:: neurophox.ml.linear 8 | :members: 9 | :undoc-members: 10 | :show-inheritance: 11 | 12 | neurophox.ml.nonlinearities module 13 | ---------------------------------- 14 | 15 | .. automodule:: neurophox.ml.nonlinearities 16 | :members: 17 | :undoc-members: 18 | :show-inheritance: 19 | -------------------------------------------------------------------------------- /doc/source/neurophox.numpy.rst: -------------------------------------------------------------------------------- 1 | neurophox.numpy package 2 | ======================= 3 | 4 | neurophox.numpy.generic module 5 | ------------------------------ 6 | 7 | .. automodule:: neurophox.numpy.generic 8 | :members: 9 | :undoc-members: 10 | :show-inheritance: 11 | 12 | neurophox.numpy.layers module 13 | ----------------------------- 14 | 15 | .. automodule:: neurophox.numpy.layers 16 | :members: 17 | :undoc-members: 18 | :show-inheritance: 19 | 20 | -------------------------------------------------------------------------------- /doc/source/neurophox.rst: -------------------------------------------------------------------------------- 1 | neurophox package 2 | ================= 3 | 4 | Subpackages 5 | ----------- 6 | 7 | .. toctree:: 8 | 9 | neurophox.components 10 | neurophox.control 11 | neurophox.ml 12 | neurophox.numpy 13 | neurophox.tensorflow 14 | neurophox.torch 15 | 16 | 17 | neurophox.helpers module 18 | ------------------------ 19 | 20 | .. automodule:: neurophox.helpers 21 | :members: 22 | :undoc-members: 23 | :show-inheritance: 24 | 25 | neurophox.initializers module 26 | ----------------------------- 27 | 28 | .. automodule:: neurophox.initializers 29 | :members: 30 | :undoc-members: 31 | :show-inheritance: 32 | 33 | neurophox.meshmodel module 34 | -------------------------- 35 | 36 | .. automodule:: neurophox.meshmodel 37 | :members: 38 | :undoc-members: 39 | :show-inheritance: 40 | 41 | -------------------------------------------------------------------------------- /doc/source/neurophox.tensorflow.rst: -------------------------------------------------------------------------------- 1 | neurophox.tensorflow package 2 | ============================ 3 | 4 | neurophox.tensorflow.generic module 5 | ----------------------------------- 6 | 7 | .. automodule:: neurophox.tensorflow.generic 8 | :members: 9 | :undoc-members: 10 | :show-inheritance: 11 | 12 | neurophox.tensorflow.layers module 13 | ---------------------------------- 14 | 15 | .. automodule:: neurophox.tensorflow.layers 16 | :members: 17 | :undoc-members: 18 | :show-inheritance: 19 | 20 | -------------------------------------------------------------------------------- /doc/source/neurophox.torch.rst: -------------------------------------------------------------------------------- 1 | neurophox.torch package 2 | ============================ 3 | 4 | neurophox.torch.generic module 5 | ----------------------------------- 6 | 7 | .. automodule:: neurophox.torch.generic 8 | :members: 9 | :undoc-members: 10 | :show-inheritance: 11 | 12 | neurophox.torch.layers module 13 | ---------------------------------- 14 | 15 | .. automodule:: neurophox.torch.layers 16 | :members: 17 | :undoc-members: 18 | :show-inheritance: 19 | -------------------------------------------------------------------------------- /neurophox/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/solgaardlab/neurophox/29a9b03a4efc241a85ecb92463e32cfd69776d23/neurophox/__init__.py -------------------------------------------------------------------------------- /neurophox/components/__init__.py: -------------------------------------------------------------------------------- 1 | from .transfermatrix import PairwiseUnitary, Beamsplitter, PhaseShiftUpper, \ 2 | PhaseShiftLower, PhaseShiftCommonMode, PhaseShiftDifferentialMode 3 | from .mzi import MZI, SMMZI, BlochMZI 4 | -------------------------------------------------------------------------------- /neurophox/components/mzi.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from typing import Union, Tuple 3 | 4 | from ..config import NP_COMPLEX 5 | from .transfermatrix import PairwiseUnitary 6 | 7 | 8 | class MZI(PairwiseUnitary): 9 | """Mach-Zehnder Interferometer 10 | 11 | Class simulating the scattering matrix formulation of an ideal phase-shifting Mach-Zehnder interferometer. 12 | This can implement any :math:`2 \\times 2` unitary operator :math:`U_2 \in \mathrm{U}(2)`. 13 | 14 | The internal phase shifts, :math:`\\theta_1, \\theta_2`, and the external phase shifts :math:`\phi_1, \phi_2`, are 15 | used to define the final unitary operator as follows (where we define :math:`\\theta := \\theta_1- \\theta_2` 16 | for convenience): 17 | 18 | In Hadamard convention, the corresponding transfer matrix is: 19 | 20 | .. math:: 21 | U_2(\\theta, \phi) = H L(\\theta_1) R(\\theta_2) H L(\\phi_1) R(\\phi_2) = e^{i\\frac{\\theta_1 + \\theta_2}{2}} 22 | \\begin{bmatrix} e^{i \phi_1}\cos \\frac{\\theta}{2} & ie^{i \phi_2}\sin \\frac{\\theta}{2} \\\\ 23 | ie^{i \phi_1}\sin \\frac{\\theta}{2} & e^{i \phi_2}\cos \\frac{\\theta}{2} \\end{bmatrix} 24 | 25 | In beamsplitter convention, the corresponding transfer matrix is: 26 | 27 | .. math:: 28 | U_2(\\theta, \phi) = B L(\\theta_1) R(\\theta_2) B L(\\phi_1) R(\\phi_2) = i e^{i\\frac{\\theta_1 + \\theta_2}{2}} 29 | \\begin{bmatrix} e^{i \phi_1}\sin \\frac{\\theta}{2} & e^{i \phi_2}\cos \\frac{\\theta}{2} \\\\ 30 | e^{i \phi_1}\cos \\frac{\\theta}{2} & -e^{i \phi_2}\sin \\frac{\\theta}{2} \\end{bmatrix} 31 | 32 | Args: 33 | internal_upper: Upper internal phase shift 34 | internal_lower: Lower internal phase shift 35 | external_upper: Upper external phase shift 36 | external_lower: Lower external phase shift 37 | hadamard: Whether to use Hadamard convention 38 | epsilon: Beamsplitter error 39 | dtype: Type-casting to use for the matrix elements 40 | """ 41 | 42 | def __init__(self, internal_upper: float, internal_lower: float, external_upper: float, external_lower: float, 43 | hadamard: bool, epsilon: Union[float, Tuple[float, float]] = 0.0, dtype=NP_COMPLEX): 44 | super(MZI, self).__init__(dtype=dtype) 45 | self.internal_upper = internal_upper 46 | self.internal_lower = internal_lower 47 | self.external_upper = external_upper 48 | self.external_lower = external_lower 49 | self.hadamard = hadamard 50 | self.epsilon = (epsilon, epsilon) if isinstance(epsilon, float) or isinstance(epsilon, int) else epsilon 51 | 52 | @property 53 | def reflectivity(self): 54 | return np.abs(self.matrix[0][0]) ** 2 55 | 56 | @property 57 | def transmissivity(self): 58 | return np.abs(self.matrix[0][1]) ** 2 59 | 60 | @property 61 | def matrix(self): 62 | return get_mzi_transfer_matrix( 63 | internal_upper=self.internal_upper, 64 | internal_lower=self.internal_lower, 65 | external_upper=self.external_upper, 66 | external_lower=self.external_lower, 67 | epsilon=self.epsilon, 68 | hadamard=self.hadamard, 69 | dtype=self.dtype 70 | ) 71 | 72 | 73 | class SMMZI(MZI): 74 | """Mach-Zehnder Interferometer (single-mode basis) 75 | 76 | Class simulating an ideal phase-shifting Mach-Zehnder interferometer. 77 | As usual in our simulation environment, we have :math:`\\theta \in [0, \pi]` and :math:`\phi \in [0, 2\pi)`. 78 | 79 | In Hadamard convention, the corresponding transfer matrix is: 80 | 81 | .. math:: 82 | U_2(\\theta, \phi) = H L(\\theta) H L(\phi) = e^{-i \\theta / 2} 83 | \\begin{bmatrix} e^{i \phi}\cos \\frac{\\theta}{2} & i\sin \\frac{\\theta}{2} \\\\ 84 | ie^{i \phi}\sin \\frac{\\theta}{2} & \cos \\frac{\\theta}{2} \\end{bmatrix} 85 | 86 | In beamsplitter convention, the corresponding transfer matrix is: 87 | 88 | .. math:: 89 | U_2(\\theta, \phi) = B L(\\theta) B L(\phi) = ie^{-i \\theta / 2} 90 | \\begin{bmatrix} e^{i \phi}\sin \\frac{\\theta}{2} & \cos \\frac{\\theta}{2} \\\\ 91 | e^{i \phi}\cos \\frac{\\theta}{2} & -\sin \\frac{\\theta}{2} \\end{bmatrix} 92 | 93 | Args: 94 | theta: Amplitude-modulating phase shift 95 | phi: External phase shift, 96 | hadamard: Whether to use Hadamard convention 97 | epsilon: Beamsplitter error 98 | dtype: Type-casting to use for the matrix elements 99 | """ 100 | 101 | def __init__(self, theta: float, phi: float, hadamard: bool, lower_theta: bool = False, 102 | lower_phi: bool = False, epsilon: Union[float, Tuple[float, float]] = 0.0, dtype=NP_COMPLEX): 103 | self.theta = theta 104 | self.phi = phi 105 | super(SMMZI, self).__init__( 106 | internal_upper=theta if not lower_theta else 0, 107 | internal_lower=theta if lower_theta else 0, 108 | external_upper=phi if not lower_phi else 0, 109 | external_lower=phi if lower_phi else 0, 110 | epsilon=epsilon, hadamard=hadamard, dtype=dtype 111 | ) 112 | 113 | @classmethod 114 | def nullify(cls, vector: np.ndarray, idx: int, lower_theta: bool = False, lower_phi: bool = False): 115 | theta = -np.arctan2(np.abs(vector[idx]), np.abs(vector[idx + 1])) * 2 116 | theta = -theta if lower_theta else theta 117 | phi = np.angle(vector[idx + 1]) - np.angle(vector[idx]) + np.pi 118 | phi = -phi if lower_phi else phi 119 | mat = cls(theta, phi, hadamard=False, 120 | lower_theta=lower_theta, lower_phi=lower_phi).givens_rotation(vector.size, idx) 121 | nullified_vector = mat @ vector 122 | return nullified_vector, mat, np.mod(theta, 2 * np.pi), np.mod(phi, 2 * np.pi) 123 | 124 | 125 | class BlochMZI(MZI): 126 | """Mach-Zehnder Interferometer (Bloch basis, named after the Bloch sphere qubit formula) 127 | 128 | Class simulating an ideal phase-shifting Mach-Zehnder interferometer. 129 | As usual in our simulation environment, we have :math:`\\theta \in [0, \pi]` and :math:`\phi \in [0, 2\pi)`. 130 | 131 | In Hadamard convention, the corresponding transfer matrix is: 132 | 133 | .. math:: 134 | U_2(\\theta, \phi) = H D(\\theta) H L(\phi) = 135 | \\begin{bmatrix} e^{i \phi}\cos \\frac{\\theta}{2} & i\sin \\frac{\\theta}{2} \\\\ 136 | ie^{i \phi}\sin \\frac{\\theta}{2} & \cos \\frac{\\theta}{2} \\end{bmatrix} 137 | 138 | In beamsplitter convention, the corresponding transfer matrix is: 139 | 140 | .. math:: 141 | U_2(\\theta, \phi) = B D(\\theta) B L(\phi) = i 142 | \\begin{bmatrix} e^{i \phi}\sin \\frac{\\theta}{2} & \cos \\frac{\\theta}{2} \\\\ 143 | e^{i \phi}\cos \\frac{\\theta}{2} & -\sin \\frac{\\theta}{2} \\end{bmatrix} 144 | 145 | Args: 146 | theta: Amplitude-modulating phase shift 147 | phi: External phase shift, 148 | hadamard: Whether to use Hadamard convention 149 | epsilon: Beamsplitter error 150 | dtype: Type-casting to use for the matrix elements 151 | """ 152 | 153 | def __init__(self, theta: float, phi: float, hadamard: bool, 154 | epsilon: Union[float, Tuple[float, float]] = 0.0, dtype=NP_COMPLEX): 155 | self.theta = theta 156 | self.phi = phi 157 | super(BlochMZI, self).__init__( 158 | internal_upper=theta / 2, 159 | internal_lower=-theta / 2, 160 | external_upper=phi, 161 | external_lower=0, 162 | epsilon=epsilon, 163 | hadamard=hadamard, 164 | dtype=dtype 165 | ) 166 | 167 | 168 | def get_mzi_transfer_matrix(internal_upper: float, internal_lower: float, external_upper: float, external_lower: float, 169 | hadamard: float, epsilon: Tuple[float, float], dtype) -> np.ndarray: 170 | """Mach-Zehnder interferometer 171 | 172 | Args: 173 | internal_upper: Upper internal phase shift 174 | internal_lower: Lower internal phase shift 175 | external_upper: Upper external phase shift 176 | external_lower: Lower external phase shift 177 | hadamard: Whether to use Hadamard convention 178 | epsilon: Beamsplitter error 179 | dtype: Type-casting to use for the matrix elements 180 | 181 | Returns: 182 | MZI transfer matrix 183 | 184 | """ 185 | cc = np.cos(np.pi / 4 + epsilon[0]) * np.cos(np.pi / 4 + epsilon[1]) 186 | cs = np.cos(np.pi / 4 + epsilon[0]) * np.sin(np.pi / 4 + epsilon[1]) 187 | sc = np.sin(np.pi / 4 + epsilon[0]) * np.cos(np.pi / 4 + epsilon[1]) 188 | ss = np.sin(np.pi / 4 + epsilon[0]) * np.sin(np.pi / 4 + epsilon[1]) 189 | iu, il, eu, el = internal_upper, internal_lower, external_upper, external_lower 190 | if hadamard: 191 | return np.array([ 192 | [(cc * np.exp(1j * iu) + ss * np.exp(1j * il)) * np.exp(1j * eu), 193 | (cs * np.exp(1j * iu) - sc * np.exp(1j * il)) * np.exp(1j * el)], 194 | [(sc * np.exp(1j * iu) - cs * np.exp(1j * il)) * np.exp(1j * eu), 195 | (ss * np.exp(1j * iu) + cc * np.exp(1j * il)) * np.exp(1j * el)] 196 | ], dtype=dtype) 197 | else: 198 | return np.array([ 199 | [(cc * np.exp(1j * iu) - ss * np.exp(1j * il)) * np.exp(1j * eu), 200 | 1j * (cs * np.exp(1j * iu) + sc * np.exp(1j * il)) * np.exp(1j * el)], 201 | [1j * (sc * np.exp(1j * iu) + cs * np.exp(1j * il)) * np.exp(1j * eu), 202 | (cc * np.exp(1j * il) - ss * np.exp(1j * iu)) * np.exp(1j * el)]], dtype=dtype) 203 | 204 | 205 | def get_tdc_transfer_matrix(kappa: float, delta: float, external_upper: float, external_lower: float, dtype) -> np.ndarray: 206 | """Tunable directional coupler 207 | 208 | Args: 209 | kappa: Phase-matched phase shift (from coupled mode theory) 210 | delta: Phase-mismatched phase shift (from coupled-mode theory) 211 | external_upper: Upper external phase shift 212 | external_lower: Lower external phase shift 213 | dtype: Type-casting to use for the matrix elements 214 | 215 | Returns: 216 | MZI transfer matrix 217 | 218 | """ 219 | k, d, eu, el = kappa, delta, external_upper, external_lower 220 | q = np.sqrt(k ** 2 + d ** 2) 221 | s, c = np.sin(q), np.sin(q) 222 | return np.array([ 223 | [(c + 1j * d * s / q) * np.exp(1j * (eu + d)), -1j * k * s / q * np.exp(1j * (eu + d))], 224 | [-1j * k * s / q * np.exp(1j * (el - d)), (c - 1j * d * s / q) * np.exp(1j * (el - d))] 225 | ], dtype=dtype) 226 | -------------------------------------------------------------------------------- /neurophox/components/transfermatrix.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | 3 | import numpy as np 4 | 5 | from ..config import NP_COMPLEX 6 | 7 | 8 | class PairwiseUnitary: 9 | """Pairwise unitary 10 | 11 | This can be considered also a two-port unitary linear optical component 12 | (e.g. unitary scattering matrix for waveguide modal interaction). 13 | 14 | This unitary matrix :math:`U_2` has the following form: 15 | 16 | .. math:: 17 | U_2 = \\begin{bmatrix} U_{11} & U_{12} \\\ U_{21} & U_{22} \\end{bmatrix}, 18 | 19 | where we define the reflectivity :math:`r = |U_{11}|^2 = |U_{22}|^2` and transmissivity 20 | :math:`t = |U_{12}|^2 = |U_{21}|^2`. We also require that the determinant :math:`|\\det U| = 1`, 21 | from which we ultimately get the property :math:`U_2^\dagger U_2 = I_2`, where :math:`I_2` is the 22 | :math:`2 \\times 2` identity matrix. All children of this class have this property. 23 | 24 | Args: 25 | dtype: Cast values as :code:`dtype` for this pairwise unitary operator 26 | 27 | """ 28 | 29 | def __init__(self, dtype=NP_COMPLEX): 30 | self.dtype = dtype 31 | 32 | @property 33 | def reflectivity(self) -> float: 34 | """ 35 | 36 | Returns: 37 | Reflectivity, :math:`r`, for the matrix :math:`U` 38 | """ 39 | raise NotImplementedError("Need to override this method in child class.") 40 | 41 | @property 42 | def transmissivity(self) -> float: 43 | """ 44 | 45 | Returns: 46 | Transmissivity, :math:`t`, for the matrix :math:`U` 47 | """ 48 | raise NotImplementedError("Need to override this method in child class.") 49 | 50 | @property 51 | def matrix(self) -> np.ndarray: 52 | """ 53 | 54 | Returns: 55 | :math:`U_2`, a :math:`2 \\times 2` unitary matrix implemented by this component 56 | """ 57 | raise NotImplementedError("Need to override this method in child class.") 58 | 59 | @property 60 | def inverse_matrix(self) -> np.ndarray: 61 | return self.matrix.conj().T 62 | 63 | def givens_rotation(self, units: int, m: int, n: Optional[int] = None) -> np.ndarray: 64 | """Givens rotation matrix :math:`T_{m, n}` which rotates any vector about axes :math:`m, n`. 65 | 66 | Args: 67 | units: Input dimension to rotate 68 | m: Lower index to rotate 69 | n: Upper index to rotate, requires :math:`n > m` (defaults to `None`, which implies :math:`n = m + 1`) 70 | 71 | Returns: 72 | The Givens rotation matrix 73 | """ 74 | if n is None: 75 | n = m + 1 76 | if not m < n: 77 | raise ValueError('Require m < n') 78 | unitary = self.matrix 79 | givens_rotation = np.eye(units, dtype=self.dtype) 80 | givens_rotation[m][m] = unitary[0, 0] 81 | givens_rotation[m][n] = unitary[0, 1] 82 | givens_rotation[n][m] = unitary[1, 0] 83 | givens_rotation[n][n] = unitary[1, 1] 84 | return givens_rotation 85 | 86 | 87 | class Beamsplitter(PairwiseUnitary): 88 | """Ideal 50/50 beamsplitter 89 | 90 | Implements the Hadamard or beamsplitter operator. 91 | 92 | Hadamard transformation, :math:`H`: 93 | 94 | .. math:: 95 | H = \\frac{1}{\sqrt{2}}\\begin{bmatrix} 1 & 1\\\ 1 & -1 \\end{bmatrix} 96 | 97 | Beamsplitter transformation :math:`B`: 98 | 99 | .. math:: 100 | B = \\frac{1}{\sqrt{2}}\\begin{bmatrix} 1 & i\\\ i & 1 \\end{bmatrix} 101 | 102 | Hadamard transformation, :math:`H_\epsilon` (with error :math:`\epsilon`): 103 | 104 | .. math:: 105 | H_\epsilon = \\frac{1}{\sqrt{2}}\\begin{bmatrix} \sqrt{1 + \epsilon} & \sqrt{1 - \epsilon}\\\ \sqrt{1 - \epsilon} & -\sqrt{1 + \epsilon} \\end{bmatrix} 106 | 107 | Beamsplitter transformation :math:`B_\epsilon` (with error :math:`\epsilon`): 108 | 109 | .. math:: 110 | B_\epsilon = \\frac{1}{\sqrt{2}}\\begin{bmatrix} \sqrt{1 + \epsilon} & i\sqrt{1 - \epsilon}\\\ i\sqrt{1 - \epsilon} & \sqrt{1 + \epsilon} \\end{bmatrix} 111 | 112 | Args: 113 | hadamard: Amplitude-modulating phase shift 114 | epsilon: Errors for beamsplitter operator 115 | dtype: Cast values as :code:`dtype` for this pairwise unitary operator 116 | 117 | """ 118 | 119 | def __init__(self, hadamard: bool, epsilon: float = 0, dtype=NP_COMPLEX): 120 | super(Beamsplitter, self).__init__(dtype=dtype) 121 | self.hadamard = hadamard 122 | self.epsilon = epsilon 123 | 124 | @property 125 | def reflectivity(self) -> float: 126 | return np.cos(np.pi / 4 + self.epsilon) ** 2 127 | 128 | @property 129 | def transmissivity(self) -> float: 130 | return np.sin(np.pi / 4 + self.epsilon) ** 2 131 | 132 | @property 133 | def matrix(self) -> np.ndarray: 134 | r = np.sqrt(self.reflectivity) 135 | t = np.sqrt(self.transmissivity) 136 | if self.hadamard: 137 | return np.array([ 138 | [r, t], 139 | [t, -r] 140 | ], dtype=self.dtype) 141 | else: 142 | return np.array([ 143 | [r, 1j * t], 144 | [1j * t, r] 145 | ], dtype=self.dtype) 146 | 147 | 148 | class PhaseShiftUpper(PairwiseUnitary): 149 | """Upper phase shift operator 150 | 151 | Implements the upper phase shift operator :math:`L(\\theta)` given phase :math:`\\theta`: 152 | 153 | .. math:: 154 | L(\\theta) = \\begin{bmatrix} e^{i\\theta} & 0\\\ 0 & 1 \\end{bmatrix} 155 | 156 | Args: 157 | phase_shift: Phase shift :math:`\\theta` 158 | dtype: Cast values as :code:`dtype` for this pairwise unitary operator 159 | 160 | """ 161 | 162 | def __init__(self, phase_shift: float, dtype=NP_COMPLEX): 163 | super(PhaseShiftUpper, self).__init__(dtype=dtype) 164 | self.phase_shift = phase_shift 165 | 166 | @property 167 | def matrix(self) -> np.ndarray: 168 | return np.array([ 169 | [np.exp(1j * self.phase_shift), 0], 170 | [0, 1] 171 | ], dtype=self.dtype) 172 | 173 | 174 | class PhaseShiftLower(PairwiseUnitary): 175 | """Lower phase shift operator 176 | 177 | Implements the lower phase shift operator :math:`R(\\theta)` given phase :math:`\\theta`: 178 | 179 | .. math:: 180 | R(\\theta) = \\begin{bmatrix} 1 & 0\\\ 0 & e^{i\\theta} \\end{bmatrix} 181 | 182 | Args: 183 | phase_shift: Phase shift :math:`\\theta` 184 | dtype: Cast values as :code:`dtype` for this pairwise unitary operator 185 | 186 | """ 187 | 188 | def __init__(self, phase_shift: float, dtype=NP_COMPLEX): 189 | super(PhaseShiftLower, self).__init__(dtype=dtype) 190 | self.phase_shift = phase_shift 191 | 192 | @property 193 | def matrix(self) -> np.ndarray: 194 | return np.array([ 195 | [1, 0], 196 | [0, np.exp(1j * self.phase_shift)] 197 | ], dtype=self.dtype) 198 | 199 | 200 | class PhaseShiftDifferentialMode(PairwiseUnitary): 201 | """Differential phase shift operator 202 | 203 | Implements the differential-mode phase shift operator :math:`D(\\theta)` given phase :math:`\\theta`: 204 | 205 | .. math:: 206 | D(\\theta) = \\begin{bmatrix} e^{i\\theta / 2} & 0\\\ 0 & e^{-i\\theta/2} \\end{bmatrix} 207 | 208 | Args: 209 | phase_shift: Phase shift :math:`\\theta` 210 | dtype: Cast values as :code:`dtype` for this pairwise unitary operator 211 | 212 | """ 213 | 214 | def __init__(self, phase_shift: float, dtype=NP_COMPLEX): 215 | super(PhaseShiftDifferentialMode, self).__init__(dtype=dtype) 216 | self.phase_shift = phase_shift 217 | 218 | @property 219 | def matrix(self) -> np.ndarray: 220 | return np.array([ 221 | [np.exp(1j * self.phase_shift / 2), 0], 222 | [0, np.exp(-1j * self.phase_shift / 2)] 223 | ], dtype=self.dtype) 224 | 225 | 226 | class PhaseShiftCommonMode(PairwiseUnitary): 227 | """Common mode phase shift operator 228 | 229 | Implements the common-mode phase shift operator :math:`C(\\theta)` given phase :math:`\\theta`: 230 | 231 | .. math:: 232 | C(\\theta) = \\begin{bmatrix} e^{i\\theta} & 0\\\ 0 & e^{i\\theta} \\end{bmatrix} 233 | 234 | Args: 235 | phase_shift: Phase shift :math:`\\theta` 236 | dtype: Cast values as :code:`dtype` for this pairwise unitary operator 237 | 238 | """ 239 | 240 | def __init__(self, phase_shift: float, dtype=NP_COMPLEX): 241 | super(PhaseShiftCommonMode, self).__init__(dtype=dtype) 242 | self.phase_shift = phase_shift 243 | 244 | @property 245 | def matrix(self) -> np.ndarray: 246 | return np.array([ 247 | [np.exp(1j * self.phase_shift), 0], 248 | [0, np.exp(1j * self.phase_shift)] 249 | ], dtype=self.dtype) 250 | -------------------------------------------------------------------------------- /neurophox/config.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import torch 4 | 5 | # Backends 6 | 7 | PYTORCH = 'torch' 8 | TFKERAS = 'tf' 9 | NUMPY = 'numpy' 10 | 11 | # Types (for memory) 12 | 13 | NP_COMPLEX = np.complex128 14 | NP_FLOAT = np.float64 15 | 16 | TF_COMPLEX = tf.complex64 17 | TF_FLOAT = tf.float32 18 | 19 | T_FLOAT = torch.double 20 | T_COMPLEX = torch.cdouble 21 | 22 | # Test seed 23 | 24 | TEST_SEED = 31415 25 | 26 | # Phase basis 27 | 28 | BLOCH = "bloch" 29 | SINGLEMODE = "sm" 30 | DEFAULT_BASIS = BLOCH 31 | -------------------------------------------------------------------------------- /neurophox/decompositions.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from .components import BlochMZI, SMMZI 3 | from .config import NP_FLOAT 4 | from typing import Callable, Tuple 5 | from .numpy import MeshNumpyLayer, RMNumpy 6 | from .meshmodel import MeshModel 7 | from .helpers import inverse_permutation 8 | 9 | 10 | def clements_decomposition(u: np.ndarray, pbar_handle: Callable = None, smmzi: bool = False) -> RMNumpy: 11 | """Clements decomposition of unitary matrix :math:`U` to output a NumPy rectangular mesh layer 12 | 13 | Args: 14 | u: unitary matrix :math:`U` to be decomposed into pairwise operators. 15 | pbar_handle: Useful for larger matrices 16 | 17 | Returns: 18 | The :code:`RMNumpy` layer that outputs the unitary :math:`U` 19 | 20 | """ 21 | u_hat = u.T.copy() 22 | n = u.shape[0] 23 | # odd and even layer dimensions 24 | theta_checkerboard = np.zeros_like(u, dtype=NP_FLOAT) 25 | phi_checkerboard = np.zeros_like(u, dtype=NP_FLOAT) 26 | phi_checkerboard = np.hstack((np.zeros((n, 1)), phi_checkerboard)) 27 | iterator = pbar_handle(range(n - 1)) if pbar_handle else range(n - 1) 28 | MZI = SMMZI if smmzi else BlochMZI 29 | for i in iterator: 30 | if i % 2: 31 | for j in range(i + 1): 32 | pairwise_index = n + j - i - 2 33 | target_row, target_col = n + j - i - 1, j 34 | theta = np.arctan(np.abs(u_hat[target_row - 1, target_col] / u_hat[target_row, target_col])) * 2 35 | phi = np.angle(u_hat[target_row, target_col] / u_hat[target_row - 1, target_col]) 36 | mzi = MZI(theta, phi, hadamard=False, dtype=np.complex128) 37 | left_multiplier = mzi.givens_rotation(units=n, m=pairwise_index) 38 | u_hat = left_multiplier @ u_hat 39 | theta_checkerboard[pairwise_index, j] = theta 40 | phi_checkerboard[pairwise_index, j] = -phi + np.pi 41 | phi_checkerboard[pairwise_index + 1, j] = np.pi 42 | else: 43 | for j in range(i + 1): 44 | pairwise_index = i - j 45 | target_row, target_col = n - j - 1, i - j 46 | theta = np.arctan(np.abs(u_hat[target_row, target_col + 1] / u_hat[target_row, target_col])) * 2 47 | phi = np.angle(-u_hat[target_row, target_col] / u_hat[target_row, target_col + 1]) 48 | mzi = BlochMZI(theta, phi, hadamard=False, dtype=np.complex128) 49 | right_multiplier = mzi.givens_rotation(units=n, m=pairwise_index) 50 | u_hat = u_hat @ right_multiplier.conj().T 51 | theta_checkerboard[pairwise_index, -j - 1] = theta 52 | phi_checkerboard[pairwise_index, -j - 1] = phi + np.pi 53 | 54 | diag_phases = np.angle(np.diag(u_hat)) 55 | theta = checkerboard_to_param(np.fliplr(theta_checkerboard), n) 56 | phi_checkerboard = np.fliplr(phi_checkerboard) 57 | if n % 2: 58 | phi_checkerboard[:, :-1] += np.fliplr(np.diag(diag_phases)) 59 | else: 60 | phi_checkerboard[:, 1:] += np.fliplr(np.diag(diag_phases)) 61 | phi_checkerboard[-1, 2::2] += np.pi / 2 # neurophox layers assume pi / 2 phase shift in even layer "bounces" 62 | phi_checkerboard[0, 2::2] += np.pi / 2 63 | 64 | gamma = phi_checkerboard[:, 0] 65 | external_phases = phi_checkerboard[:, 1:] 66 | phi, gamma = grid_common_mode_flow(external_phases, gamma=gamma) 67 | phi = checkerboard_to_param(phi, n) 68 | 69 | # for some reason, we need to adjust gamma at the end in this strange way (found via trial and error...): 70 | gamma_adj = np.zeros_like(gamma) 71 | gamma_adj[1::4] = 1 72 | gamma_adj[2::4] = 1 73 | gamma += np.pi * (1 - gamma_adj) if (n // 2) % 2 else np.pi * gamma_adj 74 | gamma = np.mod(gamma, 2 * np.pi) 75 | 76 | return RMNumpy(units=n, theta_init=theta, phi_init=phi, gamma_init=gamma) 77 | 78 | 79 | def reck_decomposition(u: np.ndarray, pbar_handle: Callable = None, 80 | lower_theta: bool = True, lower_phi: bool = True) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: 81 | """Reck decomposition of unitary matrix :math:`U` to output a NumPy triangular mesh layer (SMMZI convention only) 82 | 83 | Args: 84 | u: unitary matrix :math:`U` to be decomposed into pairwise operators. 85 | pbar_handle: Useful for larger matrices 86 | 87 | Returns: 88 | The thetas and phis for a single-mode convention MZI (each row is a diagonal network of MZIs) 89 | 90 | """ 91 | 92 | u_hat = u.T.copy() 93 | n = u.shape[0] 94 | thetas = np.zeros((n - 1, n - 1)) 95 | phis = np.zeros((n - 1, n - 1)) 96 | iterator = pbar_handle(range(n - 1)) if pbar_handle else range(n - 1) 97 | for i in iterator: 98 | for j in range(n - i - 1): 99 | _, mat, theta, phi = SMMZI.nullify(u_hat[i], n - 2 - j, lower_theta=lower_theta, lower_phi=True) 100 | u_hat = u_hat @ mat.T 101 | thetas[i, j] = theta 102 | phis[i, j] = phi 103 | return u_hat, thetas, phis 104 | 105 | 106 | def checkerboard_to_param(checkerboard: np.ndarray, units: int): 107 | param = np.zeros((units, units // 2)) 108 | if units % 2: 109 | param[::2, :] = checkerboard.T[::2, :-1:2] 110 | else: 111 | param[::2, :] = checkerboard.T[::2, ::2] 112 | param[1::2, :] = checkerboard.T[1::2, 1::2] 113 | return param 114 | 115 | 116 | def grid_common_mode_flow(external_phases: np.ndarray, gamma: np.ndarray, basis: str = "sm"): 117 | """In a grid mesh (e.g., triangular, rectangular meshes), arrange phases according to single-mode (:code:`sm`), 118 | differential mode (:code:`diff`), or max-:math:`\\pi` (:code:`maxpi`, all external phase shifts are at most 119 | :math:`\\pi`). This is done using a procedure called "common mode flow" where common modes are shifted 120 | throughout the mesh until phases are correctly set. 121 | 122 | Args: 123 | external_phases: external phases in the grid mesh 124 | gamma: input phase shifts 125 | basis: single-mode (:code:`sm`), differential mode (:code:`diff`), or max-:math:`\\pi` (:code:`maxpi`) 126 | 127 | Returns: 128 | new external phases shifts and new gamma resulting 129 | 130 | """ 131 | units, num_layers = external_phases.shape 132 | phase_shifts = np.hstack((gamma[:, np.newaxis], external_phases)).T 133 | new_phase_shifts = np.zeros_like(external_phases.T) 134 | 135 | for i in range(num_layers): 136 | current_layer = num_layers - i 137 | start_idx = (current_layer - 1) % 2 138 | end_idx = units - (current_layer + units - 1) % 2 139 | # calculate phase information 140 | upper_phase = phase_shifts[current_layer][start_idx:end_idx][::2] 141 | lower_phase = phase_shifts[current_layer][start_idx:end_idx][1::2] 142 | upper_phase = np.mod(upper_phase, 2 * np.pi) 143 | lower_phase = np.mod(lower_phase, 2 * np.pi) 144 | if basis == "sm": 145 | new_phase_shifts[-i - 1][start_idx:end_idx][::2] = upper_phase - lower_phase 146 | # assign differential phase to the single mode layer and keep common mode layer 147 | else: 148 | phase_diff = upper_phase - lower_phase 149 | phase_diff[phase_diff > np.pi] -= 2 * np.pi 150 | phase_diff[phase_diff < -np.pi] += 2 * np.pi 151 | if basis == "diff": 152 | new_phase_shifts[-i - 1][start_idx:end_idx][::2] = phase_diff / 2 153 | new_phase_shifts[-i - 1][start_idx:end_idx][1::2] = -phase_diff / 2 154 | elif basis == "pimax": 155 | new_phase_shifts[-i - 1][start_idx:end_idx][::2] = phase_diff * (phase_diff >= 0) 156 | new_phase_shifts[-i - 1][start_idx:end_idx][1::2] = -phase_diff * (phase_diff < 0) 157 | # update the previous layer with the common mode calculated for the current layer\ 158 | phase_shifts[current_layer] -= new_phase_shifts[-i - 1] 159 | phase_shifts[current_layer - 1] += np.mod(phase_shifts[current_layer], 2 * np.pi) 160 | phase_shifts[current_layer] = 0 161 | new_gamma = np.mod(phase_shifts[0], 2 * np.pi) 162 | return np.mod(new_phase_shifts.T, 2 * np.pi), new_gamma 163 | 164 | 165 | def parallel_nullification(np_layer): 166 | """Perform parallel nullification 167 | 168 | Args: 169 | np_layer: 170 | 171 | Returns: 172 | 173 | """ 174 | units, num_layers = np_layer.units, np_layer.num_layers 175 | nullification_set = np_layer.nullification_set 176 | 177 | # set the mesh to bar state 178 | theta = [] 179 | phi = [] 180 | 181 | perm_idx = np_layer.mesh.model.perm_idx 182 | num_tunable = np_layer.mesh.model.num_tunable 183 | 184 | # run the real-time O(L) algorithm 185 | for idx in range(num_layers): 186 | layer = num_layers - idx - 1 187 | if idx > 0: 188 | current_mesh = MeshNumpyLayer( 189 | MeshModel(perm_idx=perm_idx[layer + 1:], 190 | num_tunable=num_tunable[layer + 1:], 191 | basis='sm', 192 | theta_init=np.asarray(theta), 193 | phi_init=np.asarray(phi), 194 | gamma_init=np.zeros_like(np_layer.phases.gamma)) 195 | ) 196 | layer_trm = current_mesh.inverse_transform(nullification_set[layer]).squeeze() 197 | else: 198 | layer_trm = nullification_set[layer].take(inverse_permutation(perm_idx[-1])) 199 | upper_inputs = layer_trm[:-1][::2] 200 | lower_inputs = layer_trm[1:][::2] 201 | theta.insert(0, np.arctan(np.abs(upper_inputs / lower_inputs)) * 2) 202 | phi.insert(0, np.angle(upper_inputs / lower_inputs)) 203 | return MeshNumpyLayer( 204 | MeshModel(perm_idx=perm_idx, 205 | num_tunable=num_tunable, 206 | basis='sm', 207 | theta_init=np.asarray(theta), 208 | phi_init=np.asarray(phi), 209 | gamma_init=np_layer.phases.gamma.copy()) 210 | ) 211 | -------------------------------------------------------------------------------- /neurophox/helpers.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, Callable, Tuple 2 | 3 | import numpy as np 4 | import tensorflow as tf 5 | try: 6 | import torch 7 | except ImportError: 8 | # if the user did not install pytorch, just do tensorflow stuff 9 | pass 10 | 11 | from scipy.stats import multivariate_normal 12 | 13 | from .config import NP_FLOAT 14 | 15 | 16 | def to_stripe_array(nparray: np.ndarray, units: int): 17 | """ 18 | Convert a numpy array of phase shifts of size (`num_layers`, `units`) 19 | or (`batch_size`, `num_layers`, `units`) into striped array for use in 20 | general feedforward mesh architectures. 21 | 22 | Args: 23 | nparray: phase shift values for all columns 24 | units: dimension the stripe array acts on (depends on parity) 25 | 26 | Returns: 27 | A general mesh stripe array arrangement that is of size (`units`, `num_layers`) 28 | or (`batch_size`, `units`, `num_layers`) 29 | """ 30 | if len(nparray.shape) == 2: 31 | num_layers, _ = nparray.shape 32 | stripe_array = np.zeros((units - 1, num_layers), dtype=nparray.dtype) 33 | stripe_array[::2] = nparray.T 34 | stripe_array = np.vstack([stripe_array, np.zeros(shape=(1, num_layers))]) 35 | else: 36 | num_layers, _, batch_size = nparray.shape 37 | stripe_array = np.zeros((units - 1, num_layers, batch_size), dtype=nparray.dtype) 38 | stripe_array[::2] = nparray.transpose((1, 0, 2)) 39 | stripe_array = np.vstack([stripe_array, np.zeros(shape=(1, num_layers, batch_size))]) 40 | return stripe_array 41 | 42 | 43 | def to_absolute_theta(theta: np.ndarray) -> np.ndarray: 44 | theta = np.mod(theta, 2 * np.pi) 45 | theta[theta > np.pi] = 2 * np.pi - theta[theta > np.pi] 46 | return theta 47 | 48 | 49 | def get_haar_diagonal_sequence(diagonal_length, parity_odd: bool = False): 50 | odd_nums = list(diagonal_length + 1 - np.flip(np.arange(1, diagonal_length + 1, 2), axis=0)) 51 | even_nums = list(diagonal_length + 1 - np.arange(2, 2 * (diagonal_length - len(odd_nums)) + 1, 2)) 52 | nums = np.asarray(odd_nums + even_nums) 53 | if parity_odd: 54 | nums = nums[::-1] 55 | return nums 56 | 57 | 58 | def get_alpha_checkerboard(units: int, num_layers: int, include_off_mesh: bool = False, flipud=False): 59 | if units < num_layers: 60 | raise ValueError("Require units >= num_layers!") 61 | alpha_checkerboard = np.zeros((units - 1, num_layers)) 62 | diagonal_length_to_sequence = [get_haar_diagonal_sequence(i, bool(num_layers % 2)) for i in 63 | range(1, num_layers + 1)] 64 | for i in range(units - 1): 65 | for j in range(num_layers): 66 | if (i + j) % 2 == 0: 67 | if i < num_layers and j > i: 68 | diagonal_length = num_layers - np.abs(i - j) 69 | elif i > units - num_layers and j < i - units + num_layers: 70 | diagonal_length = num_layers - np.abs(i - j - units + num_layers) - 1 * (units == num_layers) 71 | else: 72 | diagonal_length = num_layers - 1 * (units == num_layers) 73 | alpha_checkerboard[i, j] = 1 if diagonal_length == 1 else \ 74 | diagonal_length_to_sequence[int(diagonal_length) - 1][min(i, j)] 75 | else: 76 | if include_off_mesh: 77 | alpha_checkerboard[i, j] = 1 78 | # symmetrize the checkerboard 79 | if units != num_layers: 80 | alpha_checkerboard = (alpha_checkerboard + np.flipud(alpha_checkerboard)) / 2 81 | return alpha_checkerboard if not flipud else np.flipud(alpha_checkerboard) 82 | 83 | 84 | def get_alpha_checkerboard_general(units: int, num_layers: int): 85 | alpha_checkerboards = [get_alpha_checkerboard(units, units, flipud=bool(n % 2 and units % 2)) 86 | for n in range(num_layers // units)] 87 | extra_layers = num_layers - num_layers // units * units 88 | if extra_layers < units: 89 | # partial checkerboard 90 | alpha_checkerboards.append( 91 | get_alpha_checkerboard(units, extra_layers, flipud=not num_layers // units % 2 and units % 2)) 92 | return np.hstack(alpha_checkerboards) 93 | 94 | 95 | def get_efficient_coarse_grain_block_sizes(units: int, tunable_layers_per_block: int = 2, use_cg_sequence: bool = True): 96 | num_blocks = int(np.rint(np.log2(units))) 97 | sampling_frequencies = [2 ** (block_num + 1) for block_num in range(num_blocks - 1)] 98 | if use_cg_sequence: 99 | sampling_frequencies = 2 ** get_haar_diagonal_sequence(num_blocks - 1) 100 | tunable_block_sizes = [tunable_layers_per_block for _ in range(num_blocks - 1)] 101 | return np.asarray(tunable_block_sizes, dtype=np.int32), np.asarray(sampling_frequencies, dtype=np.int32) 102 | 103 | 104 | def get_default_coarse_grain_block_sizes(units: int, use_cg_sequence: bool = True): 105 | num_blocks = int(np.rint(np.log2(units))) 106 | sampling_frequencies = [2 ** (block_num + 1) for block_num in range(num_blocks - 1)] 107 | if use_cg_sequence: 108 | sampling_frequencies = 2 ** get_haar_diagonal_sequence(num_blocks - 1) 109 | tunable_layer_rank = int(np.floor(units / num_blocks)) 110 | tunable_layer_rank = tunable_layer_rank + 1 if tunable_layer_rank % 2 else tunable_layer_rank 111 | tunable_block_sizes = [tunable_layer_rank for _ in range(num_blocks - 1)] 112 | tunable_block_sizes.append(units - tunable_layer_rank * (num_blocks - 1)) 113 | return np.asarray(tunable_block_sizes, dtype=np.int32), np.asarray(sampling_frequencies, dtype=np.int32) 114 | 115 | 116 | def prm_permutation(units: int, tunable_block_sizes: np.ndarray, 117 | sampling_frequencies: np.ndarray, butterfly: bool = False): 118 | grid_perms = [grid_permutation(units, tunable_block_size) for tunable_block_size in tunable_block_sizes] 119 | perms_to_concatenate = [grid_perms[0][0]] 120 | for idx, frequency in enumerate(sampling_frequencies): 121 | perm_prev = grid_perms[idx][-1] 122 | perm_next = grid_perms[idx + 1][0] 123 | perm = butterfly_layer_permutation(units, frequency) if butterfly else rectangular_permutation(units, frequency) 124 | glued_perm = glue_permutations(perm_prev, perm) 125 | glued_perm = glue_permutations(glued_perm, perm_next) 126 | perms_to_concatenate += [grid_perms[idx][1:-1], glued_perm] 127 | perms_to_concatenate.append(grid_perms[-1][1:]) 128 | return np.vstack(perms_to_concatenate) 129 | 130 | 131 | def butterfly_layer_permutation(units: int, frequency: int): 132 | if units % 2: 133 | raise NotImplementedError('Odd input dimension case not yet implemented.') 134 | frequency = frequency 135 | unpermuted_indices = np.arange(units) 136 | num_splits = units // frequency 137 | total_num_indices = num_splits * frequency 138 | unpermuted_indices_remainder = unpermuted_indices[total_num_indices:] 139 | permuted_indices = np.hstack( 140 | [np.hstack([i, i + frequency] for i in range(frequency)) + 2 * frequency * split_num 141 | for split_num in range(num_splits // 2)] + [unpermuted_indices_remainder] 142 | ) 143 | return permuted_indices.astype(np.int32) 144 | 145 | 146 | def rectangular_permutation(units: int, frequency: int): 147 | unpermuted_indices = np.arange(units) 148 | frequency_offset = np.empty((units,)) 149 | frequency_offset[::2] = -frequency 150 | frequency_offset[1::2] = frequency 151 | permuted_indices = unpermuted_indices + frequency_offset 152 | for idx in range(units): 153 | if permuted_indices[idx] < 0: 154 | permuted_indices[idx] = -1 - permuted_indices[idx] 155 | if permuted_indices[idx] > units - 1: 156 | permuted_indices[idx] = 2 * units - 1 - permuted_indices[idx] 157 | return permuted_indices.astype(np.int32) 158 | 159 | 160 | def grid_permutation(units: int, num_layers: int): 161 | ordered_idx = np.arange(units) 162 | split_num_layers = (num_layers - num_layers // 2, num_layers // 2) 163 | left_shift = np.roll(ordered_idx, -1, axis=0) 164 | right_shift = np.roll(ordered_idx, 1, axis=0) 165 | permuted_indices = np.zeros((num_layers, units)) 166 | permuted_indices[::2] = np.ones((split_num_layers[0], 1)) @ left_shift[np.newaxis, :] 167 | permuted_indices[1::2] = np.ones((split_num_layers[1], 1)) @ right_shift[np.newaxis, :] 168 | if num_layers % 2: 169 | return np.vstack((ordered_idx.astype(np.int32), 170 | permuted_indices[:-1].astype(np.int32), 171 | ordered_idx.astype(np.int32))) 172 | return np.vstack((ordered_idx.astype(np.int32), 173 | permuted_indices.astype(np.int32))) 174 | 175 | 176 | def grid_viz_permutation(units: int, num_layers: int, flip: bool = False): 177 | ordered_idx = np.arange(units) 178 | split_num_layers = (num_layers - num_layers // 2, num_layers // 2) 179 | right_shift = np.roll(ordered_idx, 1, axis=0) 180 | permuted_indices = np.zeros((num_layers, units)) 181 | if flip: 182 | permuted_indices[::2] = np.ones((split_num_layers[0], 1)) @ ordered_idx[np.newaxis, :] 183 | permuted_indices[1::2] = np.ones((split_num_layers[1], 1)) @ right_shift[np.newaxis, :] 184 | else: 185 | permuted_indices[::2] = np.ones((split_num_layers[0], 1)) @ right_shift[np.newaxis, :] 186 | permuted_indices[1::2] = np.ones((split_num_layers[1], 1)) @ ordered_idx[np.newaxis, :] 187 | return np.vstack((ordered_idx.astype(np.int32), 188 | permuted_indices[:-1].astype(np.int32), 189 | ordered_idx.astype(np.int32))) 190 | 191 | 192 | def ordered_viz_permutation(units: int, num_layers: int): 193 | ordered_idx = np.arange(units) 194 | permuted_indices = np.ones((num_layers + 1, 1)) @ ordered_idx[np.newaxis, :] 195 | return permuted_indices.astype(np.int32) 196 | 197 | 198 | def plot_complex_matrix(plt, matrix: np.ndarray): 199 | plt.figure(figsize=(15, 5), dpi=200) 200 | plt.subplot(131) 201 | plt.title('Absolute') 202 | plt.imshow(np.abs(matrix), cmap='hot') 203 | plt.colorbar(shrink=0.7) 204 | plt.subplot(132) 205 | plt.title('Real') 206 | plt.imshow(np.real(matrix), cmap='hot') 207 | plt.colorbar(shrink=0.7) 208 | plt.subplot(133) 209 | plt.title('Imag') 210 | plt.imshow(np.imag(matrix), cmap='hot') 211 | plt.colorbar(shrink=0.7) 212 | 213 | 214 | def random_gaussian_batch(batch_size: int, units: int, covariance_matrix: Optional[np.ndarray] = None, 215 | seed: Optional[int] = None) -> np.ndarray: 216 | if seed is not None: 217 | np.random.seed(seed) 218 | input_matrix = multivariate_normal.rvs( 219 | mean=np.zeros(units), 220 | cov=1 if not covariance_matrix else covariance_matrix, 221 | size=batch_size 222 | ) 223 | random_phase = np.random.rand(batch_size, units).astype(dtype=NP_FLOAT) * 2 * np.pi 224 | return input_matrix * np.exp(1j * random_phase) 225 | 226 | 227 | def glue_permutations(perm_idx_1: np.ndarray, perm_idx_2: np.ndarray): 228 | perm_idx = np.zeros_like(perm_idx_1) 229 | perm_idx[perm_idx_2] = perm_idx_1 230 | return perm_idx.astype(np.int32) 231 | 232 | 233 | def inverse_permutation(permuted_indices: np.ndarray): 234 | inv_permuted_indices = np.zeros_like(permuted_indices) 235 | for idx, perm_idx in enumerate(permuted_indices): 236 | inv_permuted_indices[perm_idx] = idx 237 | return inv_permuted_indices 238 | 239 | 240 | def pairwise_off_diag_permutation(units: int): 241 | ordered_idx = np.arange(units) 242 | perm_idx = np.zeros_like(ordered_idx) 243 | if units % 2: 244 | perm_idx[:-1][::2] = ordered_idx[1::2] 245 | perm_idx[1::2] = ordered_idx[:-1][::2] 246 | perm_idx[-1] = ordered_idx[-1] 247 | else: 248 | perm_idx[::2] = ordered_idx[1::2] 249 | perm_idx[1::2] = ordered_idx[::2] 250 | return perm_idx.astype(np.int32) 251 | 252 | 253 | def butterfly_permutation(num_layers: int): 254 | ordered_idx = np.arange(2 ** num_layers) 255 | permuted_idx = np.vstack( 256 | [butterfly_layer_permutation(2 ** num_layers, 2 ** layer) for layer in range(num_layers)] 257 | ).astype(np.int32) 258 | return np.vstack((ordered_idx.astype(np.int32), 259 | permuted_idx[1:].astype(np.int32), 260 | ordered_idx.astype(np.int32))) 261 | 262 | 263 | def neurophox_matplotlib_setup(plt): 264 | plt.rc('text', usetex=True) 265 | plt.rc('font', family='serif') 266 | # plt.rc('text', usetex=True) 267 | # plt.rc('font', **{'family': 'serif', 'serif': ['Charter']}) 268 | # plt.rcParams['mathtext.fontset'] = 'dejavuserif' 269 | plt.rcParams.update({'text.latex.preamble': [r'\usepackage{siunitx}', r'\usepackage{amsmath}']}) 270 | 271 | 272 | # Phase functions 273 | 274 | 275 | def fix_phase_tf(fixed, mask): 276 | return lambda tensor: mask * tensor + (1 - mask) * fixed 277 | 278 | 279 | def fix_phase_torch(fixed: np.ndarray, mask: np.ndarray, 280 | device: torch.device = torch.device('cpu'), dtype: torch.dtype = torch.cfloat): 281 | mask = torch.as_tensor(mask, dtype=dtype, device=device) 282 | fixed = torch.as_tensor(fixed, dtype=dtype, device=device) 283 | return lambda tensor: tensor * mask + (1 - mask) * fixed 284 | 285 | 286 | def tri_phase_tf(phase_range: float): 287 | def pcf(phase): 288 | phase = tf.math.mod(phase, 2 * phase_range) 289 | phase = tf.where(tf.greater(phase, phase_range), 290 | 2 * phase_range * tf.ones_like(phase) - phase, phase) 291 | return phase 292 | 293 | return pcf 294 | 295 | 296 | def tri_phase_torch(phase_range: float): 297 | def pcf(phase): 298 | phase = torch.fmod(phase, 2 * phase_range) 299 | phase[phase > phase_range] = 2 * phase_range - phase[phase > phase_range] 300 | return phase 301 | 302 | return pcf 303 | -------------------------------------------------------------------------------- /neurophox/initializers.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple, Union, Optional 2 | 3 | import tensorflow as tf 4 | 5 | try: 6 | import torch 7 | from torch.nn import Parameter 8 | except ImportError: 9 | # if the user did not install pytorch, just do tensorflow stuff 10 | pass 11 | 12 | import numpy as np 13 | 14 | from .config import TF_FLOAT, NP_FLOAT, TEST_SEED, T_FLOAT 15 | from .helpers import get_alpha_checkerboard_general, get_default_coarse_grain_block_sizes, \ 16 | get_efficient_coarse_grain_block_sizes 17 | from scipy.special import betaincinv 18 | 19 | 20 | class MeshPhaseInitializer: 21 | def __init__(self, units: int, num_layers: int): 22 | """ 23 | 24 | Args: 25 | units: Input dimension, :math:`N` 26 | num_layers: Number of layers :math:`L` 27 | """ 28 | self.units, self.num_layers = units, num_layers 29 | 30 | def to_np(self) -> np.ndarray: 31 | """ 32 | 33 | Returns: 34 | Initialized Numpy array 35 | """ 36 | raise NotImplementedError('Need to implement numpy initialization') 37 | 38 | def to_tf(self, phase_varname: str) -> tf.Variable: 39 | """ 40 | 41 | Returns: 42 | Initialized Tensorflow Variable 43 | """ 44 | phase_np = self.to_np() 45 | return tf.Variable( 46 | name=phase_varname, 47 | initial_value=phase_np, 48 | dtype=TF_FLOAT 49 | ) 50 | 51 | def to_torch(self, is_trainable: bool = True): 52 | """ 53 | 54 | Returns: 55 | Initialized torch Parameter 56 | """ 57 | phase_initializer = self.to_np() 58 | phase = Parameter(torch.tensor(phase_initializer, dtype=T_FLOAT), requires_grad=is_trainable) 59 | return phase 60 | 61 | 62 | class PhaseInitializer(MeshPhaseInitializer): 63 | """ 64 | User-specified initialization of rectangular and triangular mesh architectures. 65 | 66 | Args: 67 | phase: Phase to initialize 68 | units: Input dimension, :math:`N` 69 | """ 70 | 71 | def __init__(self, phase: np.ndarray, units: int): 72 | self.phase, self.units = phase, units 73 | super(PhaseInitializer, self).__init__(units, self.phase.shape[0]) 74 | 75 | def to_np(self) -> np.ndarray: 76 | return self.phase.astype(NP_FLOAT) 77 | 78 | 79 | class HaarRandomPhaseInitializer(MeshPhaseInitializer): 80 | """ 81 | Haar-random initialization of rectangular and triangular mesh architectures. 82 | 83 | Args: 84 | units: Input dimension, :math:`N` 85 | num_layers: Number of layers, :math:`L` 86 | hadamard: Whether to use Hadamard convention 87 | tri: Initializer for the triangular mesh architecture 88 | """ 89 | 90 | def __init__(self, units: int, num_layers: int = None, hadamard: bool = False, tri: bool = False): 91 | self.tri = tri 92 | if self.tri: 93 | self.num_layers = 2 * units - 3 94 | else: 95 | self.num_layers = units if not num_layers else num_layers 96 | self.hadamard = hadamard 97 | super(HaarRandomPhaseInitializer, self).__init__(units, self.num_layers) 98 | 99 | def to_np(self) -> np.ndarray: 100 | theta_0, theta_1 = get_haar_theta(self.units, self.num_layers, hadamard=self.hadamard, tri=self.tri) 101 | theta = np.zeros((self.num_layers, self.units // 2)) 102 | theta[::2, :] = theta_0 103 | if self.units % 2: 104 | theta[1::2, :] = theta_1 105 | else: 106 | theta[1::2, :-1] = theta_1 107 | return theta.astype(NP_FLOAT) 108 | 109 | 110 | class PRMPhaseInitializer(MeshPhaseInitializer): 111 | def __init__(self, units: int, hadamard: bool, tunable_layers_per_block: Optional[int] = None): 112 | """ 113 | A useful initialization of permuting mesh architectures based on the Haar random initialization above. 114 | This currently only works if using default permuting mesh architecture or setting :math:`tunable_layers_per_block`. 115 | 116 | Args: 117 | units: Input dimension, :math:`N` 118 | hadamard: Whether to use Hadamard convention 119 | tunable_layers_per_block: Number of tunable layers per block (same behavior as :code:`PermutingRectangularMeshModel`). 120 | """ 121 | self.tunable_block_sizes, _ = get_default_coarse_grain_block_sizes(units) if tunable_layers_per_block is None \ 122 | else get_efficient_coarse_grain_block_sizes(units, tunable_layers_per_block) 123 | self.hadamard = hadamard 124 | self.num_layers = int(np.sum(self.tunable_block_sizes)) 125 | super(PRMPhaseInitializer, self).__init__(units, self.num_layers) 126 | 127 | def to_np(self) -> np.ndarray: 128 | thetas = [] 129 | for block_size in self.tunable_block_sizes: 130 | theta_0, theta_1 = get_haar_theta(self.units, block_size, hadamard=self.hadamard) 131 | theta = np.zeros((block_size, self.units // 2)) 132 | theta[::2, :] = theta_0 133 | if self.units % 2: 134 | theta[1::2, :] = theta_1 135 | else: 136 | theta[1::2, :-1] = theta_1 137 | thetas.append(theta.astype(NP_FLOAT)) 138 | return np.vstack(thetas) 139 | 140 | 141 | class UniformRandomPhaseInitializer(MeshPhaseInitializer): 142 | def __init__(self, units: int, num_layers: int, max_phase, min_phase: float = 0): 143 | """ 144 | Defines a uniform random initializer up to some maximum phase, 145 | e.g. :math:`\\theta \in [0, \pi]` or :math:`\phi \in [0, 2\pi)`. 146 | 147 | Args: 148 | units: Input dimension, :math:`N`. 149 | num_layers: Number of layers, :math:`L`. 150 | max_phase: Maximum phase 151 | min_phase: Minimum phase (usually 0) 152 | """ 153 | self.units = units 154 | self.num_layers = units 155 | self.max_phase = max_phase 156 | self.min_phase = min_phase 157 | super(UniformRandomPhaseInitializer, self).__init__(units, num_layers) 158 | 159 | def to_np(self) -> np.ndarray: 160 | phase = (self.max_phase - self.min_phase) * np.random.rand(self.num_layers, self.units // 2) + self.min_phase 161 | return phase.astype(NP_FLOAT) 162 | 163 | 164 | class ConstantPhaseInitializer(MeshPhaseInitializer): 165 | def __init__(self, units: int, num_layers: int, constant_phase: float): 166 | """ 167 | 168 | Args: 169 | units: Input dimension, :math:`N` 170 | num_layers: Number of layers, :math:`L` 171 | constant_phase: The constant phase to set all array elements 172 | """ 173 | self.constant_phase = constant_phase 174 | super(ConstantPhaseInitializer, self).__init__(units, num_layers) 175 | 176 | def to_np(self) -> np.ndarray: 177 | return self.constant_phase * np.ones((self.num_layers, self.units // 2)) 178 | 179 | 180 | def get_haar_theta(units: int, num_layers: int, hadamard: bool, 181 | tri: bool = False) -> Union[Tuple[np.ndarray, np.ndarray], 182 | Tuple[tf.Variable, tf.Variable], 183 | tf.Variable]: 184 | if tri: 185 | alpha_rows = np.repeat(np.linspace(1, units - 1, units - 1)[:, np.newaxis], units * 2 - 3, axis=1).T 186 | theta_0_root = 2 * alpha_rows[::2, ::2] 187 | theta_1_root = 2 * alpha_rows[1::2, 1::2] 188 | else: 189 | alpha_checkerboard = get_alpha_checkerboard_general(units, num_layers) 190 | theta_0_root = 2 * alpha_checkerboard.T[::2, ::2] 191 | theta_1_root = 2 * alpha_checkerboard.T[1::2, 1::2] 192 | theta_0_init = 2 * np.arcsin(np.random.rand(*theta_0_root.shape) ** (1 / theta_0_root)) 193 | theta_1_init = 2 * np.arcsin(np.random.rand(*theta_1_root.shape) ** (1 / theta_1_root)) 194 | if not hadamard: 195 | theta_0_init = np.pi - theta_0_init 196 | theta_1_init = np.pi - theta_1_init 197 | return theta_0_init.astype(dtype=NP_FLOAT), theta_1_init.astype(dtype=NP_FLOAT) 198 | 199 | 200 | def get_ortho_haar_theta(units: int, num_layers: int, 201 | hadamard: bool) -> Union[Tuple[np.ndarray, np.ndarray], 202 | Tuple[tf.Variable, tf.Variable], 203 | tf.Variable]: 204 | alpha_checkerboard = get_alpha_checkerboard_general(units, num_layers) 205 | theta_0_root = alpha_checkerboard.T[::2, ::2] - 1 206 | theta_1_root = alpha_checkerboard.T[1::2, 1::2] - 1 207 | theta_0_init = 2 * np.arcsin(betaincinv(0.5 * theta_0_root, 0.5, np.random.rand(*theta_0_root.shape))) 208 | theta_1_init = 2 * np.arcsin(betaincinv(0.5 * theta_1_root, 0.5, np.random.rand(*theta_1_root.shape))) 209 | if not hadamard: 210 | theta_0_init = np.pi - theta_0_init 211 | theta_1_init = np.pi - theta_1_init 212 | return theta_0_init.astype(dtype=NP_FLOAT), theta_1_init.astype(dtype=NP_FLOAT) 213 | 214 | 215 | def get_initializer(units: int, num_layers: int, initializer_name: str, 216 | hadamard: bool = False, testing: bool = False) -> MeshPhaseInitializer: 217 | if testing: 218 | np.random.seed(TEST_SEED) 219 | initializer_name_to_initializer = { 220 | 'haar_rect': HaarRandomPhaseInitializer(units, num_layers, hadamard), 221 | 'haar_tri': HaarRandomPhaseInitializer(units, num_layers, hadamard, tri=True), 222 | 'haar_prm': PRMPhaseInitializer(units, hadamard=hadamard), 223 | 'random_phi': UniformRandomPhaseInitializer(units, num_layers, 2 * np.pi), 224 | 'random_gamma': UniformRandomPhaseInitializer(2 * units, 1, 2 * np.pi), 225 | 'constant_gamma': UniformRandomPhaseInitializer(2 * units, 1, 0.0), 226 | 'constant_max_gamma': UniformRandomPhaseInitializer(2 * units, 1, 2 * np.pi), 227 | 'random_constant': ConstantPhaseInitializer(units, num_layers, np.pi * np.random.rand()), 228 | 'random_theta': UniformRandomPhaseInitializer(units, num_layers, np.pi), 229 | 'constant_phi': ConstantPhaseInitializer(units, num_layers, 0.0), 230 | 'constant_max_phi': ConstantPhaseInitializer(units, num_layers, 2 * np.pi), 231 | 'bar': ConstantPhaseInitializer(units, num_layers, 0.0 if hadamard else np.pi), 232 | 'cross': ConstantPhaseInitializer(units, num_layers, np.pi if hadamard else 0), 233 | 'transmissive': UniformRandomPhaseInitializer(units, num_layers, 234 | min_phase=np.pi / 2, 235 | max_phase=np.pi) 236 | } 237 | return initializer_name_to_initializer[initializer_name] 238 | -------------------------------------------------------------------------------- /neurophox/meshmodel.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, Union, Tuple, List 2 | import numpy as np 3 | 4 | try: 5 | import torch 6 | from torch.nn import Parameter 7 | except ImportError: 8 | pass 9 | from .helpers import butterfly_permutation, grid_permutation, to_stripe_array, prm_permutation, \ 10 | get_efficient_coarse_grain_block_sizes, get_default_coarse_grain_block_sizes 11 | from .initializers import get_initializer, MeshPhaseInitializer, PhaseInitializer 12 | from .config import BLOCH, TEST_SEED 13 | 14 | 15 | class MeshModel: 16 | """Any feedforward mesh model of :math:`N` inputs/outputs and :math:`L` layers. 17 | 18 | Args: 19 | perm_idx: A numpy array of :math:`N \\times L` permutation indices for all layers of the mesh 20 | hadamard: Whether to use Hadamard convention 21 | num_tunable: A numpy array of :math:`L` integers, where for layer :math:`\ell`, :math:`M_\ell \leq \\lfloor N / 2\\rfloor`, used to defined the phase shift mask. 22 | bs_error: Beamsplitter error (ignore for pure machine learning applications) 23 | testing: Use a seed for randomizing error (ignore for pure machine learning applications) 24 | use_different_errors: Use different errors for the left and right beamsplitter errors 25 | theta_init: Initializer for :code:`theta` (:math:`\\boldsymbol{\\theta}` or :math:`\\theta_{n\ell}`) 26 | a :code:`str`, :code:`ndarray`, or tuple of the form :code:`(theta_init, theta_fn)`. 27 | phi_init: Initializer for :code:`phi` (:math:`\\boldsymbol{\\phi}` or :math:`\\phi_{n\ell}`): 28 | a :code:`str`, :code:`ndarray`, or tuple of the form :code:`(phi_init, phi_fn)`. 29 | gamma_init: Initializer for :code:`gamma` (:math:`\\boldsymbol{\\gamma}` or :math:`\\gamma_{n}`): 30 | a :code:`str`, :code:`ndarray`, or tuple of the form :code:`(gamma_init, gamma_fn)`. 31 | basis: Phase basis to use for controlling each pairwise unitary (simulated interferometer) in the mesh 32 | """ 33 | 34 | def __init__(self, perm_idx: np.ndarray, hadamard: bool = False, num_tunable: Optional[np.ndarray] = None, 35 | bs_error: float = 0.0, testing: bool = False, use_different_errors: bool = False, 36 | theta_init: Union[str, tuple, np.ndarray] = "random_theta", 37 | phi_init: Union[str, tuple, np.ndarray] = "random_phi", 38 | gamma_init: Union[str, np.ndarray] = "random_gamma", basis: str = BLOCH): 39 | self.units = perm_idx.shape[1] 40 | self.num_layers = perm_idx.shape[0] - 1 41 | self.perm_idx = perm_idx 42 | self.num_tunable = num_tunable if num_tunable is not None else self.units // 2 * np.ones((self.num_layers,)) 43 | self.hadamard = hadamard 44 | self.bs_error = bs_error 45 | self.testing = testing 46 | self.use_different_errors = use_different_errors 47 | self.mask = np.zeros((self.num_layers, self.units // 2)) 48 | self.theta_init, self.theta_fn = (theta_init, None) if not isinstance(theta_init, tuple) else theta_init 49 | self.phi_init, self.phi_fn = (phi_init, None) if not isinstance(phi_init, tuple) else phi_init 50 | self.gamma_init, self.gamma_fn = (gamma_init, None) if not isinstance(gamma_init, tuple) else gamma_init 51 | self.basis = basis 52 | for layer in range(self.num_layers): 53 | self.mask[layer][:int(self.num_tunable[layer])] = 1 54 | if self.num_tunable.shape[0] != self.num_layers: 55 | raise ValueError("num_mzis, perm_idx num_layers mismatch.") 56 | if self.units < 2: 57 | raise ValueError("units must be at least 2.") 58 | 59 | @property 60 | def init(self) -> Tuple[MeshPhaseInitializer, MeshPhaseInitializer, MeshPhaseInitializer]: 61 | """ 62 | Returns: 63 | Initializers for :math:`\\boldsymbol{\\theta}, \\boldsymbol{\\phi}, \gamma_n`. 64 | """ 65 | if not isinstance(self.theta_init, np.ndarray): 66 | theta_init = get_initializer(self.units, self.num_layers, self.theta_init, self.hadamard, self.testing) 67 | else: 68 | theta_init = PhaseInitializer(self.theta_init, self.units) 69 | if not isinstance(self.phi_init, np.ndarray): 70 | phi_init = get_initializer(self.units, self.num_layers, self.phi_init, self.hadamard, self.testing) 71 | else: 72 | phi_init = PhaseInitializer(self.phi_init, self.units) 73 | if not isinstance(self.gamma_init, np.ndarray): 74 | gamma_init = get_initializer(self.units, self.num_layers, self.gamma_init, self.hadamard, self.testing) 75 | else: 76 | gamma_init = PhaseInitializer(self.gamma_init, self.units) 77 | return theta_init, phi_init, gamma_init 78 | 79 | @property 80 | def mzi_error_matrices(self) -> Tuple[np.ndarray, np.ndarray]: 81 | """ 82 | 83 | Returns: 84 | Error numpy arrays for Numpy :code:`MeshNumpyLayer` 85 | """ 86 | if self.testing: 87 | np.random.seed(TEST_SEED) 88 | mask = self.mask if self.mask is not None else np.ones((self.num_layers, self.units // 2)) 89 | if isinstance(self.bs_error, float) or isinstance(self.bs_error, int): 90 | e_l = np.random.randn(self.num_layers, self.units // 2) * self.bs_error 91 | if self.use_different_errors: 92 | if self.testing: 93 | np.random.seed(TEST_SEED + 1) 94 | e_r = np.random.randn(self.num_layers, self.units // 2) * self.bs_error 95 | else: 96 | e_r = e_l 97 | elif isinstance(self.bs_error, np.ndarray): 98 | if self.bs_error.shape != self.mask.shape: 99 | raise AttributeError('bs_error.shape and mask.shape should be the same.') 100 | e_l = e_r = self.bs_error 101 | elif isinstance(self.bs_error, tuple): 102 | if self.bs_error[0].shape != self.mask.shape or self.bs_error[1].shape != self.mask.shape: 103 | raise AttributeError('bs_error.shape and mask.shape should be the same.') 104 | e_l, e_r = self.bs_error 105 | else: 106 | raise TypeError('bs_error must be float, ndarray or (ndarray, ndarray).') 107 | return e_l * mask, e_r * mask 108 | 109 | @property 110 | def mzi_error_tensors(self): 111 | e_l, e_r = self.mzi_error_matrices 112 | 113 | cc = 2 * to_stripe_array(np.cos(np.pi / 4 + e_l) * np.cos(np.pi / 4 + e_r), self.units) 114 | cs = 2 * to_stripe_array(np.cos(np.pi / 4 + e_l) * np.sin(np.pi / 4 + e_r), self.units) 115 | sc = 2 * to_stripe_array(np.sin(np.pi / 4 + e_l) * np.cos(np.pi / 4 + e_r), self.units) 116 | ss = 2 * to_stripe_array(np.sin(np.pi / 4 + e_l) * np.sin(np.pi / 4 + e_r), self.units) 117 | 118 | return ss, cs, sc, cc 119 | 120 | 121 | class RectangularMeshModel(MeshModel): 122 | """Rectangular mesh 123 | 124 | The rectangular mesh contains :math:`N` inputs/outputs and :math:`L` layers in rectangular grid arrangement 125 | of pairwise unitary operators to implement :math:`U \in \mathrm{U}(N)`. 126 | 127 | Args: 128 | units: Input dimension, :math:`N` 129 | num_layers: Number of layers, :math:`L` 130 | hadamard: Hadamard convention 131 | bs_error: Beamsplitter layer 132 | basis: Phase basis to use for controlling each pairwise unitary (simulated interferometer) in the mesh 133 | theta_init: Initializer for :code:`theta` (:math:`\\boldsymbol{\\theta}` or :math:`\\theta_{n\ell}`), 134 | see :code:`MeshModel`. 135 | phi_init: Initializer for :code:`phi` (:math:`\\boldsymbol{\\phi}` or :math:`\\phi_{n\ell}`), 136 | see :code:`MeshModel`. 137 | gamma_init: Initializer for :code:`gamma` (:math:`\\boldsymbol{\\gamma}` or :math:`\\gamma_{n}`), 138 | see :code:`MeshModel`. 139 | """ 140 | 141 | def __init__(self, units: int, num_layers: int = None, hadamard: bool = False, bs_error: float = 0.0, 142 | basis: str = BLOCH, theta_init: Union[str, tuple, np.ndarray] = "haar_rect", 143 | phi_init: Union[str, tuple, np.ndarray] = "random_phi", 144 | gamma_init: Union[str, tuple, np.ndarray] = "random_gamma"): 145 | self.num_layers = num_layers if num_layers else units 146 | perm_idx = grid_permutation(units, self.num_layers).astype(np.int32) 147 | num_mzis = (np.ones((self.num_layers,)) * units // 2).astype(np.int32) 148 | num_mzis[1::2] = (units - 1) // 2 149 | super(RectangularMeshModel, self).__init__(perm_idx, 150 | hadamard=hadamard, 151 | bs_error=bs_error, 152 | num_tunable=num_mzis, 153 | theta_init=theta_init, 154 | phi_init=phi_init, 155 | gamma_init=gamma_init, 156 | basis=basis) 157 | 158 | 159 | class TriangularMeshModel(MeshModel): 160 | """Triangular mesh 161 | 162 | The triangular mesh contains :math:`N` inputs/outputs and :math:`L = 2N - 3` layers in triangular grid arrangement 163 | of pairwise unitary operators to implement any :math:`U \in \mathrm{U}(N)`. 164 | 165 | Args: 166 | units: Input dimension, :math:`N` 167 | hadamard: Hadamard convention 168 | bs_error: Beamsplitter layer 169 | basis: Phase basis to use for controlling each pairwise unitary (simulated interferometer) in the mesh 170 | theta_init: Initializer for :code:`theta` (:math:`\\boldsymbol{\\theta}` or :math:`\\theta_{n\ell}`), 171 | see :code:`MeshModel`. 172 | phi_init: Initializer for :code:`phi` (:math:`\\boldsymbol{\\phi}` or :math:`\\phi_{n\ell}`), 173 | see :code:`MeshModel`. 174 | gamma_init: Initializer for :code:`gamma` (:math:`\\boldsymbol{\\gamma}` or :math:`\\gamma_{n}`), 175 | see :code:`MeshModel`. 176 | """ 177 | 178 | def __init__(self, units: int, hadamard: bool = False, bs_error: float = 0.0, basis: str = BLOCH, 179 | theta_init: Union[str, tuple, np.ndarray] = "haar_tri", 180 | phi_init: Union[str, tuple, np.ndarray] = "random_phi", 181 | gamma_init: Union[str, tuple, np.ndarray] = "random_gamma"): 182 | perm_idx = grid_permutation(units, 2 * units - 3).astype(np.int32) 183 | num_mzis = ((np.hstack([np.arange(1, units), np.arange(units - 2, 0, -1)]) + 1) // 2).astype(np.int32) 184 | super(TriangularMeshModel, self).__init__(perm_idx, 185 | hadamard=hadamard, 186 | bs_error=bs_error, 187 | num_tunable=num_mzis, 188 | theta_init=theta_init, 189 | phi_init=phi_init, 190 | gamma_init=gamma_init, 191 | basis=basis) 192 | 193 | 194 | class ButterflyMeshModel(MeshModel): 195 | """Butterfly mesh 196 | 197 | The butterfly mesh contains :math:`L` layers and :math:`N = 2^L` inputs/outputs to implement :math:`U \in \mathrm{U}(N)`. 198 | Unlike the triangular and full (:math:`L = N`) rectangular mesh, the butterfly mesh is not universal. However, 199 | it has attractive properties for efficient machine learning and compact photonic implementations of unitary mesh models. 200 | 201 | Args: 202 | num_layers: Number of layers, :math:`L` 203 | hadamard: Hadamard convention 204 | bs_error: Beamsplitter layer 205 | theta_init: Initializer for :code:`theta` (:math:`\\boldsymbol{\\theta}` or :math:`\\theta_{n\ell}`), 206 | see :code:`MeshModel`. 207 | phi_init: Initializer for :code:`phi` (:math:`\\boldsymbol{\\phi}` or :math:`\\phi_{n\ell}`), 208 | see :code:`MeshModel`. 209 | gamma_init: Initializer for :code:`gamma` (:math:`\\boldsymbol{\\gamma}` or :math:`\\gamma_{n}`), 210 | see :code:`MeshModel`. 211 | """ 212 | 213 | def __init__(self, num_layers: int, hadamard: bool = False, 214 | bs_error: float = 0.0, basis: str = BLOCH, 215 | theta_init: Union[str, tuple, np.ndarray] = "random_theta", 216 | phi_init: Union[str, tuple, np.ndarray] = "random_phi", 217 | gamma_init: Union[str, tuple, np.ndarray] = "random_gamma"): 218 | super(ButterflyMeshModel, self).__init__(butterfly_permutation(num_layers), 219 | hadamard=hadamard, 220 | bs_error=bs_error, 221 | basis=basis, 222 | theta_init=theta_init, 223 | phi_init=phi_init, 224 | gamma_init=gamma_init) 225 | 226 | 227 | class PermutingRectangularMeshModel(MeshModel): 228 | """Permuting rectangular mesh model 229 | 230 | Args: 231 | units: Input dimension, :math:`N` 232 | tunable_layers_per_block: The number of tunable layers per block (overrides `num_tunable_layers_list`, `sampling_frequencies`) 233 | num_tunable_layers_list: Number of tunable layers in each block in order from left to right 234 | sampling_frequencies: Frequencies of sampling frequencies between the tunable layers 235 | bs_error: Photonic error in the beamsplitter 236 | hadamard: Whether to use hadamard convention (otherwise use beamsplitter convention) 237 | theta_init: Initializer for :code:`theta` (:math:`\\boldsymbol{\\theta}` or :math:`\\theta_{n\ell}`), 238 | see :code:`MeshModel`. 239 | phi_init: Initializer for :code:`phi` (:math:`\\boldsymbol{\\phi}` or :math:`\\phi_{n\ell}`), 240 | see :code:`MeshModel`. 241 | gamma_init: Initializer for :code:`gamma` (:math:`\\boldsymbol{\\gamma}` or :math:`\\gamma_{n}`), 242 | see :code:`MeshModel`. 243 | """ 244 | 245 | def __init__(self, units: int, tunable_layers_per_block: int = None, 246 | num_tunable_layers_list: Optional[List[int]] = None, sampling_frequencies: Optional[List[int]] = None, 247 | bs_error: float = 0.0, hadamard: bool = False, theta_init: Union[str, tuple, np.ndarray] = 'haar_prm', 248 | phi_init: Union[str, tuple, np.ndarray] = 'random_phi', 249 | gamma_init: Union[str, tuple, np.ndarray] = 'random_gamma'): 250 | 251 | if tunable_layers_per_block is not None: 252 | self.block_sizes, self.sampling_frequencies = get_efficient_coarse_grain_block_sizes( 253 | units=units, 254 | tunable_layers_per_block=tunable_layers_per_block 255 | ) 256 | elif sampling_frequencies is None or num_tunable_layers_list is None: 257 | self.block_sizes, self.sampling_frequencies = get_default_coarse_grain_block_sizes(units) 258 | else: 259 | self.block_sizes, self.sampling_frequencies = num_tunable_layers_list, sampling_frequencies 260 | 261 | num_mzis_list = [] 262 | for block_size in self.block_sizes: 263 | num_mzis_list.append((np.ones((block_size,)) * units // 2).astype(np.int32)) 264 | num_mzis_list[-1][1::2] = (units - 1) // 2 265 | num_mzis = np.hstack(num_mzis_list) 266 | 267 | super(PermutingRectangularMeshModel, self).__init__( 268 | perm_idx=prm_permutation(units=units, tunable_block_sizes=self.block_sizes, 269 | sampling_frequencies=self.sampling_frequencies, butterfly=False), 270 | num_tunable=num_mzis, 271 | hadamard=hadamard, 272 | bs_error=bs_error, 273 | theta_init=theta_init, 274 | phi_init=phi_init, 275 | gamma_init=gamma_init 276 | ) 277 | -------------------------------------------------------------------------------- /neurophox/ml/__init__.py: -------------------------------------------------------------------------------- 1 | from .linear import LinearMultiModelRunner 2 | from .nonlinearities import * 3 | -------------------------------------------------------------------------------- /neurophox/ml/linear.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, Callable, List, Union 2 | 3 | import numpy as np 4 | import tensorflow as tf 5 | import pickle 6 | 7 | from ..config import TF_COMPLEX, NP_COMPLEX 8 | from ..helpers import random_gaussian_batch 9 | from ..tensorflow import MeshLayer, SVD 10 | 11 | 12 | def complex_mse(y_true: tf.Tensor, y_pred: tf.Tensor): 13 | """ 14 | 15 | Args: 16 | y_true: The true labels, :math:`V \in \mathbb{C}^{B \\times N}` 17 | y_pred: The true labels, :math:`\\widehat{V} \in \mathbb{C}^{B \\times N}` 18 | 19 | Returns: 20 | The complex mean squared error :math:`\\boldsymbol{e} \in \mathbb{R}^B`, 21 | where given example :math:`\\widehat{V}_i \in \mathbb{C}^N`, 22 | we have :math:`e_i = \\frac{\|V_i - \\widehat{V}_i\|^2}{N}`. 23 | 24 | """ 25 | real_loss = tf.losses.mse(tf.math.real(y_true), tf.math.real(y_pred)) 26 | imag_loss = tf.losses.mse(tf.math.imag(y_true), tf.math.imag(y_pred)) 27 | return (real_loss + imag_loss) / 2 28 | 29 | 30 | def normalized_fidelity(u: tf.Tensor, u_hat: tf.Tensor): 31 | """ 32 | 33 | Args: 34 | u: The true (target) unitary, :math:`U \in \mathrm{U}(N)` 35 | u_hat: The estimated unitary (not necessarily unitary), :math:`\\widehat{U}` 36 | 37 | Returns: 38 | The normalized fidelity independent of the norm of :math:`\\widehat{U}`. 39 | 40 | """ 41 | 42 | def trf(x: tf.Tensor, y=None): 43 | y = x if y is None else y 44 | trace = tf.linalg.trace(tf.transpose(tf.math.conj(x)) @ y) 45 | return tf.math.real(trace) ** 2 + tf.math.imag(trace) ** 2 46 | 47 | return trf(u_hat, u) / trf(u_hat) 48 | 49 | 50 | class LinearMultiModelRunner: 51 | """ 52 | Complex mean square error linear optimization experiment that can run and track multiple model optimizations in parallel. 53 | 54 | Args: 55 | experiment_name: Name of the experiment 56 | layer_names: List of layer names 57 | layers: List of transformer layers 58 | optimizer: Optimizer for all layers or list of optimizers for each layer 59 | batch_size: Batch size for the optimization 60 | iterations_per_epoch: Iterations per epoch 61 | iterations_per_tb_update: Iterations per update of TensorBoard 62 | logdir: Logging directory for TensorBoard to track losses of each layer (default to `None` for no logging) 63 | train_on_test: Use same training and testing set 64 | store_params: Store params during the training for visualization later 65 | """ 66 | def __init__(self, experiment_name: str, layer_names: List[str], 67 | layers: List[MeshLayer], optimizer: Union[tf.keras.optimizers.Optimizer, List[tf.keras.optimizers.Optimizer]], 68 | batch_size: int, iterations_per_epoch: int=50, iterations_per_tb_update: int=5, 69 | logdir: Optional[str]=None, train_on_test: bool=False, store_params: bool=True): # e.g., logdir=/data/tensorboard/neurophox/ 70 | self.losses = {name: [] for name in layer_names} 71 | self.results = {name: [] for name in layer_names} 72 | self.layer_names = layer_names 73 | self.layers = layers 74 | self.optimizers = optimizer if isinstance(optimizer, List) else [optimizer for _ in layer_names] 75 | if not (len(layer_names) == len(layers) and len(layers) == len(self.optimizers)): 76 | raise ValueError("layer_names, layers, and optimizers must all be the same length") 77 | self.batch_size = batch_size 78 | self.iters = 0 79 | self.iterations_per_epoch = iterations_per_epoch 80 | self.iterations_per_tb_update = iterations_per_tb_update 81 | self.experiment_name = experiment_name 82 | self.logdir = logdir 83 | self.train_on_test = train_on_test 84 | self.store_params = store_params 85 | if self.logdir: 86 | self.summary_writers = {name: tf.summary.create_file_writer( 87 | f'{self.logdir}/{experiment_name}/{name}/' 88 | ) for name in layer_names} 89 | 90 | def iterate(self, target_unitary: np.ndarray, cost_fn: Callable = complex_mse): 91 | """ 92 | Run gradient update toward a target unitary :math:`U`. 93 | 94 | Args: 95 | target_unitary: Target unitary, :math:`U`. 96 | cost_fn: Cost function for linear model (default to complex mean square error) 97 | 98 | """ 99 | if self.train_on_test: 100 | x_train, y_train = tf.eye(self.layers[0].units, dtype=TF_COMPLEX),\ 101 | tf.convert_to_tensor(target_unitary, dtype=TF_COMPLEX) 102 | else: 103 | x_train, y_train = generate_keras_batch(self.layers[0].units, target_unitary, self.batch_size) 104 | for name, layer, optimizer in zip(self.layer_names, self.layers, self.optimizers): 105 | with tf.GradientTape() as tape: 106 | loss = cost_fn(layer(x_train), y_train) 107 | grads = tape.gradient(loss, layer.trainable_variables) 108 | optimizer.apply_gradients(grads_and_vars=zip(grads, layer.trainable_variables)) 109 | self.losses[name].append(tf.reduce_sum( 110 | cost_fn(layer(tf.eye(layer.units, dtype=TF_COMPLEX)), 111 | tf.convert_to_tensor(target_unitary.astype(NP_COMPLEX)))).numpy() 112 | ) 113 | if self.iters % self.iterations_per_tb_update and self.logdir: 114 | self.update_tensorboard(name) 115 | if self.iters % self.iterations_per_epoch == 0 and self.store_params: 116 | if not isinstance(layer, SVD): 117 | phases = layer.phases 118 | mask = layer.mesh.model.mask 119 | estimate = layer(tf.eye(layer.units, dtype=TF_COMPLEX)).numpy() 120 | self.results[name].append({ 121 | "theta_list": phases.theta.param_list(mask), 122 | "phi_list": phases.phi.param_list(mask), 123 | "theta_checkerboard": phases.theta.checkerboard_arrangement, 124 | "phi_checkerboard": phases.phi.checkerboard_arrangement, 125 | "gamma": phases.gamma, 126 | "estimate_mag": np.abs(estimate), 127 | "error_mag": np.abs(estimate - target_unitary) 128 | }) 129 | 130 | self.iters += 1 131 | 132 | def update_tensorboard(self, name: str): 133 | """ 134 | Update TensorBoard variables. 135 | 136 | Args: 137 | name: Layer name corresponding to variables that are updated 138 | 139 | """ 140 | with self.summary_writers[name].as_default(): 141 | tf.summary.scalar(f'loss-{self.experiment_name}', self.losses[name][-1], step=self.iters) 142 | 143 | def run(self, num_epochs: int, target_unitary: np.ndarray, pbar: Optional[Callable]=None): 144 | """ 145 | 146 | Args: 147 | num_epochs: Number of epochs (defined in terms of `iterations_per_epoch`) 148 | target_unitary: Target unitary, :math:`U`. 149 | pbar: Progress bar (tqdm recommended) 150 | 151 | """ 152 | iterator = pbar(range(num_epochs * self.iterations_per_epoch)) if pbar else range(num_epochs * self.iterations_per_epoch) 153 | for _ in iterator: 154 | self.iterate(target_unitary) 155 | 156 | def save(self, savepath: str): 157 | """ 158 | Save results for the multi-model runner to pickle file. 159 | 160 | Args: 161 | savepath: Path to save results. 162 | """ 163 | with open(f"{savepath}/{self.experiment_name}.p", "wb") as f: 164 | pickle.dump({"losses": self.losses, "results": self.results}, f) 165 | 166 | 167 | def generate_keras_batch(units, target_unitary, batch_size): 168 | x_train = random_gaussian_batch(batch_size=batch_size, units=units) 169 | y_train = x_train @ target_unitary 170 | return tf.convert_to_tensor(x_train, dtype=TF_COMPLEX), tf.convert_to_tensor(y_train, dtype=TF_COMPLEX) 171 | -------------------------------------------------------------------------------- /neurophox/ml/nonlinearities.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from ..config import TF_COMPLEX 3 | 4 | 5 | def cnormsq(inputs: tf.Tensor) -> tf.Tensor: 6 | """ 7 | 8 | Args: 9 | inputs: The input tensor, :math:`V`. 10 | 11 | Returns: 12 | An output tensor that performs elementwise absolute value squared operation, :math:`f(V) = |V|^2`. 13 | 14 | """ 15 | return tf.cast(tf.square(tf.math.real(inputs)) + tf.square(tf.math.imag(inputs)), TF_COMPLEX) 16 | 17 | 18 | def cnorm(inputs: tf.Tensor) -> tf.Tensor: 19 | """ 20 | 21 | Args: 22 | inputs: The input tensor, :math:`V`. 23 | 24 | Returns: 25 | An output tensor that performs elementwise absolute value operation, :math:`f(V) = |V|`. 26 | 27 | """ 28 | return tf.cast(tf.sqrt(tf.square(tf.math.real(inputs)) + tf.square(tf.math.imag(inputs))), TF_COMPLEX) 29 | -------------------------------------------------------------------------------- /neurophox/numpy/__init__.py: -------------------------------------------------------------------------------- 1 | from .layers import * 2 | -------------------------------------------------------------------------------- /neurophox/numpy/layers.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from typing import Optional, List 3 | 4 | from .generic import MeshNumpyLayer, MeshPhases 5 | from ..config import DEFAULT_BASIS 6 | from ..meshmodel import RectangularMeshModel, TriangularMeshModel, PermutingRectangularMeshModel, ButterflyMeshModel 7 | from ..helpers import grid_viz_permutation 8 | 9 | 10 | class RMNumpy(MeshNumpyLayer): 11 | def __init__(self, units: int, num_layers: int = None, hadamard: bool = False, basis: str = DEFAULT_BASIS, 12 | bs_error: float = 0.0, theta_init="haar_rect", 13 | phi_init="random_phi", gamma_init="random_gamma"): 14 | """Rectangular mesh network layer for unitary operators implemented in numpy 15 | 16 | Args: 17 | units: The dimension of the unitary matrix (:math:`N`) 18 | num_layers: The number of layers (:math:`L`) of the mesh 19 | hadamard: Hadamard convention for the beamsplitters 20 | basis: Phase basis to use 21 | bs_error: Beamsplitter split ratio error 22 | theta_init: Initializer name for :code:`theta` (:math:`\\boldsymbol{\\theta}` or :math:`\\theta_{n\ell}`) 23 | phi_init: Initializer name for :code:`phi` (:math:`\\boldsymbol{\\phi}` or :math:`\\phi_{n\ell}`) 24 | gamma_init: Initializer name for :code:`gamma` (:math:`\\boldsymbol{\\gamma}` or :math:`\\gamma_{n}`) 25 | """ 26 | super(RMNumpy, self).__init__( 27 | RectangularMeshModel(units, num_layers, hadamard, bs_error, basis, 28 | theta_init, phi_init, gamma_init) 29 | ) 30 | 31 | def propagate(self, inputs: np.ndarray, explicit: bool = False, 32 | viz_perm_idx: Optional[np.ndarray] = None) -> np.ndarray: 33 | viz_perm_idx = grid_viz_permutation(self.units, self.num_layers, 34 | flip=explicit) if viz_perm_idx is None else viz_perm_idx 35 | # viz_perm_idx = None 36 | return super(RMNumpy, self).propagate(inputs, explicit, viz_perm_idx) 37 | 38 | def inverse_propagate(self, inputs: np.ndarray, explicit: bool = False, 39 | viz_perm_idx: Optional[np.ndarray] = None) -> np.ndarray: 40 | viz_perm_idx = grid_viz_permutation(self.units, self.num_layers, 41 | flip=explicit) if viz_perm_idx is None else viz_perm_idx 42 | # viz_perm_idx = None 43 | return super(RMNumpy, self).inverse_propagate(inputs, explicit, viz_perm_idx) 44 | 45 | 46 | class TMNumpy(MeshNumpyLayer): 47 | def __init__(self, units: int, hadamard: bool = False, basis: str = DEFAULT_BASIS, 48 | bs_error: float = 0.0, 49 | theta_init="haar_tri", phi_init="random_phi", gamma_init="random_gamma"): 50 | """Triangular mesh network layer for unitary operators implemented in numpy 51 | 52 | Args: 53 | units: The dimension of the unitary matrix (:math:`N`) 54 | hadamard: Hadamard convention for the beamsplitters 55 | basis: Phase basis to use 56 | bs_error: Beamsplitter split ratio error 57 | """ 58 | super(TMNumpy, self).__init__( 59 | TriangularMeshModel(units, hadamard, bs_error, basis, 60 | theta_init, phi_init, gamma_init) 61 | ) 62 | 63 | def propagate(self, inputs: np.ndarray, explicit: bool = False, 64 | viz_perm_idx: Optional[np.ndarray] = None) -> np.ndarray: 65 | viz_perm_idx = grid_viz_permutation(self.units, self.num_layers) if viz_perm_idx is None else viz_perm_idx 66 | return super(TMNumpy, self).propagate(inputs, explicit, viz_perm_idx) 67 | 68 | def inverse_propagate(self, inputs: np.ndarray, explicit: bool = False, 69 | viz_perm_idx: Optional[np.ndarray] = None) -> np.ndarray: 70 | viz_perm_idx = grid_viz_permutation(self.units, self.num_layers) if viz_perm_idx is None else viz_perm_idx 71 | return super(TMNumpy, self).inverse_propagate(inputs, explicit, viz_perm_idx) 72 | 73 | 74 | class BMNumpy(MeshNumpyLayer): 75 | def __init__(self, num_layers: int, basis: str = DEFAULT_BASIS, 76 | bs_error: float = 0.0, hadamard: bool = False, theta_init: Optional[str] = 'random_theta', 77 | phi_init: Optional[str] = 'random_phi'): 78 | """Butterfly mesh unitary layer (currently, only :math:`2^L` units allowed) 79 | 80 | Args: 81 | num_layers: The number of layers (:math:`L`), with dimension of the unitary matrix set to (:math:`N = 2^L`) 82 | bs_error: Photonic error in the beamsplitter 83 | hadamard: Hadamard convention for the beamsplitters 84 | theta_init: Initializer name for :code:`theta` (:math:`\\boldsymbol{\\theta}` or :math:`\\theta_{n\ell}`) 85 | phi_init: Initializer name for :code:`phi` (:math:`\\boldsymbol{\\phi}` or :math:`\\phi_{n\ell}`) 86 | """ 87 | super(BMNumpy, self).__init__( 88 | ButterflyMeshModel(num_layers, hadamard, bs_error, basis, theta_init, phi_init) 89 | ) 90 | 91 | 92 | class PRMNumpy(MeshNumpyLayer): 93 | def __init__(self, units: int, tunable_layers_per_block: int = None, 94 | num_tunable_layers_list: Optional[List[int]] = None, sampling_frequencies: Optional[List[int]] = None, 95 | bs_error: float = 0.0, hadamard: bool = False, theta_init: Optional[str] = 'haar_prm', 96 | phi_init: Optional[str] = 'random_phi'): 97 | """Permuting rectangular mesh unitary layer 98 | 99 | Args: 100 | units: The dimension of the unitary matrix (:math:`N`) to be modeled by this transformer 101 | tunable_layers_per_block: The number of tunable layers per block (overrides `num_tunable_layers_list`, `sampling_frequencies`) 102 | num_tunable_layers_list: Number of tunable layers in each block in order from left to right 103 | sampling_frequencies: Frequencies of sampling frequencies between the tunable layers 104 | bs_error: Photonic error in the beamsplitter 105 | hadamard: Hadamard convention for the beamsplitters 106 | theta_init: Initializer name for :code:`theta` (:math:`\\boldsymbol{\\theta}` or :math:`\\theta_{n\ell}`) 107 | phi_init: Initializer name for :code:`phi` (:math:`\\boldsymbol{\\phi}` or :math:`\\phi_{n\ell}`) 108 | """ 109 | if theta_init == 'haar_prm' and tunable_layers_per_block is not None: 110 | raise NotImplementedError('haar_prm initializer is incompatible with setting tunable_layers_per_block.') 111 | super(PRMNumpy, self).__init__( 112 | PermutingRectangularMeshModel(units, tunable_layers_per_block, num_tunable_layers_list, 113 | sampling_frequencies, bs_error, hadamard, 114 | theta_init, phi_init) 115 | ) 116 | -------------------------------------------------------------------------------- /neurophox/tensorflow/__init__.py: -------------------------------------------------------------------------------- 1 | from .layers import * 2 | 3 | NAME_TO_TRANSFORMER_LAYER = { 4 | 'rm': RM, 5 | 'tm': TM, 6 | 'prm': PRM, 7 | 'svd': SVD 8 | } 9 | -------------------------------------------------------------------------------- /neurophox/tensorflow/generic.py: -------------------------------------------------------------------------------- 1 | from typing import List, Tuple, Optional, Callable 2 | 3 | import tensorflow as tf 4 | from tensorflow.keras.layers import Layer, Activation 5 | import numpy as np 6 | 7 | from ..numpy.generic import MeshPhases 8 | from ..meshmodel import MeshModel 9 | from ..helpers import pairwise_off_diag_permutation, plot_complex_matrix, inverse_permutation 10 | from ..config import TF_COMPLEX, BLOCH, SINGLEMODE 11 | 12 | 13 | class TransformerLayer(Layer): 14 | """Base transformer class for transformer layers (invertible functions, usually linear) 15 | 16 | Args: 17 | units: Dimension of the input to be transformed by the transformer 18 | activation: Nonlinear activation function (:code:`None` if there's no nonlinearity) 19 | """ 20 | def __init__(self, units: int, activation: Activation = None, **kwargs): 21 | self.units = units 22 | self.activation = activation 23 | super(TransformerLayer, self).__init__(**kwargs) 24 | 25 | def transform(self, inputs: tf.Tensor) -> tf.Tensor: 26 | """ 27 | Transform inputs using layer (needs to be overwritten by child classes) 28 | 29 | Args: 30 | inputs: Inputs to be transformed by layer 31 | 32 | Returns: 33 | Transformed inputs 34 | """ 35 | raise NotImplementedError("Needs to be overwritten by child class.") 36 | 37 | def inverse_transform(self, outputs: tf.Tensor) -> tf.Tensor: 38 | """ 39 | Transform outputs using layer 40 | 41 | Args: 42 | outputs: Outputs to be inverse-transformed by layer 43 | 44 | Returns: 45 | Transformed outputs 46 | """ 47 | raise NotImplementedError("Needs to be overwritten by child class.") 48 | 49 | def call(self, inputs, training=None, mask=None): 50 | outputs = self.transform(inputs) 51 | if self.activation: 52 | outputs = self.activation(outputs) 53 | return outputs 54 | 55 | @property 56 | def matrix(self): 57 | """ 58 | Shortcut of :code:`transformer.transform(np.eye(self.units))` 59 | 60 | Returns: 61 | Matrix implemented by layer 62 | """ 63 | identity_matrix = np.eye(self.units, dtype=np.complex64) 64 | return self.transform(identity_matrix).numpy() 65 | 66 | @property 67 | def inverse_matrix(self): 68 | """ 69 | Shortcut of :code:`transformer.inverse_transform(np.eye(self.units))` 70 | 71 | Returns: 72 | Inverse matrix implemented by layer 73 | """ 74 | identity_matrix = np.eye(self.units, dtype=np.complex64) 75 | return self.inverse_transform(identity_matrix).numpy() 76 | 77 | def plot(self, plt): 78 | """ 79 | Plot :code:`transformer.matrix`. 80 | 81 | Args: 82 | plt: :code:`matplotlib.pyplot` for plotting 83 | """ 84 | plot_complex_matrix(plt, self.matrix) 85 | 86 | 87 | class CompoundTransformerLayer(TransformerLayer): 88 | """Compound transformer class for unitary matrices 89 | 90 | Args: 91 | units: Dimension of the input to be transformed by the transformer 92 | transformer_list: List of :class:`Transformer` objects to apply to the inputs 93 | is_complex: Whether the input to be transformed is complex 94 | """ 95 | def __init__(self, units: int, transformer_list: List[TransformerLayer]): 96 | self.transformer_list = transformer_list 97 | super(CompoundTransformerLayer, self).__init__(units=units) 98 | 99 | def transform(self, inputs: tf.Tensor) -> tf.Tensor: 100 | """Inputs are transformed by :math:`L` transformer layers 101 | :math:`T^{(\ell)} \in \mathbb{C}^{N \\times N}` as follows: 102 | 103 | .. math:: 104 | V_{\mathrm{out}} = V_{\mathrm{in}} \prod_{\ell=1}^L T_\ell, 105 | 106 | where :math:`V_{\mathrm{out}}, V_{\mathrm{in}} \in \mathbb{C}^{M \\times N}`. 107 | 108 | Args: 109 | inputs: Input batch represented by the matrix :math:`V_{\mathrm{in}} \in \mathbb{C}^{M \\times N}` 110 | 111 | Returns: 112 | Transformed :code:`inputs`, :math:`V_{\mathrm{out}}` 113 | """ 114 | outputs = inputs 115 | for transformer in self.transformer_list: 116 | outputs = transformer.transform(outputs) 117 | return outputs 118 | 119 | def inverse_transform(self, outputs: tf.Tensor) -> tf.Tensor: 120 | """Outputs are inverse-transformed by :math:`L` transformer layers :math:`T^{(\ell)} \in \mathbb{C}^{N \\times N}` as follows: 121 | 122 | .. math:: 123 | V_{\mathrm{in}} = V_{\mathrm{out}} \prod_{\ell=L}^1 T_\ell^{-1}, 124 | 125 | where :math:`V_{\mathrm{out}}, V_{\mathrm{in}} \in \mathbb{C}^{M \\times N}`. 126 | 127 | Args: 128 | outputs: Output batch represented by the matrix :math:`V_{\mathrm{out}} \in \mathbb{C}^{M \\times N}` 129 | 130 | Returns: 131 | Transformed :code:`outputs`, :math:`V_{\mathrm{in}}` 132 | """ 133 | inputs = outputs 134 | for transformer in self.transformer_list[::-1]: 135 | inputs = transformer.inverse_transform(inputs) 136 | return inputs 137 | 138 | 139 | class PermutationLayer(TransformerLayer): 140 | """Permutation layer 141 | 142 | Args: 143 | permuted_indices: order of indices for the permutation matrix (efficient permutation representation) 144 | """ 145 | def __init__(self, permuted_indices: np.ndarray): 146 | super(PermutationLayer, self).__init__(units=permuted_indices.shape[0]) 147 | self.permuted_indices = np.asarray(permuted_indices, dtype=np.int32) 148 | self.inv_permuted_indices = inverse_permutation(self.permuted_indices) 149 | 150 | def transform(self, inputs: tf.Tensor): 151 | """ 152 | Performs the permutation for this layer represented by :math:`P` defined by `permuted_indices`: 153 | 154 | .. math:: 155 | V_{\mathrm{out}} = V_{\mathrm{in}} P, 156 | 157 | where :math:`P` is any :math:`N`-dimensional permutation and 158 | :math:`V_{\mathrm{out}}, V_{\mathrm{in}} \in \mathbb{C}^{M \\times N}`. 159 | 160 | Args: 161 | inputs: Input batch represented by the matrix :math:`V_{\mathrm{in}} \in \mathbb{C}^{M \\times N}` 162 | 163 | Returns: 164 | Permuted :code:`inputs`, :math:`V_{\mathrm{out}}` 165 | """ 166 | return tf.gather(inputs, self.permuted_indices, axis=-1) 167 | 168 | def inverse_transform(self, outputs: tf.Tensor): 169 | """ 170 | Performs the inverse permutation for this layer represented by :math:`P^{-1}` defined by `inv_permuted_indices`: 171 | 172 | .. math:: 173 | V_{\mathrm{in}} = V_{\mathrm{out}} P^{-1}, 174 | 175 | where :math:`P` is any :math:`N`-dimensional permutation and 176 | :math:`V_{\mathrm{out}}, V_{\mathrm{in}} \in \mathbb{C}^{M \\times N}`. 177 | 178 | Args: 179 | outputs: :code:`outputs` batch represented by the matrix :math:`V_{\mathrm{out}} \in \mathbb{C}^{M \\times N}` 180 | 181 | Returns: 182 | Permuted :code:`outputs`, :math:`V_{\mathrm{in}}` 183 | """ 184 | return tf.gather(outputs, self.inv_permuted_indices, axis=-1) 185 | 186 | 187 | class MeshVerticalLayer(TransformerLayer): 188 | """ 189 | Args: 190 | diag: the diagonal terms to multiply 191 | off_diag: the off-diagonal terms to multiply 192 | left_perm: the permutation for the mesh vertical layer (prior to the coupling operation) 193 | right_perm: the right permutation for the mesh vertical layer 194 | (usually for the final layer and after the coupling operation) 195 | """ 196 | def __init__(self, pairwise_perm_idx: np.ndarray, diag: tf.Tensor, off_diag: tf.Tensor, 197 | right_perm: PermutationLayer = None, left_perm: PermutationLayer = None): 198 | self.diag = diag 199 | self.off_diag = off_diag 200 | self.left_perm = left_perm 201 | self.right_perm = right_perm 202 | self.pairwise_perm_idx = pairwise_perm_idx 203 | super(MeshVerticalLayer, self).__init__(pairwise_perm_idx.shape[0]) 204 | 205 | def transform(self, inputs: tf.Tensor): 206 | """ 207 | Propagate :code:`inputs` through single layer :math:`\ell < L` 208 | (where :math:`U_\ell` represents the matrix for layer :math:`\ell`): 209 | 210 | .. math:: 211 | V_{\mathrm{out}} = V_{\mathrm{in}} U^{(\ell')}, 212 | 213 | Args: 214 | inputs: :code:`inputs` batch represented by the matrix :math:`V_{\mathrm{in}} \in \mathbb{C}^{M \\times N}` 215 | 216 | Returns: 217 | Propaged :code:`inputs` through single layer :math:`\ell` to form an array 218 | :math:`V_{\mathrm{out}} \in \mathbb{C}^{M \\times N}`. 219 | """ 220 | outputs = inputs if self.left_perm is None else self.left_perm.transform(inputs) 221 | outputs = outputs * self.diag + tf.gather(outputs * self.off_diag, self.pairwise_perm_idx, axis=-1) 222 | return outputs if self.right_perm is None else self.right_perm.transform(outputs) 223 | 224 | def inverse_transform(self, outputs: tf.Tensor): 225 | """ 226 | Inverse-propagate :code:`inputs` through single layer :math:`\ell < L` 227 | (where :math:`U_\ell` represents the matrix for layer :math:`\ell`): 228 | 229 | .. math:: 230 | V_{\mathrm{in}} = V_{\mathrm{out}} (U^{(\ell')})^\dagger, 231 | 232 | Args: 233 | outputs: :code:`outputs` batch represented by the matrix :math:`V_{\mathrm{out}} \in \mathbb{C}^{M \\times N}` 234 | 235 | Returns: 236 | Inverse propaged :code:`outputs` through single layer :math:`\ell` to form an array 237 | :math:`V_{\mathrm{in}} \in \mathbb{C}^{M \\times N}`. 238 | """ 239 | inputs = outputs if self.right_perm is None else self.right_perm.inverse_transform(outputs) 240 | diag = tf.math.conj(self.diag) 241 | off_diag = tf.gather(tf.math.conj(self.off_diag), self.pairwise_perm_idx, axis=-1) 242 | inputs = inputs * diag + tf.gather(inputs * off_diag, self.pairwise_perm_idx, axis=-1) 243 | return inputs if self.left_perm is None else self.left_perm.inverse_transform(inputs) 244 | 245 | 246 | class MeshParamTensorflow: 247 | """A class that cleanly arranges parameters into a specific arrangement that can be used to simulate any mesh 248 | 249 | Args: 250 | param: parameter to arrange in mesh 251 | units: number of inputs/outputs of the mesh 252 | """ 253 | def __init__(self, param: tf.Tensor, units: int): 254 | self.param = param 255 | self.units = units 256 | 257 | @property 258 | def single_mode_arrangement(self): 259 | """ 260 | The single-mode arrangement based on the :math:`L(\\theta)` transfer matrix for :code:`PhaseShiftUpper` 261 | is one where elements of `param` are on the even rows and all odd rows are zero. 262 | 263 | In particular, given the :code:`param` array 264 | :math:`\\boldsymbol{\\theta} = [\\boldsymbol{\\theta}_1, \\boldsymbol{\\theta}_2, \ldots \\boldsymbol{\\theta}_M]^T`, 265 | where :math:`\\boldsymbol{\\theta}_m` represent row vectors and :math:`M = \\lfloor\\frac{N}{2}\\rfloor`, the single-mode arrangement has the stripe array form 266 | :math:`\widetilde{\\boldsymbol{\\theta}} = [\\boldsymbol{\\theta}_1, \\boldsymbol{0}, \\boldsymbol{\\theta}_2, \\boldsymbol{0}, \ldots \\boldsymbol{\\theta}_N, \\boldsymbol{0}]^T`. 267 | where :math:`\widetilde{\\boldsymbol{\\theta}} \in \mathbb{R}^{N \\times L}` defines the :math:`\\boldsymbol{\\theta}` of the final mesh 268 | and :math:`\\boldsymbol{0}` represents an array of zeros of the same size as :math:`\\boldsymbol{\\theta}_n`. 269 | 270 | Returns: 271 | Single-mode arrangement array of phases 272 | 273 | """ 274 | tensor_t = tf.transpose(self.param) 275 | stripe_tensor = tf.reshape(tf.concat((tensor_t, tf.zeros_like(tensor_t)), 1), 276 | shape=(tensor_t.shape[0] * 2, tensor_t.shape[1])) 277 | if self.units % 2: 278 | return tf.concat([stripe_tensor, tf.zeros(shape=(1, tensor_t.shape[1]))], axis=0) 279 | else: 280 | return stripe_tensor 281 | 282 | @property 283 | def common_mode_arrangement(self) -> tf.Tensor: 284 | """ 285 | The common-mode arrangement based on the :math:`C(\\theta)` transfer matrix for :code:`PhaseShiftCommonMode` 286 | is one where elements of `param` are on the even rows and repeated on respective odd rows. 287 | 288 | In particular, given the :code:`param` array 289 | :math:`\\boldsymbol{\\theta} = [\\boldsymbol{\\theta}_1, \\boldsymbol{\\theta}_2, \ldots \\boldsymbol{\\theta}_M]^T`, 290 | where :math:`\\boldsymbol{\\theta}_n` represent row vectors and :math:`M = \\lfloor\\frac{N}{2}\\rfloor`, 291 | the common-mode arrangement has the stripe array form 292 | :math:`\\widetilde{\\boldsymbol{\\theta}} = [\\boldsymbol{\\theta}_1, \\boldsymbol{\\theta}_1,\\boldsymbol{\\theta}_2, \\boldsymbol{\\theta}_2, \ldots \\boldsymbol{\\theta}_N, \\boldsymbol{\\theta}_N]^T`. 293 | where :math:`\widetilde{\\boldsymbol{\\theta}} \in \mathbb{R}^{N \\times L}` defines the :math:`\\boldsymbol{\\theta}` of the final mesh. 294 | 295 | 296 | Returns: 297 | Common-mode arrangement array of phases 298 | 299 | """ 300 | phases = self.single_mode_arrangement 301 | return phases + _roll_tensor(phases) 302 | 303 | @property 304 | def differential_mode_arrangement(self) -> tf.Tensor: 305 | """ 306 | The differential-mode arrangement is based on the :math:`D(\\theta)` transfer matrix 307 | for :code:`PhaseShiftDifferentialMode`. 308 | 309 | Given the :code:`param` array 310 | :math:`\\boldsymbol{\\theta} = [\cdots \\boldsymbol{\\theta}_m \cdots]^T`, 311 | where :math:`\\boldsymbol{\\theta}_n` represent row vectors and :math:`M = \\lfloor\\frac{N}{2}\\rfloor`, the differential-mode arrangement has the form 312 | :math:`\\widetilde{\\boldsymbol{\\theta}} = \\left[\cdots \\frac{\\boldsymbol{\\theta}_m}{2}, -\\frac{\\boldsymbol{\\theta}_m}{2} \cdots \\right]^T`. 313 | where :math:`\widetilde{\\boldsymbol{\\theta}} \in \mathbb{R}^{N \\times L}` defines the :math:`\\boldsymbol{\\theta}` of the final mesh. 314 | 315 | Returns: 316 | Differential-mode arrangement array of phases 317 | 318 | """ 319 | phases = self.single_mode_arrangement 320 | return phases / 2 - _roll_tensor(phases / 2) 321 | 322 | def __add__(self, other): 323 | return MeshParamTensorflow(self.param + other.param, self.units) 324 | 325 | def __sub__(self, other): 326 | return MeshParamTensorflow(self.param - other.param, self.units) 327 | 328 | def __mul__(self, other): 329 | return MeshParamTensorflow(self.param * other.param, self.units) 330 | 331 | 332 | class MeshPhasesTensorflow: 333 | """Organizes the phases in the mesh into appropriate arrangements 334 | 335 | Args: 336 | theta: Array to be converted to :math:`\\boldsymbol{\\theta}` 337 | phi: Array to be converted to :math:`\\boldsymbol{\\phi}` 338 | gamma: Array to be converted to :math:`\\boldsymbol{\gamma}` 339 | mask: Mask over values of :code:`theta` and :code:`phi` that are not in bar state 340 | basis: Phase basis to use 341 | hadamard: Whether to use Hadamard convention 342 | theta_fn: TF-friendly phi function call to reparametrize phi (example use cases: see `neurophox.helpers`). 343 | By default, use identity function. 344 | phi_fn: TF-friendly phi function call to reparametrize phi (example use cases: see `neurophox.helpers`). 345 | By default, use identity function. 346 | gamma_fn: TF-friendly gamma function call to reparametrize gamma (example use cases: see `neurophox.helpers`). 347 | phase_loss_fn: Incorporate phase shift-dependent loss into the model. 348 | The function is of the form phase_loss_fn(phases), 349 | which returns the loss 350 | 351 | """ 352 | def __init__(self, theta: tf.Variable, phi: tf.Variable, mask: np.ndarray, gamma: tf.Variable, units: int, 353 | basis: str = SINGLEMODE, hadamard: bool = False, 354 | theta_fn: Optional[Callable] = None, phi_fn: Optional[Callable] = None, 355 | gamma_fn: Optional[Callable] = None, 356 | phase_loss_fn: Optional[Callable[[tf.Tensor], tf.Tensor]] = None): 357 | self.mask = mask if mask is not None else np.ones_like(theta) 358 | self.theta_fn = (lambda x: x) if theta_fn is None else theta_fn 359 | self.phi_fn = (lambda x: x) if phi_fn is None else phi_fn 360 | self.gamma_fn = (lambda x: x) if gamma_fn is None else gamma_fn 361 | self.theta = MeshParamTensorflow(self.theta_fn(theta) * mask + (1 - mask) * (1 - hadamard) * np.pi, units=units) 362 | self.phi = MeshParamTensorflow(self.phi_fn(phi) * mask + (1 - mask) * (1 - hadamard) * np.pi, units=units) 363 | self.gamma = self.gamma_fn(gamma) 364 | self.basis = basis 365 | self.phase_loss_fn = (lambda x: 0) if phase_loss_fn is None else phase_loss_fn 366 | self.phase_fn = lambda phase: tf.complex(tf.cos(phase), tf.sin(phase)) * (1 - _to_complex(self.phase_loss_fn(phase))) 367 | self.input_phase_shift_layer = self.phase_fn(self.gamma) 368 | if self.theta.param.shape != self.phi.param.shape: 369 | raise ValueError("Internal phases (theta) and external phases (phi) need to have the same shape.") 370 | 371 | @property 372 | def internal_phase_shifts(self): 373 | """ 374 | 375 | The internal phase shift matrix of the mesh corresponds to an `L \\times N` array of phase shifts 376 | (in between beamsplitters, thus internal) where :math:`L` is number of layers and :math:`N` is number of inputs/outputs 377 | 378 | Returns: 379 | Internal phase shift matrix corresponding to :math:`\\boldsymbol{\\theta}` 380 | """ 381 | if self.basis == BLOCH: 382 | return self.theta.differential_mode_arrangement 383 | elif self.basis == SINGLEMODE: 384 | return self.theta.single_mode_arrangement 385 | else: 386 | raise NotImplementedError(f"{self.basis} is not yet supported or invalid.") 387 | 388 | @property 389 | def external_phase_shifts(self): 390 | """The external phase shift matrix of the mesh corresponds to an `L \\times N` array of phase shifts 391 | (outside of beamsplitters, thus external) where :math:`L` is number of layers and :math:`N` is number of inputs/outputs 392 | 393 | Returns: 394 | External phase shift matrix corresponding to :math:`\\boldsymbol{\\phi}` 395 | """ 396 | if self.basis == BLOCH or self.basis == SINGLEMODE: 397 | return self.phi.single_mode_arrangement 398 | else: 399 | raise NotImplementedError(f"{self.basis} is not yet supported or invalid.") 400 | 401 | @property 402 | def internal_phase_shift_layers(self): 403 | """Elementwise applying complex exponential to :code:`internal_phase_shifts`. 404 | 405 | Returns: 406 | Internal phase shift layers corresponding to :math:`\\boldsymbol{\\theta}` 407 | """ 408 | return self.phase_fn(self.internal_phase_shifts) 409 | 410 | @property 411 | def external_phase_shift_layers(self): 412 | """Elementwise applying complex exponential to :code:`external_phase_shifts`. 413 | 414 | Returns: 415 | External phase shift layers corresponding to :math:`\\boldsymbol{\\phi}` 416 | """ 417 | return self.phase_fn(self.external_phase_shifts) 418 | 419 | 420 | class Mesh: 421 | def __init__(self, model: MeshModel): 422 | """General mesh network layer defined by `neurophox.meshmodel.MeshModel` 423 | 424 | Args: 425 | model: The `MeshModel` model of the mesh network (e.g., rectangular, triangular, custom, etc.) 426 | """ 427 | self.model = model 428 | self.units, self.num_layers = self.model.units, self.model.num_layers 429 | self.pairwise_perm_idx = pairwise_off_diag_permutation(self.units) 430 | ss, cs, sc, cc = self.model.mzi_error_tensors 431 | self.ss, self.cs, self.sc, self.cc = tf.constant(ss, dtype=TF_COMPLEX), tf.constant(cs, dtype=TF_COMPLEX), \ 432 | tf.constant(sc, dtype=TF_COMPLEX), tf.constant(cc, dtype=TF_COMPLEX) 433 | self.perm_layers = [PermutationLayer(self.model.perm_idx[layer]) for layer in range(self.num_layers + 1)] 434 | 435 | def mesh_layers(self, phases: MeshPhasesTensorflow) -> List[MeshVerticalLayer]: 436 | """ 437 | 438 | Args: 439 | phases: The :code:`MeshPhasesTensorflow` object containing :math:`\\boldsymbol{\\theta}, \\boldsymbol{\\phi}, \\boldsymbol{\\gamma}` 440 | 441 | Returns: 442 | List of mesh layers to be used by any instance of :code:`MeshLayer` 443 | """ 444 | internal_psl = phases.internal_phase_shift_layers 445 | external_psl = phases.external_phase_shift_layers 446 | # smooth trick to efficiently perform the layerwise coupling computation 447 | 448 | if self.model.hadamard: 449 | s11 = self.cc * internal_psl + self.ss * _roll_tensor(internal_psl, up=True) 450 | s22 = _roll_tensor(self.ss * internal_psl + self.cc * _roll_tensor(internal_psl, up=True)) 451 | s12 = _roll_tensor(self.cs * internal_psl - self.sc * _roll_tensor(internal_psl, up=True)) 452 | s21 = self.sc * internal_psl - self.cs * _roll_tensor(internal_psl, up=True) 453 | else: 454 | s11 = self.cc * internal_psl - self.ss * _roll_tensor(internal_psl, up=True) 455 | s22 = _roll_tensor(-self.ss * internal_psl + self.cc * _roll_tensor(internal_psl, up=True)) 456 | s12 = 1j * _roll_tensor(self.cs * internal_psl + self.sc * _roll_tensor(internal_psl, up=True)) 457 | s21 = 1j * (self.sc * internal_psl + self.cs * _roll_tensor(internal_psl, up=True)) 458 | 459 | diag_layers = external_psl * (s11 + s22) / 2 460 | off_diag_layers = _roll_tensor(external_psl) * (s21 + s12) / 2 461 | 462 | if self.units % 2: 463 | diag_layers = tf.concat((diag_layers[:-1], tf.ones_like(diag_layers[-1:])), axis=0) 464 | 465 | diag_layers, off_diag_layers = tf.transpose(diag_layers), tf.transpose(off_diag_layers) 466 | 467 | mesh_layers = [MeshVerticalLayer(self.pairwise_perm_idx, diag_layers[0], off_diag_layers[0], 468 | self.perm_layers[1], self.perm_layers[0])] 469 | for layer in range(1, self.num_layers): 470 | mesh_layers.append(MeshVerticalLayer(self.pairwise_perm_idx, diag_layers[layer], off_diag_layers[layer], 471 | self.perm_layers[layer + 1])) 472 | return mesh_layers 473 | 474 | 475 | class MeshLayer(TransformerLayer): 476 | """Mesh network layer for unitary operators implemented in numpy 477 | 478 | Args: 479 | mesh_model: The `MeshModel` model of the mesh network (e.g., rectangular, triangular, custom, etc.) 480 | activation: Nonlinear activation function (:code:`None` if there's no nonlinearity) 481 | incoherent: Use an incoherent representation for the layer (no phase coherent between respective inputs...) 482 | phases: Initialize with phases (overrides mesh model initialization) 483 | """ 484 | 485 | def __init__(self, mesh_model: MeshModel, activation: Activation = None, 486 | incoherent: bool = False, 487 | phase_loss_fn: Optional[Callable[[tf.Tensor], tf.Tensor]] = None, 488 | **kwargs): 489 | self.mesh = Mesh(mesh_model) 490 | self.units, self.num_layers = self.mesh.units, self.mesh.num_layers 491 | self.incoherent = incoherent 492 | self.phase_loss_fn = phase_loss_fn 493 | super(MeshLayer, self).__init__(self.units, activation=activation, **kwargs) 494 | theta_init, phi_init, gamma_init = self.mesh.model.init 495 | self.theta, self.phi, self.gamma = theta_init.to_tf("theta"), phi_init.to_tf("phi"), gamma_init.to_tf("gamma") 496 | self.theta_fn, self.phi_fn, self.gamma_fn = self.mesh.model.theta_fn, self.mesh.model.phi_fn, self.mesh.model.gamma_fn 497 | 498 | @tf.function 499 | def transform(self, inputs: tf.Tensor) -> tf.Tensor: 500 | """ 501 | Performs the operation (where :math:`U` represents the matrix for this layer): 502 | 503 | .. math:: 504 | V_{\mathrm{out}} = V_{\mathrm{in}} U, 505 | 506 | where :math:`U \in \mathrm{U}(N)` and :math:`V_{\mathrm{out}}, V_{\mathrm{in}} \in \mathbb{C}^{M \\times N}`. 507 | 508 | Args: 509 | inputs: :code:`inputs` batch represented by the matrix :math:`V_{\mathrm{in}} \in \mathbb{C}^{M \\times N}` 510 | 511 | Returns: 512 | Transformed :code:`inputs`, :math:`V_{\mathrm{out}}` 513 | """ 514 | _inputs = np.eye(self.units, dtype=np.complex64) if self.incoherent else inputs 515 | mesh_phases, mesh_layers = self.phases_and_layers 516 | outputs = _inputs * mesh_phases.input_phase_shift_layer 517 | for layer in range(self.num_layers): 518 | outputs = mesh_layers[layer].transform(outputs) 519 | if self.incoherent: 520 | power_matrix = tf.math.real(outputs) ** 2 + tf.math.imag(outputs) ** 2 521 | power_inputs = tf.math.real(inputs) ** 2 + tf.math.imag(inputs) ** 2 522 | outputs = power_inputs @ power_matrix 523 | return tf.complex(tf.sqrt(outputs), tf.zeros_like(outputs)) 524 | return outputs 525 | 526 | @tf.function 527 | def inverse_transform(self, outputs: tf.Tensor) -> tf.Tensor: 528 | """ 529 | Performs the operation (where :math:`U` represents the matrix for this layer): 530 | 531 | .. math:: 532 | V_{\mathrm{in}} = V_{\mathrm{out}} U^\dagger, 533 | 534 | where :math:`U \in \mathrm{U}(N)` and :math:`V_{\mathrm{out}}, V_{\mathrm{in}} \in \mathbb{C}^{M \\times N}`. 535 | 536 | Args: 537 | outputs: :code:`outputs` batch represented by the matrix :math:`V_{\mathrm{out}} \in \mathbb{C}^{M \\times N}` 538 | 539 | Returns: 540 | Inverse transformed :code:`outputs`, :math:`V_{\mathrm{in}}` 541 | """ 542 | mesh_phases, mesh_layers = self.phases_and_layers 543 | inputs = outputs 544 | for layer in reversed(range(self.num_layers)): 545 | inputs = mesh_layers[layer].inverse_transform(inputs) 546 | inputs = inputs * tf.math.conj(mesh_phases.input_phase_shift_layer) 547 | return inputs 548 | 549 | @property 550 | def phases_and_layers(self) -> Tuple[MeshPhasesTensorflow, List[MeshVerticalLayer]]: 551 | """ 552 | 553 | Returns: 554 | Phases and layers for this mesh layer 555 | """ 556 | mesh_phases = MeshPhasesTensorflow( 557 | theta=self.theta, phi=self.phi, gamma=self.gamma, 558 | theta_fn=self.theta_fn, phi_fn=self.phi_fn, gamma_fn=self.gamma_fn, 559 | mask=self.mesh.model.mask, hadamard=self.mesh.model.hadamard, 560 | units=self.units, basis=self.mesh.model.basis, phase_loss_fn=self.phase_loss_fn, 561 | ) 562 | mesh_layers = self.mesh.mesh_layers(mesh_phases) 563 | return mesh_phases, mesh_layers 564 | 565 | @property 566 | def phases(self) -> MeshPhases: 567 | """ 568 | 569 | Returns: 570 | The :code:`MeshPhases` object for this layer 571 | """ 572 | return MeshPhases( 573 | theta=self.theta.numpy() * self.mesh.model.mask, 574 | phi=self.phi.numpy() * self.mesh.model.mask, 575 | mask=self.mesh.model.mask, 576 | gamma=self.gamma.numpy() 577 | ) 578 | 579 | 580 | def _roll_tensor(tensor: tf.Tensor, up=False): 581 | # a complex number-friendly roll that works on gpu 582 | if up: 583 | return tf.concat([tensor[1:], tensor[tf.newaxis, 0]], axis=0) 584 | return tf.concat([tensor[tf.newaxis, -1], tensor[:-1]], axis=0) 585 | 586 | 587 | def _to_complex(tensor: tf.Tensor): 588 | if isinstance(tensor, tf.Tensor) and tensor.dtype == tf.float32: 589 | return tf.complex(tensor, tf.zeros_like(tensor)) 590 | else: 591 | return tensor 592 | -------------------------------------------------------------------------------- /neurophox/tensorflow/layers.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, List, Dict, Union, Callable 2 | 3 | import tensorflow as tf 4 | from tensorflow.keras.layers import Activation 5 | import numpy as np 6 | 7 | from .generic import TransformerLayer, MeshLayer, CompoundTransformerLayer, PermutationLayer 8 | from ..meshmodel import RectangularMeshModel, TriangularMeshModel, PermutingRectangularMeshModel, ButterflyMeshModel 9 | from ..helpers import rectangular_permutation, butterfly_layer_permutation 10 | from ..config import DEFAULT_BASIS, TF_FLOAT, TF_COMPLEX 11 | 12 | 13 | class RM(MeshLayer): 14 | """Rectangular mesh network layer for unitary operators implemented in tensorflow 15 | 16 | Args: 17 | units: The dimension of the unitary matrix (:math:`N`) 18 | num_layers: The number of layers (:math:`L`) of the mesh 19 | hadamard: Hadamard convention for the beamsplitters 20 | basis: Phase basis to use 21 | bs_error: Beamsplitter split ratio error 22 | theta_init: Initializer for :code:`theta` (:math:`\\boldsymbol{\\theta}` or :math:`\\theta_{n\ell}`) 23 | phi_init: Initializer for :code:`phi` (:math:`\\boldsymbol{\\phi}` or :math:`\\phi_{n\ell}`) 24 | gamma_init: Initializer for :code:`gamma` (:math:`\\boldsymbol{\\gamma}` or :math:`\\gamma_{n}`) 25 | phase_loss_fn: Phase loss function for the layer 26 | activation: Nonlinear activation function (:code:`None` if there's no nonlinearity) 27 | """ 28 | 29 | def __init__(self, units: int, num_layers: int = None, hadamard: bool = False, incoherent: bool = False, 30 | basis: str = DEFAULT_BASIS, bs_error: float = 0.0, 31 | theta_init: Union[str, tuple, np.ndarray] = "haar_rect", 32 | phi_init: Union[str, tuple, np.ndarray] = "random_phi", 33 | gamma_init: Union[str, tuple, np.ndarray] = "random_gamma", 34 | phase_loss_fn: Optional[Callable[[tf.Tensor], tf.Tensor]] = None, 35 | activation: Activation = None, **kwargs): 36 | super(RM, self).__init__( 37 | RectangularMeshModel(units, num_layers, hadamard, bs_error, basis, theta_init, phi_init, gamma_init), 38 | activation=activation, incoherent=incoherent, 39 | phase_loss_fn=phase_loss_fn, **kwargs 40 | ) 41 | 42 | 43 | class TM(MeshLayer): 44 | """Triangular mesh network layer for unitary operators implemented in tensorflow 45 | 46 | Args: 47 | units: The dimension of the unitary matrix (:math:`N`) 48 | hadamard: Hadamard convention for the beamsplitters 49 | basis: Phase basis to use 50 | bs_error: Beamsplitter split ratio error 51 | theta_init: Initializer for :code:`theta` (:math:`\\boldsymbol{\\theta}` or :math:`\\theta_{n\ell}`) 52 | phi_init: Initializer for :code:`phi` (:math:`\\boldsymbol{\\phi}` or :math:`\\phi_{n\ell}`) 53 | gamma_init: Initializer for :code:`gamma` (:math:`\\boldsymbol{\\gamma}` or :math:`\\gamma_{n}`) 54 | phase_loss_fn: Phase loss function for the layer 55 | activation: Nonlinear activation function (:code:`None` if there's no nonlinearity) 56 | """ 57 | 58 | def __init__(self, units: int, hadamard: bool = False, incoherent: bool = False, basis: str = DEFAULT_BASIS, 59 | bs_error: float = 0.0, theta_init: Union[str, tuple, np.ndarray] = "haar_tri", 60 | phi_init: Union[str, tuple, np.ndarray] = "random_phi", 61 | gamma_init: Union[str, tuple, np.ndarray] = "random_gamma", 62 | phase_loss_fn: Optional[Callable[[tf.Tensor], tf.Tensor]] = None, 63 | activation: Activation = None, **kwargs): 64 | super(TM, self).__init__( 65 | TriangularMeshModel(units, hadamard, bs_error, basis, theta_init, phi_init, gamma_init), 66 | activation=activation, incoherent=incoherent, phase_loss_fn=phase_loss_fn, **kwargs 67 | ) 68 | 69 | 70 | class PRM(MeshLayer): 71 | """Permuting rectangular mesh unitary layer 72 | 73 | Args: 74 | units: The dimension of the unitary matrix (:math:`N`) to be modeled by this transformer 75 | tunable_layers_per_block: The number of tunable layers per block (overrides :code:`num_tunable_layers_list`, :code:`sampling_frequencies`) 76 | num_tunable_layers_list: Number of tunable layers in each block in order from left to right 77 | sampling_frequencies: Frequencies of sampling frequencies between the tunable layers 78 | is_trainable: Whether the parameters are trainable 79 | bs_error: Photonic error in the beamsplitter 80 | theta_init: Initializer for :code:`theta` (:math:`\\boldsymbol{\\theta}` or :math:`\\theta_{n\ell}`) 81 | phi_init: Initializer for :code:`phi` (:math:`\\boldsymbol{\\phi}` or :math:`\\phi_{n\ell}`) 82 | gamma_init: Initializer for :code:`gamma` (:math:`\\boldsymbol{\\gamma}` or :math:`\\gamma_{n}`) 83 | phase_loss_fn: Phase loss function for the layer 84 | activation: Nonlinear activation function (:code:`None` if there's no nonlinearity) 85 | """ 86 | 87 | def __init__(self, units: int, tunable_layers_per_block: int = None, 88 | num_tunable_layers_list: Optional[List[int]] = None, sampling_frequencies: Optional[List[int]] = None, 89 | bs_error: float = 0.0, hadamard: bool = False, incoherent: bool = False, 90 | theta_init: Union[str, tuple, np.ndarray] = "haar_prm", 91 | phi_init: Union[str, tuple, np.ndarray] = "random_phi", 92 | gamma_init: Union[str, tuple, np.ndarray] = "random_gamma", 93 | phase_loss_fn: Optional[Callable[[tf.Tensor], tf.Tensor]] = None, 94 | activation: Activation = None, **kwargs): 95 | if theta_init == 'haar_prm' and tunable_layers_per_block is not None: 96 | raise NotImplementedError('haar_prm initializer is incompatible with setting tunable_layers_per_block.') 97 | super(PRM, self).__init__( 98 | PermutingRectangularMeshModel(units, tunable_layers_per_block, num_tunable_layers_list, 99 | sampling_frequencies, bs_error, hadamard, theta_init, phi_init, gamma_init), 100 | activation=activation, incoherent=incoherent, phase_loss_fn=phase_loss_fn, **kwargs) 101 | 102 | 103 | class BM(MeshLayer): 104 | """Butterfly mesh unitary layer 105 | 106 | Args: 107 | units: The dimension of the unitary matrix (:math:`N`) 108 | hadamard: Hadamard convention for the beamsplitters 109 | basis: Phase basis to use 110 | bs_error: Beamsplitter split ratio error 111 | theta_init: Initializer for :code:`theta` (:math:`\\boldsymbol{\\theta}` or :math:`\\theta_{n\ell}`) 112 | phi_init: Initializer for :code:`phi` (:math:`\\boldsymbol{\\phi}` or :math:`\\phi_{n\ell}`) 113 | phase_loss_fn: Phase loss function for the layer 114 | activation: Nonlinear activation function (:code:`None` if there's no nonlinearity) 115 | """ 116 | 117 | def __init__(self, num_layers: int, hadamard: bool = False, incoherent: bool = False, basis: str = DEFAULT_BASIS, 118 | bs_error: float = 0.0, theta_init: Union[str, tuple, np.ndarray] = "random_theta", 119 | phi_init: Union[str, tuple, np.ndarray] = "random_phi", 120 | phase_loss_fn: Optional[Callable[[tf.Tensor], tf.Tensor]] = None, 121 | activation: Activation = None, **kwargs): 122 | super(BM, self).__init__( 123 | ButterflyMeshModel(num_layers, hadamard, bs_error, basis, theta_init, phi_init), 124 | activation=activation, incoherent=incoherent, phase_loss_fn=phase_loss_fn, **kwargs 125 | ) 126 | 127 | 128 | class SVD(CompoundTransformerLayer): 129 | """Singular value decomposition transformer for implementing a matrix. 130 | 131 | Notes: 132 | SVD requires you specify the unitary transformers used to implement the SVD in `unitary_transformer_dict`, 133 | specifying transformer name and arguments for that transformer. 134 | 135 | Args: 136 | units: The number of inputs (:math:`M`) of the :math:`M \\times N` matrix to be modelled by the SVD 137 | mesh_dict: The name and properties of the mesh layer used for the SVD 138 | output_units: The dimension of the output (:math:`N`) of the :math:`M \\times N` matrix to be modelled by the SVD 139 | pos_singular_values: Whether to allow only positive singular values 140 | activation: Nonlinear activation function (:code:`None` if there's no nonlinearity) 141 | """ 142 | 143 | def __init__(self, units: int, mesh_dict: Dict, output_units: Optional[int] = None, 144 | pos_singular_values: bool = False, activation: Activation = None): 145 | self.units = units 146 | self.output_units = output_units if output_units is not None else units 147 | if output_units != units and output_units is not None: 148 | raise NotImplementedError("Still working out a clean implementation of non-square linear operators.") 149 | self.mesh_name = mesh_dict['name'] 150 | self.mesh_properties = mesh_dict.get('properties', {}) 151 | self.pos = pos_singular_values 152 | 153 | mesh_name2layer = { 154 | 'rm': RM, 155 | 'prm': PRM, 156 | 'tm': TM 157 | } 158 | 159 | self.v = mesh_name2layer[self.mesh_name](units=units, name="v", **self.mesh_properties) 160 | self.diag = Diagonal(units, output_units=output_units, pos=self.pos) 161 | self.u = mesh_name2layer[self.mesh_name](units=units, name="u", **self.mesh_properties) 162 | 163 | self.activation = activation 164 | 165 | super(SVD, self).__init__( 166 | units=self.units, 167 | transformer_list=[self.v, self.diag, self.u] 168 | ) 169 | 170 | 171 | class DiagonalPhaseLayer(TransformerLayer): 172 | """Diagonal matrix of phase shifts 173 | 174 | Args: 175 | units: Dimension of the input (number of input waveguide ports), :math:`N` 176 | """ 177 | 178 | def __init__(self, units: int, **kwargs): 179 | super(DiagonalPhaseLayer, self).__init__(units=units) 180 | self.gamma = tf.Variable( 181 | name="gamma", 182 | initial_value=tf.constant(2 * np.pi * np.random.rand(units), dtype=TF_FLOAT), 183 | dtype=TF_FLOAT, 184 | **kwargs 185 | ) 186 | self.diag_vec = tf.complex(tf.cos(self.gamma), tf.sin(self.gamma)) 187 | self.inv_diag_vec = tf.complex(tf.cos(-self.gamma), tf.sin(-self.gamma)) 188 | self.variables.append(self.gamma) 189 | 190 | @tf.function 191 | def transform(self, inputs: tf.Tensor): 192 | return self.diag_vec * inputs 193 | 194 | @tf.function 195 | def inverse_transform(self, outputs: tf.Tensor): 196 | return self.inv_diag_vec * outputs 197 | 198 | 199 | class Diagonal(TransformerLayer): 200 | """Diagonal matrix of gains and losses (not necessarily real) 201 | 202 | Args: 203 | units: Dimension of the input (number of input waveguide ports), :math:`N` 204 | is_complex: Whether to use complex values or not 205 | output_units: Dimension of the output (number of output waveguide ports), :math:`M`. 206 | If :math:`M < N`, remove last :math:`N - M` elements. 207 | If :math:`M > N`, pad with :math:`M - N` zeros. 208 | pos: Enforce positive definite matrix (only positive singular values) 209 | 210 | """ 211 | 212 | def __init__(self, units: int, is_complex: bool = True, output_units: Optional[int] = None, 213 | pos: bool = False, **kwargs): 214 | super(Diagonal, self).__init__(units=units, **kwargs) 215 | self.output_dim = output_units if output_units is not None else units 216 | self.pos = pos 217 | self.is_complex = is_complex 218 | singular_value_dim = min(self.units, self.output_dim) 219 | self.sigma = tf.Variable( 220 | name="sigma", 221 | initial_value=tf.constant(2 * np.pi * np.random.randn(singular_value_dim), dtype=TF_FLOAT), 222 | dtype=TF_FLOAT 223 | ) 224 | 225 | @tf.function 226 | def transform(self, inputs: tf.Tensor) -> tf.Tensor: 227 | sigma = tf.abs(self.sigma) if self.pos else self.sigma 228 | diag_vec = tf.cast(sigma, TF_COMPLEX) if self.is_complex else sigma 229 | if self.output_dim == self.units: 230 | return diag_vec * inputs 231 | elif self.output_dim < self.units: 232 | return diag_vec * inputs[:self.output_dim] 233 | else: 234 | return tf.pad(diag_vec * inputs, tf.constant([[0, 0], [0, self.output_dim - self.units]])) 235 | 236 | @tf.function 237 | def inverse_transform(self, outputs: tf.Tensor) -> tf.Tensor: 238 | sigma = tf.abs(self.sigma) if self.pos else self.sigma 239 | inv_diag_vec = tf.cast(1 / sigma, TF_COMPLEX) if self.is_complex else 1 / sigma 240 | if self.output_dim == self.units: 241 | return inv_diag_vec * outputs 242 | elif self.output_dim > self.units: 243 | return inv_diag_vec * outputs[:self.units] 244 | else: 245 | return tf.pad(inv_diag_vec * outputs, tf.constant([[0, 0], [0, self.units - self.output_dim]])) 246 | 247 | 248 | class RectangularPerm(PermutationLayer): 249 | """Rectangular permutation layer 250 | 251 | The rectangular permutation layer for a frequency :math:`f` corresponds effectively is equivalent to adding 252 | :math:`f` layers of cross state MZIs in a grid configuration to the existing mesh. 253 | 254 | Args: 255 | units: Dimension of the input (number of input waveguide ports), :math:`N` 256 | frequency: Frequency of interacting mesh wires (waveguides) 257 | """ 258 | 259 | def __init__(self, units: int, frequency: int): 260 | self.frequency = frequency 261 | super(RectangularPerm, self).__init__( 262 | permuted_indices=rectangular_permutation(units, frequency)) 263 | 264 | 265 | class ButterflyPerm(PermutationLayer): 266 | """Butterfly (FFT) permutation layer 267 | 268 | The butterfly or FFT permutation for a frequency :math:`f` corresponds to switching all inputs 269 | that are :math:`f` inputs apart. This works most cleanly in a butterfly mesh architecture where 270 | the number of inputs, :math:`N`, and the frequencies, :math:`f` are powers of two. 271 | 272 | Args: 273 | units: Dimension of the input (number of input waveguide ports), :math:`N` 274 | frequency: Frequency of interacting mesh wires (waveguides) 275 | """ 276 | 277 | def __init__(self, units: int, frequency: int): 278 | self.frequency = frequency 279 | super(ButterflyPerm, self).__init__(permuted_indices=butterfly_layer_permutation(units, frequency)) 280 | -------------------------------------------------------------------------------- /neurophox/torch/__init__.py: -------------------------------------------------------------------------------- 1 | from .generic import MeshTorchLayer, MeshVerticalLayer 2 | from .layers import * 3 | -------------------------------------------------------------------------------- /neurophox/torch/generic.py: -------------------------------------------------------------------------------- 1 | from typing import List, Optional, Callable 2 | 3 | import torch 4 | from torch.nn import Module, Parameter 5 | import numpy as np 6 | 7 | from ..numpy.generic import MeshPhases 8 | from ..config import BLOCH, SINGLEMODE 9 | from ..meshmodel import MeshModel 10 | from ..helpers import pairwise_off_diag_permutation, plot_complex_matrix 11 | 12 | 13 | class TransformerLayer(Module): 14 | """Base transformer class for transformer layers (invertible functions, usually linear) 15 | 16 | Args: 17 | units: Dimension of the input to be transformed by the transformer 18 | activation: Nonlinear activation function (:code:`None` if there's no nonlinearity) 19 | """ 20 | 21 | def __init__(self, units: int, is_trainable: bool = False): 22 | super(TransformerLayer, self).__init__() 23 | self.units = units 24 | self.is_trainable = is_trainable 25 | 26 | def transform(self, inputs: torch.Tensor) -> torch.Tensor: 27 | return inputs 28 | 29 | def inverse_transform(self, outputs: torch.Tensor) -> torch.Tensor: 30 | return outputs 31 | 32 | def matrix(self) -> np.ndarray: 33 | torch_matrix = self.transform(np.eye(self.units, dtype=np.complex64)).cpu().detach().numpy() 34 | return torch_matrix 35 | 36 | def inverse_matrix(self) -> np.ndarray: 37 | torch_matrix = self.inverse_transform(np.eye(self.units, dtype=np.complex64)).cpu().detach().numpy() 38 | return torch_matrix 39 | 40 | def plot(self, plt): 41 | plot_complex_matrix(plt, self.matrix()) 42 | 43 | def forward(self, x): 44 | return self.transform(x) 45 | 46 | 47 | class CompoundTransformerLayer(TransformerLayer): 48 | def __init__(self, units: int, transformer_list: List[TransformerLayer], is_trainable: bool = False): 49 | self.transformer_list = transformer_list 50 | super(CompoundTransformerLayer, self).__init__(units=units, is_trainable=is_trainable) 51 | 52 | def transform(self, inputs: torch.Tensor) -> torch.Tensor: 53 | outputs = inputs 54 | for transformer in self.transformer_list: 55 | outputs = transformer.transform(outputs) 56 | return outputs 57 | 58 | def inverse_transform(self, outputs: torch.Tensor) -> torch.Tensor: 59 | inputs = outputs 60 | for transformer in self.transformer_list[::-1]: 61 | inputs = transformer.inverse_transform(inputs) 62 | return inputs 63 | 64 | 65 | class PermutationLayer(TransformerLayer): 66 | def __init__(self, permuted_indices: np.ndarray): 67 | super(PermutationLayer, self).__init__(units=permuted_indices.shape[0]) 68 | self.units = permuted_indices.shape[0] 69 | self.permuted_indices = np.asarray(permuted_indices, dtype=np.long) 70 | self.inv_permuted_indices = np.zeros_like(self.permuted_indices) 71 | for idx, perm_idx in enumerate(self.permuted_indices): 72 | self.inv_permuted_indices[perm_idx] = idx 73 | 74 | def transform(self, inputs: torch.Tensor): 75 | return inputs[..., self.permuted_indices] 76 | 77 | def inverse_transform(self, outputs: torch.Tensor): 78 | return outputs[..., self.inv_permuted_indices] 79 | 80 | 81 | class MeshVerticalLayer(TransformerLayer): 82 | """ 83 | Args: 84 | diag: the diagonal terms to multiply 85 | off_diag: the off-diagonal terms to multiply 86 | left_perm: the permutation for the mesh vertical layer (prior to the coupling operation) 87 | right_perm: the right permutation for the mesh vertical layer 88 | (usually for the final layer and after the coupling operation) 89 | """ 90 | 91 | def __init__(self, pairwise_perm_idx: np.ndarray, diag: torch.Tensor, off_diag: torch.Tensor, 92 | right_perm: PermutationLayer = None, left_perm: PermutationLayer = None): 93 | self.diag = diag 94 | self.off_diag = off_diag 95 | self.pairwise_perm_idx = pairwise_perm_idx 96 | super(MeshVerticalLayer, self).__init__(pairwise_perm_idx.shape[0]) 97 | self.left_perm = left_perm 98 | self.right_perm = right_perm 99 | 100 | def transform(self, inputs: torch.Tensor) -> torch.Tensor: 101 | """ 102 | Propagate :code:`inputs` through single layer :math:`\ell < L` 103 | (where :math:`U_\ell` represents the matrix for layer :math:`\ell`): 104 | 105 | .. math:: 106 | V_{\mathrm{out}} = V_{\mathrm{in}} U^{(\ell')}, 107 | 108 | Args: 109 | inputs: :code:`inputs` batch represented by the matrix :math:`V_{\mathrm{in}} \in \mathbb{C}^{M \\times N}` 110 | 111 | Returns: 112 | Propaged :code:`inputs` through single layer :math:`\ell` to form an array 113 | :math:`V_{\mathrm{out}} \in \mathbb{C}^{M \\times N}`. 114 | """ 115 | if isinstance(inputs, np.ndarray): 116 | inputs = torch.tensor(inputs, dtype=torch.cfloat, device=self.diag.device) 117 | outputs = inputs if self.left_perm is None else self.left_perm(inputs) 118 | diag_out = outputs * self.diag 119 | off_diag_out = outputs * self.off_diag 120 | outputs = diag_out + off_diag_out[..., self.pairwise_perm_idx] 121 | outputs = outputs if self.right_perm is None else self.right_perm(outputs) 122 | return outputs 123 | 124 | def inverse_transform(self, outputs: torch.Tensor) -> torch.Tensor: 125 | """ 126 | Inverse-propagate :code:`inputs` through single layer :math:`\ell < L` 127 | (where :math:`U_\ell` represents the matrix for layer :math:`\ell`): 128 | 129 | .. math:: 130 | V_{\mathrm{in}} = V_{\mathrm{out}} (U^{(\ell')})^\dagger, 131 | 132 | Args: 133 | outputs: :code:`outputs` batch represented by the matrix :math:`V_{\mathrm{out}} \in \mathbb{C}^{M \\times N}` 134 | 135 | Returns: 136 | Inverse propaged :code:`outputs` through single layer :math:`\ell` to form an array 137 | :math:`V_{\mathrm{in}} \in \mathbb{C}^{M \\times N}`. 138 | """ 139 | if isinstance(outputs, np.ndarray): 140 | outputs = torch.tensor(outputs, dtype=torch.cfloat, device=self.diag.device) 141 | inputs = outputs if self.right_perm is None else self.right_perm.inverse_transform(outputs) 142 | diag = self.diag.conj() 143 | off_diag = self.off_diag[..., self.pairwise_perm_idx].conj() 144 | inputs = inputs * diag + (inputs * off_diag)[..., self.pairwise_perm_idx] 145 | inputs = inputs if self.left_perm is None else self.left_perm.inverse_transform(inputs) 146 | return inputs 147 | 148 | 149 | class MeshParamTorch: 150 | """A class that cleanly arranges parameters into a specific arrangement that can be used to simulate any mesh 151 | 152 | Args: 153 | param: parameter to arrange in mesh 154 | units: number of inputs/outputs of the mesh 155 | """ 156 | 157 | def __init__(self, param: torch.Tensor, units: int): 158 | self.param = param 159 | self.units = units 160 | 161 | @property 162 | def single_mode_arrangement(self) -> torch.Tensor: 163 | """ 164 | The single-mode arrangement based on the :math:`L(\\theta)` transfer matrix for :code:`PhaseShiftUpper` 165 | is one where elements of `param` are on the even rows and all odd rows are zero. 166 | 167 | In particular, given the :code:`param` array 168 | :math:`\\boldsymbol{\\theta} = [\\boldsymbol{\\theta}_1, \\boldsymbol{\\theta}_2, \ldots \\boldsymbol{\\theta}_M]^T`, 169 | where :math:`\\boldsymbol{\\theta}_m` represent row vectors and :math:`M = \\lfloor\\frac{N}{2}\\rfloor`, the single-mode arrangement has the stripe array form 170 | :math:`\widetilde{\\boldsymbol{\\theta}} = [\\boldsymbol{\\theta}_1, \\boldsymbol{0}, \\boldsymbol{\\theta}_2, \\boldsymbol{0}, \ldots \\boldsymbol{\\theta}_N, \\boldsymbol{0}]^T`. 171 | where :math:`\widetilde{\\boldsymbol{\\theta}} \in \mathbb{R}^{N \\times L}` defines the :math:`\\boldsymbol{\\theta}` of the final mesh 172 | and :math:`\\boldsymbol{0}` represents an array of zeros of the same size as :math:`\\boldsymbol{\\theta}_n`. 173 | 174 | Returns: 175 | Single-mode arrangement array of phases 176 | 177 | """ 178 | num_layers = self.param.shape[0] 179 | tensor_t = self.param.t() 180 | stripe_tensor = torch.zeros(self.units, num_layers, dtype=torch.float32, device=self.param.device) 181 | if self.units % 2: 182 | stripe_tensor[:-1][::2] = tensor_t 183 | else: 184 | stripe_tensor[::2] = tensor_t 185 | return stripe_tensor 186 | 187 | @property 188 | def common_mode_arrangement(self) -> torch.Tensor: 189 | """ 190 | The common-mode arrangement based on the :math:`C(\\theta)` transfer matrix for :code:`PhaseShiftCommonMode` 191 | is one where elements of `param` are on the even rows and repeated on respective odd rows. 192 | 193 | In particular, given the :code:`param` array 194 | :math:`\\boldsymbol{\\theta} = [\\boldsymbol{\\theta}_1, \\boldsymbol{\\theta}_2, \ldots \\boldsymbol{\\theta}_M]^T`, 195 | where :math:`\\boldsymbol{\\theta}_n` represent row vectors and :math:`M = \\lfloor\\frac{N}{2}\\rfloor`, the common-mode arrangement has the stripe array form 196 | :math:`\\widetilde{\\boldsymbol{\\theta}} = [\\boldsymbol{\\theta}_1, \\boldsymbol{\\theta}_1, \\boldsymbol{\\theta}_2, \\boldsymbol{\\theta}_2, \ldots \\boldsymbol{\\theta}_N, \\boldsymbol{\\theta}_N]^T`. 197 | where :math:`\widetilde{\\boldsymbol{\\theta}} \in \mathbb{R}^{N \\times L}` defines the :math:`\\boldsymbol{\\theta}` of the final mesh. 198 | 199 | Returns: 200 | Common-mode arrangement array of phases 201 | """ 202 | phases = self.single_mode_arrangement 203 | return phases + phases.roll(1, 0) 204 | 205 | @property 206 | def differential_mode_arrangement(self) -> torch.Tensor: 207 | """ 208 | The differential-mode arrangement is based on the :math:`D(\\theta)` transfer matrix 209 | for :code:`PhaseShiftDifferentialMode`. 210 | 211 | Given the :code:`param` array 212 | :math:`\\boldsymbol{\\theta} = [\cdots \\boldsymbol{\\theta}_m \cdots]^T`, 213 | where :math:`\\boldsymbol{\\theta}_n` represent row vectors and :math:`M = \\lfloor\\frac{N}{2}\\rfloor`, the differential-mode arrangement has the form 214 | :math:`\\widetilde{\\boldsymbol{\\theta}} = \\left[\cdots \\frac{\\boldsymbol{\\theta}_m}{2}, -\\frac{\\boldsymbol{\\theta}_m}{2} \cdots \\right]^T`. 215 | where :math:`\widetilde{\\boldsymbol{\\theta}} \in \mathbb{R}^{N \\times L}` defines the :math:`\\boldsymbol{\\theta}` of the final mesh. 216 | 217 | Returns: 218 | Differential-mode arrangement array of phases 219 | 220 | """ 221 | phases = self.single_mode_arrangement 222 | return phases / 2 - phases.roll(1, 0) / 2 223 | 224 | def __add__(self, other): 225 | return MeshParamTorch(self.param + other.param, self.units) 226 | 227 | def __sub__(self, other): 228 | return MeshParamTorch(self.param - other.param, self.units) 229 | 230 | def __mul__(self, other): 231 | return MeshParamTorch(self.param * other.param, self.units) 232 | 233 | 234 | class MeshPhasesTorch: 235 | """Organizes the phases in the mesh into appropriate arrangements 236 | 237 | Args: 238 | theta: Array to be converted to :math:`\\boldsymbol{\\theta}` 239 | phi: Array to be converted to :math:`\\boldsymbol{\\phi}` 240 | gamma: Array to be converted to :math:`\\boldsymbol{\gamma}` 241 | mask: Mask over values of :code:`theta` and :code:`phi` that are not in bar state 242 | basis: Phase basis to use 243 | hadamard: Whether to use Hadamard convention 244 | theta_fn: Pytorch-friendly phi function call to reparametrize phi (example use cases: 245 | can use a mask to keep some values of theta fixed or always bound theta between 0 and pi). 246 | By default, use identity function. 247 | phi_fn: Pytorch-friendly phi function call to reparametrize phi (example use cases: 248 | can use a mask to keep some values of phi fixed or always bound phi between 0 and 2 * pi). 249 | By default, use identity function. 250 | gamma_fn: Pytorch-friendly gamma function call to reparametrize gamma (example use cases: 251 | can use a mask to keep some values of gamma fixed or always bound gamma between 0 and 2 * pi). 252 | By default, use identity function. 253 | phase_loss_fn: Incorporate phase shift-dependent loss into the model. 254 | The function is of the form phase_loss_fn(phases), 255 | which returns the loss 256 | """ 257 | 258 | def __init__(self, theta: Parameter, phi: Parameter, mask: np.ndarray, gamma: Parameter, units: int, 259 | basis: str = SINGLEMODE, hadamard: bool = False, theta_fn: Optional[Callable] = None, 260 | phi_fn: Optional[Callable] = None, gamma_fn: Optional[Callable] = None, 261 | phase_loss_fn: Optional[Callable[[torch.Tensor], torch.Tensor]] = None): 262 | self.mask = mask if mask is not None else np.ones_like(theta) 263 | mask = torch.as_tensor(mask, dtype=theta.dtype, device=theta.device) 264 | self.theta_fn = (lambda x: x) if theta_fn is None else theta_fn 265 | self.phi_fn = (lambda x: x) if phi_fn is None else phi_fn 266 | self.gamma_fn = (lambda x: x) if gamma_fn is None else gamma_fn 267 | self.theta = MeshParamTorch(self.theta_fn(theta) * mask + (1 - mask) * (1 - hadamard) * np.pi, units=units) 268 | self.phi = MeshParamTorch(self.phi_fn(phi) * mask + (1 - mask) * (1 - hadamard) * np.pi, units=units) 269 | self.gamma = self.gamma_fn(gamma) 270 | self.basis = basis 271 | self.phase_fn = lambda phase: torch.as_tensor(1 - phase_loss_fn(phase)) * phasor(phase) \ 272 | if phase_loss_fn is not None else phasor(phase) 273 | self.input_phase_shift_layer = self.phase_fn(gamma) 274 | if self.theta.param.shape != self.phi.param.shape: 275 | raise ValueError("Internal phases (theta) and external phases (phi) need to have the same shape.") 276 | 277 | @property 278 | def internal_phase_shifts(self): 279 | """The internal phase shift matrix of the mesh corresponds to an `L \\times N` array of phase shifts 280 | (in between beamsplitters, thus internal) where :math:`L` is number of layers 281 | and :math:`N` is number of inputs/outputs. 282 | 283 | Returns: 284 | Internal phase shift matrix corresponding to :math:`\\boldsymbol{\\theta}` 285 | """ 286 | if self.basis == BLOCH: 287 | return self.theta.differential_mode_arrangement 288 | elif self.basis == SINGLEMODE: 289 | return self.theta.single_mode_arrangement 290 | else: 291 | raise NotImplementedError(f"{self.basis} is not yet supported or invalid.") 292 | 293 | @property 294 | def external_phase_shifts(self): 295 | """The external phase shift matrix of the mesh corresponds to an `L \\times N` array of phase shifts 296 | (outside of beamsplitters, thus external) where :math:`L` is number of layers and 297 | :math:`N` is number of inputs/outputs. 298 | 299 | Returns: 300 | External phase shift matrix corresponding to :math:`\\boldsymbol{\\phi}` 301 | """ 302 | if self.basis == BLOCH or self.basis == SINGLEMODE: 303 | return self.phi.single_mode_arrangement 304 | else: 305 | raise NotImplementedError(f"{self.basis} is not yet supported or invalid.") 306 | 307 | @property 308 | def internal_phase_shift_layers(self): 309 | """Elementwise applying complex exponential to :code:`internal_phase_shifts`. 310 | 311 | Returns: 312 | Internal phase shift layers corresponding to :math:`\\boldsymbol{\\theta}` 313 | """ 314 | return self.phase_fn(self.internal_phase_shifts) 315 | 316 | @property 317 | def external_phase_shift_layers(self): 318 | """Elementwise applying complex exponential to :code:`external_phase_shifts`. 319 | 320 | Returns: 321 | External phase shift layers corresponding to :math:`\\boldsymbol{\\phi}` 322 | """ 323 | return self.phase_fn(self.external_phase_shifts) 324 | 325 | 326 | class MeshTorchLayer(TransformerLayer): 327 | """Mesh network layer for unitary operators implemented in numpy 328 | 329 | Args: 330 | mesh_model: The model of the mesh network (e.g., rectangular, triangular, butterfly) 331 | """ 332 | 333 | def __init__(self, mesh_model: MeshModel): 334 | super(MeshTorchLayer, self).__init__(mesh_model.units) 335 | self.mesh_model = mesh_model 336 | ss, cs, sc, cc = self.mesh_model.mzi_error_tensors 337 | ss, cs, sc, cc = torch.as_tensor(ss, dtype=torch.float32), torch.as_tensor(cs, dtype=torch.float32), \ 338 | torch.as_tensor(sc, dtype=torch.float32), torch.as_tensor(cc, dtype=torch.float32) 339 | self.register_buffer("ss", ss) 340 | self.register_buffer("cs", cs) 341 | self.register_buffer("sc", sc) 342 | self.register_buffer("cc", cc) 343 | theta_init, phi_init, gamma_init = self.mesh_model.init 344 | self.units, self.num_layers = self.mesh_model.units, self.mesh_model.num_layers 345 | self.theta, self.phi, self.gamma = theta_init.to_torch(), phi_init.to_torch(), gamma_init.to_torch() 346 | self.theta_fn, self.phi_fn, self.gamma_fn = self.mesh_model.theta_fn, self.mesh_model.phi_fn,\ 347 | self.mesh_model.gamma_fn 348 | self.pairwise_perm_idx = pairwise_off_diag_permutation(self.units) 349 | self.perm_layers = [PermutationLayer(self.mesh_model.perm_idx[layer]) for layer in range(self.num_layers + 1)] 350 | 351 | def transform(self, inputs: torch.Tensor) -> torch.Tensor: 352 | """ 353 | Performs the operation (where :math:`U` represents the matrix for this layer): 354 | 355 | .. math:: 356 | V_{\mathrm{out}} = V_{\mathrm{in}} U, 357 | 358 | where :math:`U \in \mathrm{U}(N)` and :math:`V_{\mathrm{out}}, V_{\mathrm{in}} \in \mathbb{C}^{M \\times N}`. 359 | 360 | Args: 361 | inputs: :code:`inputs` batch represented by the matrix :math:`V_{\mathrm{in}} \in \mathbb{C}^{M \\times N}` 362 | 363 | Returns: 364 | Transformed :code:`inputs`, :math:`V_{\mathrm{out}}` 365 | """ 366 | mesh_phases = MeshPhasesTorch( 367 | theta=self.theta, phi=self.phi, gamma=self.gamma, 368 | theta_fn=self.theta_fn, phi_fn=self.phi_fn, gamma_fn=self.gamma_fn, 369 | mask=self.mesh_model.mask, hadamard=self.mesh_model.hadamard, 370 | units=self.units, basis=self.mesh_model.basis 371 | ) 372 | mesh_layers = self.mesh_layers(mesh_phases) 373 | if isinstance(inputs, np.ndarray): 374 | inputs = torch.tensor(inputs, dtype=torch.cfloat, device=self.theta.device) 375 | outputs = inputs * mesh_phases.input_phase_shift_layer 376 | for mesh_layer in mesh_layers: 377 | outputs = mesh_layer(outputs) 378 | return outputs 379 | 380 | def inverse_transform(self, outputs: torch.Tensor) -> torch.Tensor: 381 | """ 382 | Performs the operation (where :math:`U` represents the matrix for this layer): 383 | 384 | .. math:: 385 | V_{\mathrm{in}} = V_{\mathrm{out}} U^\dagger, 386 | 387 | where :math:`U \in \mathrm{U}(N)` and :math:`V_{\mathrm{out}}, V_{\mathrm{in}} \in \mathbb{C}^{M \\times N}`. 388 | 389 | Args: 390 | outputs: :code:`outputs` batch represented by the matrix :math:`V_{\mathrm{out}} \in \mathbb{C}^{M \\times N}` 391 | 392 | Returns: 393 | Inverse transformed :code:`outputs`, :math:`V_{\mathrm{in}}` 394 | """ 395 | mesh_phases = MeshPhasesTorch( 396 | theta=self.theta, phi=self.phi, gamma=self.gamma, 397 | mask=self.mesh_model.mask, hadamard=self.mesh_model.hadamard, 398 | units=self.units, basis=self.mesh_model.basis 399 | ) 400 | mesh_layers = self.mesh_layers(mesh_phases) 401 | inputs = torch.tensor(outputs, dtype=torch.cfloat, device=self.theta.device) if isinstance(outputs, np.ndarray) else outputs 402 | for layer in reversed(range(self.num_layers)): 403 | inputs = mesh_layers[layer].inverse_transform(inputs) 404 | inputs = inputs * mesh_phases.input_phase_shift_layer.conj() 405 | return inputs 406 | 407 | def adjoint_transform(self, inputs: torch.Tensor) -> torch.Tensor: 408 | return self.inverse_transform(inputs) 409 | 410 | @property 411 | def phases(self) -> MeshPhases: 412 | """ 413 | 414 | Returns: 415 | The :code:`MeshPhases` object for this layer 416 | """ 417 | return MeshPhases( 418 | theta=self.theta.detach().numpy() * self.mesh_model.mask, 419 | phi=self.phi.detach().numpy() * self.mesh_model.mask, 420 | mask=self.mesh_model.mask, 421 | gamma=self.gamma.detach().numpy() 422 | ) 423 | 424 | def mesh_layers(self, phases: MeshPhasesTorch) -> List[MeshVerticalLayer]: 425 | """ 426 | 427 | Args: 428 | phases: The :code:`MeshPhasesTensorflow` object containing :math:`\\boldsymbol{\\theta}, \\boldsymbol{\\phi}, \\boldsymbol{\\gamma}` 429 | 430 | Returns: 431 | List of mesh layers to be used by any instance of :code:`MeshLayer` 432 | """ 433 | internal_psl = phases.internal_phase_shift_layers 434 | external_psl = phases.external_phase_shift_layers 435 | 436 | # smooth trick to efficiently perform the layerwise coupling computation 437 | 438 | if self.mesh_model.hadamard: 439 | s11 = self.cc * internal_psl + self.ss * internal_psl.roll(-1, 0) 440 | s22 = (self.ss * internal_psl + self.cc * internal_psl.roll(-1, 0)).roll(1, 0) 441 | s12 = (self.cs * internal_psl - self.sc * internal_psl.roll(-1, 0)).roll(1, 0) 442 | s21 = self.sc * internal_psl - self.cs * internal_psl.roll(-1, 0) 443 | else: 444 | s11 = self.cc * internal_psl - self.ss * internal_psl.roll(-1, 0) 445 | s22 = (-self.ss * internal_psl + self.cc * internal_psl.roll(-1, 0)).roll(1, 0) 446 | s12 = 1j * (self.cs * internal_psl + self.sc * internal_psl.roll(-1, 0)).roll(1, 0) 447 | s21 = 1j * (self.sc * internal_psl + self.cs * internal_psl.roll(-1, 0)) 448 | 449 | diag_layers = external_psl * (s11 + s22) / 2 450 | off_diag_layers = external_psl.roll(1, 0) * (s21 + s12) / 2 451 | 452 | if self.units % 2: 453 | diag_layers = torch.cat((diag_layers[:-1], torch.ones_like(diag_layers[-1:])), dim=0) 454 | 455 | diag_layers, off_diag_layers = diag_layers.t(), off_diag_layers.t() 456 | 457 | mesh_layers = [MeshVerticalLayer( 458 | self.pairwise_perm_idx, diag_layers[0], off_diag_layers[0], self.perm_layers[1], self.perm_layers[0])] 459 | for layer in range(1, self.num_layers): 460 | mesh_layers.append(MeshVerticalLayer(self.pairwise_perm_idx, diag_layers[layer], 461 | off_diag_layers[layer], self.perm_layers[layer + 1])) 462 | 463 | return mesh_layers 464 | 465 | 466 | def phasor(phase: torch.Tensor): 467 | return phase.cos() + 1j * phase.sin() 468 | -------------------------------------------------------------------------------- /neurophox/torch/layers.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, List, Union 2 | 3 | from .generic import MeshTorchLayer, PermutationLayer 4 | from ..meshmodel import RectangularMeshModel, TriangularMeshModel, PermutingRectangularMeshModel, ButterflyMeshModel 5 | from ..helpers import rectangular_permutation, butterfly_layer_permutation 6 | from ..config import DEFAULT_BASIS 7 | 8 | import numpy as np 9 | 10 | 11 | class RMTorch(MeshTorchLayer): 12 | """Rectangular mesh network layer for unitary operators implemented in tensorflow 13 | 14 | Args: 15 | units: The dimension of the unitary matrix (:math:`N`) 16 | num_layers: The number of layers (:math:`L`) of the mesh 17 | hadamard: Hadamard convention for the beamsplitters 18 | basis: Phase basis to use 19 | bs_error: Beamsplitter split ratio error 20 | theta_init: Initializer for :code:`theta` (:math:`\\boldsymbol{\\theta}` or :math:`\\theta_{n\ell}`) 21 | a :code:`str`, :code:`ndarray`, or tuple of the form :code:`(theta_init, theta_fn)`. 22 | phi_init: Initializer for :code:`phi` (:math:`\\boldsymbol{\\phi}` or :math:`\\phi_{n\ell}`): 23 | a :code:`str`, :code:`ndarray`, or tuple of the form :code:`(phi_init, phi_fn)`. 24 | gamma_init: Initializer for :code:`gamma` (:math:`\\boldsymbol{\\gamma}` or :math:`\\gamma_{n}`): 25 | a :code:`str`, :code:`ndarray`, or tuple of the form :code:`(gamma_init, gamma_fn)`. 26 | """ 27 | 28 | def __init__(self, units: int, num_layers: int = None, hadamard: bool = False, basis: str = DEFAULT_BASIS, 29 | bs_error: float = 0.0, theta_init: Union[str, tuple, np.ndarray] = "haar_rect", 30 | phi_init: Union[str, tuple, np.ndarray] = "random_phi", 31 | gamma_init: Union[str, tuple, np.ndarray] = "random_gamma"): 32 | super(RMTorch, self).__init__( 33 | RectangularMeshModel(units, num_layers, hadamard, bs_error, basis, 34 | theta_init, phi_init, gamma_init)) 35 | 36 | 37 | class TMTorch(MeshTorchLayer): 38 | """Triangular mesh network layer for unitary operators implemented in tensorflow 39 | 40 | Args: 41 | units: The dimension of the unitary matrix (:math:`N`) 42 | hadamard: Hadamard convention for the beamsplitters 43 | basis: Phase basis to use 44 | bs_error: Beamsplitter split ratio error 45 | theta_init: Initializer for :code:`theta` (:math:`\\boldsymbol{\\theta}` or :math:`\\theta_{n\ell}`) 46 | a :code:`str`, :code:`ndarray`, or tuple of the form :code:`(theta_init, theta_fn)`. 47 | phi_init: Initializer for :code:`phi` (:math:`\\boldsymbol{\\phi}` or :math:`\\phi_{n\ell}`): 48 | a :code:`str`, :code:`ndarray`, or tuple of the form :code:`(phi_init, phi_fn)`. 49 | gamma_init: Initializer for :code:`gamma` (:math:`\\boldsymbol{\\gamma}` or :math:`\\gamma_{n}`): 50 | a :code:`str`, :code:`ndarray`, or tuple of the form :code:`(gamma_init, gamma_fn)`. 51 | """ 52 | 53 | def __init__(self, units: int, hadamard: bool = False, basis: str = DEFAULT_BASIS, 54 | bs_error: float = 0.0, theta_init: Union[str, tuple, np.ndarray] = "haar_rect", 55 | phi_init: Union[str, tuple, np.ndarray] = "random_phi", 56 | gamma_init: Union[str, tuple, np.ndarray] = "random_gamma"): 57 | super(TMTorch, self).__init__( 58 | TriangularMeshModel(units, hadamard, bs_error, basis, 59 | theta_init, phi_init, gamma_init) 60 | ) 61 | 62 | 63 | class PRMTorch(MeshTorchLayer): 64 | """Permuting rectangular mesh unitary layer 65 | 66 | Args: 67 | units: The dimension of the unitary matrix (:math:`N`) to be modeled by this transformer 68 | tunable_layers_per_block: The number of tunable layers per block (overrides :code:`num_tunable_layers_list`, :code:`sampling_frequencies`) 69 | num_tunable_layers_list: Number of tunable layers in each block in order from left to right 70 | sampling_frequencies: Frequencies of sampling frequencies between the tunable layers 71 | bs_error: Photonic error in the beamsplitter 72 | theta_init: Initializer for :code:`theta` (:math:`\\boldsymbol{\\theta}` or :math:`\\theta_{n\ell}`) 73 | a :code:`str`, :code:`ndarray`, or tuple of the form :code:`(theta_init, theta_fn)`. 74 | phi_init: Initializer for :code:`phi` (:math:`\\boldsymbol{\\phi}` or :math:`\\phi_{n\ell}`): 75 | a :code:`str`, :code:`ndarray`, or tuple of the form :code:`(phi_init, phi_fn)`. 76 | gamma_init: Initializer for :code:`gamma` (:math:`\\boldsymbol{\\gamma}` or :math:`\\gamma_{n}`): 77 | a :code:`str`, :code:`ndarray`, or tuple of the form :code:`(gamma_init, gamma_fn)`. 78 | """ 79 | 80 | def __init__(self, units: int, tunable_layers_per_block: int = None, 81 | num_tunable_layers_list: Optional[List[int]] = None, sampling_frequencies: Optional[List[int]] = None, 82 | bs_error: float = 0.0, hadamard: bool = False, 83 | theta_init: Union[str, tuple, np.ndarray] = "haar_rect", 84 | phi_init: Union[str, tuple, np.ndarray] = "random_phi", 85 | gamma_init: Union[str, tuple, np.ndarray] = "random_gamma"): 86 | if theta_init == 'haar_prm' and tunable_layers_per_block is not None: 87 | raise NotImplementedError('haar_prm initializer is incompatible with setting tunable_layers_per_block.') 88 | super(PRMTorch, self).__init__( 89 | PermutingRectangularMeshModel(units, tunable_layers_per_block, num_tunable_layers_list, 90 | sampling_frequencies, bs_error, hadamard, 91 | theta_init, phi_init, gamma_init) 92 | ) 93 | 94 | 95 | class BMTorch(MeshTorchLayer): 96 | """Butterfly mesh unitary layer 97 | 98 | Args: 99 | hadamard: Hadamard convention for the beamsplitters 100 | basis: Phase basis to use 101 | bs_error: Beamsplitter split ratio error 102 | theta_init: Initializer for :code:`theta` (:math:`\\boldsymbol{\\theta}` or :math:`\\theta_{n\ell}`) 103 | a :code:`str`, :code:`ndarray`, or tuple of the form :code:`(theta_init, theta_fn)`. 104 | phi_init: Initializer for :code:`phi` (:math:`\\boldsymbol{\\phi}` or :math:`\\phi_{n\ell}`): 105 | a :code:`str`, :code:`ndarray`, or tuple of the form :code:`(phi_init, phi_fn)`. 106 | gamma_init: Initializer for :code:`gamma` (:math:`\\boldsymbol{\\gamma}` or :math:`\\gamma_{n}`): 107 | a :code:`str`, :code:`ndarray`, or tuple of the form :code:`(gamma_init, gamma_fn)`. 108 | """ 109 | 110 | def __init__(self, num_layers: int, hadamard: bool = False, basis: str = DEFAULT_BASIS, 111 | bs_error: float = 0.0, theta_init: Union[str, tuple, np.ndarray] = "haar_rect", 112 | phi_init: Union[str, tuple, np.ndarray] = "random_phi", 113 | gamma_init: Union[str, tuple, np.ndarray] = "random_gamma"): 114 | super(BMTorch, self).__init__( 115 | ButterflyMeshModel(num_layers, hadamard, bs_error, basis, theta_init, phi_init, gamma_init) 116 | ) 117 | 118 | 119 | class RectangularPerm(PermutationLayer): 120 | """Rectangular permutation layer 121 | 122 | The rectangular permutation layer for a frequency :math:`f` corresponds effectively is equivalent to adding 123 | :math:`f` layers of cross state MZIs in a grid configuration to the existing mesh. 124 | 125 | Args: 126 | units: Dimension of the input (number of input waveguide ports), :math:`N` 127 | frequency: Frequency of interacting mesh wires (waveguides) 128 | """ 129 | 130 | def __init__(self, units: int, frequency: int): 131 | self.frequency = frequency 132 | super(RectangularPerm, self).__init__(permuted_indices=rectangular_permutation(units, frequency)) 133 | 134 | 135 | class ButterflyPerm(PermutationLayer): 136 | """Butterfly (FFT) permutation layer 137 | 138 | The butterfly or FFT permutation for a frequency :math:`f` corresponds to switching all inputs 139 | that are :math:`f` inputs apart. This works most cleanly in a butterfly mesh architecture where 140 | the number of inputs, :math:`N`, and the frequencies, :math:`f` are powers of two. 141 | 142 | Args: 143 | units: Dimension of the input (number of input waveguide ports), :math:`N` 144 | frequency: Frequency of interacting mesh wires (waveguides) 145 | """ 146 | 147 | def __init__(self, units: int, frequency: int): 148 | self.frequency = frequency 149 | super(ButterflyPerm, self).__init__(permuted_indices=butterfly_layer_permutation(units, frequency)) 150 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | numpy>=1.16 2 | scipy 3 | tensorflow>=2.2.0 4 | torch>=1.7 -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | from setuptools import setup, find_packages 3 | 4 | project_name = "neurophox" 5 | 6 | requirements = [ 7 | "numpy>=1.16", 8 | "scipy", 9 | "tensorflow>=2.0", 10 | "torch>=1.3" 11 | ] 12 | 13 | setup( 14 | name=project_name, 15 | version="0.1.0-beta.0", 16 | packages=find_packages(), 17 | description='A simulation framework for unitary neural networks and photonic devices', 18 | author='Sunil Pai', 19 | author_email='sunilpai@stanford.edu', 20 | license='MIT', 21 | url="https://github.com/solgaardlab/neurophox", 22 | classifiers=( 23 | 'Development Status :: 3 - Beta', 24 | 'Intended Audience :: Developers', 25 | 'Natural Language :: English', 26 | 'License :: OSI Approved :: MIT License', 27 | 'Programming Language :: Python', 28 | 'Programming Language :: Python :: 3', 29 | 'Programming Language :: Python :: 3.6', 30 | 'Programming Language :: Python :: 3.7', 31 | ), 32 | install_requires=requirements 33 | ) 34 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/solgaardlab/neurophox/29a9b03a4efc241a85ecb92463e32cfd69776d23/tests/__init__.py -------------------------------------------------------------------------------- /tests/test_components.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import tensorflow as tf 3 | 4 | from neurophox.components import * 5 | from neurophox.config import NP_COMPLEX, TEST_SEED 6 | 7 | TEST_DIMENSIONS = [7, 8] 8 | IDENTITY = np.eye(2, dtype=NP_COMPLEX) 9 | np.random.seed(TEST_SEED) 10 | RANDOM_PHASE_SHIFT = float(np.pi * 2 * np.random.rand()) 11 | np.random.seed(TEST_SEED) 12 | RANDOM_THETA = float(np.pi * 2 * np.random.rand()) 13 | 14 | 15 | class PhaseShiftTest(tf.test.TestCase): 16 | def test_upper(self): 17 | ps = PhaseShiftUpper(RANDOM_PHASE_SHIFT) 18 | ps_inv = PhaseShiftUpper(-RANDOM_PHASE_SHIFT) 19 | self.assertAllClose(ps.matrix @ ps.inverse_matrix, IDENTITY) 20 | self.assertAllClose(ps.matrix.conj(), ps.inverse_matrix) 21 | self.assertAllClose(ps.matrix.conj(), ps_inv.matrix) 22 | 23 | def test_lower(self): 24 | ps = PhaseShiftLower(RANDOM_PHASE_SHIFT) 25 | ps_inv = PhaseShiftLower(-RANDOM_PHASE_SHIFT) 26 | self.assertAllClose(ps.matrix @ ps.inverse_matrix, IDENTITY) 27 | self.assertAllClose(ps.matrix.conj(), ps.inverse_matrix) 28 | self.assertAllClose(ps.matrix.conj(), ps_inv.matrix) 29 | 30 | def test_common_mode(self): 31 | ps = PhaseShiftCommonMode(RANDOM_PHASE_SHIFT) 32 | ps_inv = PhaseShiftCommonMode(-RANDOM_PHASE_SHIFT) 33 | self.assertAllClose(ps.matrix @ ps.inverse_matrix, IDENTITY) 34 | self.assertAllClose(ps.matrix, np.exp(1j * RANDOM_PHASE_SHIFT) * IDENTITY) 35 | self.assertAllClose(ps.matrix.conj(), ps.inverse_matrix) 36 | self.assertAllClose(ps.matrix.conj(), ps_inv.matrix) 37 | 38 | def test_differential_mode(self): 39 | ps = PhaseShiftDifferentialMode(RANDOM_PHASE_SHIFT) 40 | ps_inv = PhaseShiftDifferentialMode(-RANDOM_PHASE_SHIFT) 41 | self.assertAllClose(ps.matrix @ ps.inverse_matrix, IDENTITY) 42 | self.assertAllClose(ps.matrix.conj(), ps.inverse_matrix) 43 | self.assertAllClose(ps.matrix.conj(), ps_inv.matrix) 44 | 45 | 46 | class BeamsplitterTest(tf.test.TestCase): 47 | def test_(self): 48 | for hadamard in [True, False]: 49 | for epsilon in [0, 0.1]: 50 | bs = Beamsplitter(hadamard=hadamard, epsilon=epsilon) 51 | self.assertAllClose(bs.matrix @ bs.inverse_matrix, IDENTITY) 52 | self.assertAllClose(bs.matrix.conj().T, bs.inverse_matrix) 53 | if epsilon == 0: 54 | self.assertAllClose(np.abs(bs.matrix ** 2), 0.5 * np.ones_like(bs.matrix)) 55 | self.assertAllClose(np.linalg.det(bs.matrix), -1 if hadamard else 1) 56 | 57 | 58 | class MZITest(tf.test.TestCase): 59 | def test_mzi(self): 60 | for hadamard in [True, False]: 61 | for epsilon in [0, 0.1]: 62 | mzi = BlochMZI(theta=RANDOM_THETA, phi=RANDOM_PHASE_SHIFT, hadamard=hadamard, epsilon=epsilon) 63 | bs = Beamsplitter(hadamard=hadamard, epsilon=epsilon) 64 | internal_ps = PhaseShiftDifferentialMode(RANDOM_THETA) 65 | external_ps_upper = PhaseShiftUpper(RANDOM_PHASE_SHIFT) 66 | self.assertAllClose(mzi.matrix @ mzi.inverse_matrix, IDENTITY) 67 | self.assertAllClose(mzi.matrix.conj().T, mzi.inverse_matrix) 68 | self.assertAllClose(np.linalg.det(mzi.matrix), np.exp(1j * RANDOM_PHASE_SHIFT)) 69 | self.assertAllClose(mzi.matrix, bs.matrix @ internal_ps.matrix @ bs.matrix @ external_ps_upper.matrix) 70 | 71 | def test_smmzi(self): 72 | for hadamard in [True, False]: 73 | for epsilon in [0, 0.1]: 74 | mzi = SMMZI(theta=RANDOM_THETA, phi=RANDOM_PHASE_SHIFT, hadamard=hadamard, epsilon=epsilon) 75 | bs = Beamsplitter(hadamard=hadamard, epsilon=epsilon) 76 | internal_ps = PhaseShiftUpper(RANDOM_THETA) 77 | external_ps_upper = PhaseShiftUpper(RANDOM_PHASE_SHIFT) 78 | self.assertAllClose(mzi.matrix @ mzi.inverse_matrix, IDENTITY) 79 | self.assertAllClose(mzi.matrix.conj().T, mzi.inverse_matrix) 80 | self.assertAllClose(np.linalg.det(mzi.matrix), np.exp(1j * (RANDOM_PHASE_SHIFT + RANDOM_THETA))) 81 | self.assertAllClose(mzi.matrix, bs.matrix @ internal_ps.matrix @ bs.matrix @ external_ps_upper.matrix) 82 | 83 | 84 | if __name__ == '__main__': 85 | tf.test.main() 86 | -------------------------------------------------------------------------------- /tests/test_decompositions.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | 4 | from neurophox.numpy import RMNumpy, TMNumpy, BMNumpy 5 | from neurophox.decompositions import parallel_nullification, clements_decomposition, reck_decomposition 6 | 7 | TEST_DIMENSIONS = [7, 8] 8 | TEST_LAYERS = [3, 4] 9 | 10 | 11 | class ParallelNullificationTest(tf.test.TestCase): 12 | def test_rm(self): 13 | for units in TEST_DIMENSIONS: 14 | rm = RMNumpy(units=units, basis='sm') 15 | reconstructed_rm = parallel_nullification(rm) 16 | self.assertAllClose(rm.matrix, reconstructed_rm.matrix) 17 | 18 | def test_tm(self): 19 | for units in TEST_DIMENSIONS: 20 | tm = TMNumpy(units=units, basis='sm') 21 | reconstructed_tm = parallel_nullification(tm) 22 | self.assertAllClose(tm.matrix, reconstructed_tm.matrix) 23 | 24 | def test_bm(self): 25 | for num_layers in TEST_LAYERS: 26 | bm = BMNumpy(num_layers=num_layers, basis='sm') 27 | reconstructed_bm = parallel_nullification(bm) 28 | self.assertAllClose(bm.matrix, reconstructed_bm.matrix) 29 | 30 | 31 | class ClementsDecompositionTest(tf.test.TestCase): 32 | def test(self): 33 | for units in TEST_DIMENSIONS: 34 | rm = RMNumpy(units=units) 35 | reconstructed_rm = clements_decomposition(rm.matrix) 36 | self.assertAllClose(rm.matrix, reconstructed_rm.matrix) 37 | 38 | 39 | class ReckDecompositionTest(tf.test.TestCase): 40 | def test(self): 41 | for units in TEST_DIMENSIONS: 42 | tm = TMNumpy(units=units) 43 | nullified, _, _ = reck_decomposition(tm.matrix) 44 | self.assertAllClose(np.trace(np.abs(nullified)), units) 45 | 46 | 47 | 48 | if __name__ == '__main__': 49 | tf.test.main() 50 | -------------------------------------------------------------------------------- /tests/test_fieldpropagation.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | from neurophox.numpy import RMNumpy 4 | from neurophox.helpers import random_gaussian_batch 5 | 6 | DIMS = [7, 8] 7 | 8 | 9 | class FieldPropagationTest(tf.test.TestCase): 10 | def test(self): 11 | for units in DIMS: 12 | identity = np.eye(units) 13 | rdlayer = RMNumpy(units=units) 14 | fields = rdlayer.propagate(identity) 15 | inverse_fields = rdlayer.inverse_propagate(rdlayer.transform(identity)) 16 | self.assertAllClose(fields, inverse_fields) 17 | 18 | def test_explicit(self): 19 | for units in DIMS: 20 | identity = np.eye(units) 21 | rdlayer = RMNumpy(units=units) 22 | fields = rdlayer.propagate(identity, explicit=True) 23 | inverse_fields = rdlayer.inverse_propagate(rdlayer.transform(identity), explicit=True) 24 | self.assertAllClose(fields, inverse_fields) 25 | 26 | 27 | class FieldPropagationBatchTest(tf.test.TestCase): 28 | def test(self): 29 | for units in DIMS: 30 | batch = random_gaussian_batch(batch_size=units, units=units) 31 | rdlayer = RMNumpy(units=units) 32 | fields = rdlayer.propagate(batch) 33 | inverse_fields = rdlayer.inverse_propagate(rdlayer.transform(batch)) 34 | self.assertAllClose(fields, inverse_fields) 35 | 36 | def test_explicit(self): 37 | for units in DIMS: 38 | identity = random_gaussian_batch(batch_size=units, units=units) 39 | rdlayer = RMNumpy(units=units) 40 | fields = rdlayer.propagate(identity, explicit=True) 41 | inverse_fields = rdlayer.inverse_propagate(rdlayer.transform(identity), explicit=True) 42 | self.assertAllClose(fields, inverse_fields) 43 | -------------------------------------------------------------------------------- /tests/test_layers.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import pytest 4 | 5 | from neurophox.config import TF_COMPLEX, NP_COMPLEX 6 | import itertools 7 | 8 | from neurophox.numpy import RMNumpy, TMNumpy, PRMNumpy, BMNumpy, MeshNumpyLayer 9 | from neurophox.tensorflow import RM, TM, PRM, BM, MeshLayer 10 | from neurophox.torch import RMTorch, TMTorch, PRMTorch, BMTorch, MeshTorchLayer 11 | 12 | TEST_DIMENSIONS = [7, 8] 13 | 14 | 15 | class RMLayerTest(tf.test.TestCase): 16 | def test_tf(self): 17 | for units, bs_error, hadamard in itertools.product(TEST_DIMENSIONS, (0, 0.1), (True, False)): 18 | with self.subTest(units=units, bs_error=bs_error, hadamard=hadamard): 19 | identity_matrix = tf.eye(units, dtype=TF_COMPLEX) 20 | rm = RM( 21 | units=units, 22 | hadamard=hadamard, 23 | bs_error=bs_error 24 | ) 25 | self.assertAllClose(rm.inverse_transform(rm.matrix), identity_matrix) 26 | self.assertAllClose(rm.matrix.conj().T, rm.inverse_matrix) 27 | 28 | # def test_t(self): 29 | # for units, bs_error, hadamard in itertools.product(TEST_DIMENSIONS, (0, 0.1), (True, False)): 30 | # with self.subTest(units=units, bs_error=bs_error, hadamard=hadamard): 31 | # identity_matrix = np.eye(units, dtype=NP_COMPLEX) 32 | # rm = RMTorch( 33 | # units=units, 34 | # hadamard=hadamard, 35 | # bs_error=bs_error 36 | # ) 37 | # self.assertAllClose(rm.matrix().conj().T @ rm.matrix(), identity_matrix) 38 | # self.assertAllClose(rm.matrix().conj().T, rm.inverse_matrix()) 39 | 40 | 41 | class PRMLayerTest(tf.test.TestCase): 42 | def test_tf(self): 43 | for units, bs_error, hadamard in itertools.product(TEST_DIMENSIONS, (0, 0.1), (True, False)): 44 | with self.subTest(units=units, bs_error=bs_error, hadamard=hadamard): 45 | identity_matrix = tf.eye(units, dtype=TF_COMPLEX) 46 | prm = PRM( 47 | units=units, 48 | hadamard=hadamard, 49 | bs_error=bs_error 50 | ) 51 | self.assertAllClose(prm.inverse_transform(prm.matrix), identity_matrix) 52 | self.assertAllClose(prm.matrix.conj().T, prm.inverse_matrix) 53 | 54 | # def test_t(self): 55 | # for units, bs_error, hadamard in itertools.product(TEST_DIMENSIONS, (0, 0.1), (True, False)): 56 | # with self.subTest(units=units, bs_error=bs_error, hadamard=hadamard): 57 | # identity_matrix = np.eye(units, dtype=NP_COMPLEX) 58 | # prm = PRMTorch( 59 | # units=units, 60 | # hadamard=hadamard, 61 | # bs_error=bs_error 62 | # ) 63 | # self.assertAllClose(prm.matrix().conj().T @ prm.matrix(), identity_matrix) 64 | # self.assertAllClose(prm.matrix().conj().T, prm.inverse_matrix()) 65 | 66 | 67 | class TMLayerTest(tf.test.TestCase): 68 | def test_tf(self): 69 | for units, bs_error, hadamard in itertools.product(TEST_DIMENSIONS, (0, 0.1), (True, False)): 70 | with self.subTest(units=units, bs_error=bs_error, hadamard=hadamard): 71 | identity_matrix = tf.eye(units, dtype=TF_COMPLEX) 72 | tm = TM( 73 | units=units, 74 | hadamard=hadamard, 75 | bs_error=bs_error 76 | ) 77 | self.assertAllClose(tm.inverse_transform(tm.matrix), identity_matrix) 78 | self.assertAllClose(tm.matrix.conj().T, tm.inverse_matrix) 79 | 80 | # def test_t(self): 81 | # for units, bs_error, hadamard in itertools.product(TEST_DIMENSIONS, (0, 0.1), (True, False)): 82 | # with self.subTest(units=units, bs_error=bs_error, hadamard=hadamard): 83 | # identity_matrix = np.eye(units, dtype=NP_COMPLEX) 84 | # tm = TMTorch( 85 | # units=units, 86 | # hadamard=hadamard, 87 | # bs_error=bs_error 88 | # ) 89 | # self.assertAllClose(tm.matrix().conj().T @ tm.matrix(), identity_matrix) 90 | # self.assertAllClose(tm.matrix().conj().T, tm.inverse_matrix()) 91 | 92 | 93 | class BMLayerTest(tf.test.TestCase): 94 | def test_tf(self): 95 | for units, bs_error, hadamard in itertools.product(TEST_DIMENSIONS, (0, 0.1), (True, False)): 96 | with self.subTest(units=units, bs_error=bs_error, hadamard=hadamard): 97 | identity_matrix = tf.eye(2 ** units, dtype=TF_COMPLEX) 98 | bm = BM( 99 | num_layers=units, 100 | hadamard=hadamard, 101 | bs_error=bs_error 102 | ) 103 | self.assertAllClose(bm.inverse_transform(bm.matrix), identity_matrix) 104 | self.assertAllClose(bm.matrix.conj().T, bm.inverse_matrix) 105 | 106 | # def test_t(self): 107 | # for units, bs_error, hadamard in itertools.product(TEST_DIMENSIONS, (0, 0.1), (True, False)): 108 | # with self.subTest(units=units, bs_error=bs_error, hadamard=hadamard): 109 | # identity_matrix = np.eye(2 ** units, dtype=NP_COMPLEX) 110 | # bm = BMTorch( 111 | # num_layers=units, 112 | # hadamard=hadamard, 113 | # bs_error=bs_error 114 | # ) 115 | # self.assertAllClose(bm.matrix().conj().T @ bm.matrix(), identity_matrix) 116 | # self.assertAllClose(bm.matrix().conj().T, bm.inverse_matrix()) 117 | 118 | 119 | class CorrespondenceTest(tf.test.TestCase): 120 | def test_rm(self): 121 | for units, bs_error, hadamard in itertools.product(TEST_DIMENSIONS, (0, 0.1), (True, False)): 122 | with self.subTest(units=units, bs_error=bs_error, hadamard=hadamard): 123 | rm_np = RMNumpy( 124 | units=units, 125 | hadamard=hadamard, 126 | bs_error=bs_error 127 | ) 128 | test_correspondence(self, rm_np) 129 | 130 | def test_tm(self): 131 | for units, bs_error, hadamard in itertools.product(TEST_DIMENSIONS, (0, 0.1), (True, False)): 132 | with self.subTest(units=units, bs_error=bs_error, hadamard=hadamard): 133 | tm_np = TMNumpy( 134 | units=units, 135 | hadamard=hadamard, 136 | bs_error=bs_error 137 | ) 138 | test_correspondence(self, tm_np) 139 | 140 | def test_prm(self): 141 | for units, bs_error, hadamard in itertools.product(TEST_DIMENSIONS, (0, 0.1), (True, False)): 142 | with self.subTest(units=units, bs_error=bs_error, hadamard=hadamard): 143 | prm_np = PRMNumpy( 144 | units=units, 145 | hadamard=hadamard, 146 | bs_error=bs_error 147 | ) 148 | test_correspondence(self, prm_np) 149 | 150 | def test_bm(self): 151 | for units, bs_error, hadamard in itertools.product(TEST_DIMENSIONS, (0, 0.1), (True, False)): 152 | with self.subTest(units=units, bs_error=bs_error, hadamard=hadamard): 153 | bm_np = BMNumpy( 154 | num_layers=units, 155 | hadamard=hadamard, 156 | bs_error=bs_error 157 | ) 158 | test_correspondence(self, bm_np) 159 | 160 | 161 | @pytest.mark.skip(reason="helper function") 162 | def test_correspondence(test_case: tf.test.TestCase, np_layer: MeshNumpyLayer): 163 | # set the testing to true and rebuild layer! 164 | np_layer._setup(testing=True) 165 | tf_layer = MeshLayer(np_layer.mesh.model) 166 | # t_layer = MeshTorchLayer(np_layer.mesh.model) 167 | test_case.assertAllClose(tf_layer.matrix, np_layer.matrix) 168 | test_case.assertAllClose(tf_layer.matrix.conj().T, np_layer.inverse_matrix) 169 | # test_case.assertAllClose(t_layer.matrix(), np_layer.matrix) 170 | 171 | 172 | if __name__ == '__main__': 173 | tf.test.main() 174 | -------------------------------------------------------------------------------- /tests/test_phasecontrolfn.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import torch 3 | import numpy as np 4 | import itertools 5 | from scipy.stats import unitary_group 6 | from neurophox.config import TF_COMPLEX 7 | from neurophox.helpers import fix_phase_tf, tri_phase_tf, fix_phase_torch 8 | from neurophox.tensorflow import MeshLayer 9 | from neurophox.torch import RMTorch 10 | from neurophox.numpy import RMNumpy 11 | from neurophox.meshmodel import RectangularMeshModel 12 | from neurophox.ml.linear import complex_mse 13 | from torch.autograd.functional import jacobian 14 | 15 | TEST_CASES = itertools.product([7, 8], (0, 0.1), (True, False)) 16 | 17 | 18 | class PhaseControlTest(tf.test.TestCase): 19 | def test_mask_pcf(self): 20 | def random_mask_init(x, use_torch=False): 21 | x_mask = np.ones_like(x).flatten() 22 | x_mask[:x_mask.size // 2] = 0 23 | np.random.shuffle(x_mask) 24 | x_mask = np.reshape(x_mask, x.shape) 25 | fix_phase = fix_phase_torch if use_torch else fix_phase_tf 26 | return (x, fix_phase(x, x_mask)), x_mask 27 | 28 | for units, bs_error, hadamard in TEST_CASES: 29 | with self.subTest(units=units, bs_error=bs_error, hadamard=hadamard): 30 | np.random.seed(0) 31 | target = unitary_group.rvs(units, random_state=0) 32 | identity_matrix = tf.eye(units, dtype=TF_COMPLEX) 33 | rm_numpy = RMNumpy(units) 34 | theta_init, theta_mask = random_mask_init(rm_numpy.theta) 35 | phi_init, phi_mask = random_mask_init(rm_numpy.phi) 36 | gamma_init, gamma_mask = random_mask_init(rm_numpy.gamma) 37 | mesh_model = RectangularMeshModel( 38 | units=units, 39 | hadamard=hadamard, 40 | bs_error=bs_error, 41 | theta_init=theta_init, 42 | phi_init=phi_init, 43 | gamma_init=gamma_init 44 | ) 45 | rm = MeshLayer(mesh_model) 46 | with tf.GradientTape() as tape: 47 | loss = complex_mse(rm(identity_matrix), target) 48 | grads = tape.gradient(loss, rm.trainable_variables) 49 | # t_loss = torch.nn.MSELoss(reduction='mean') 50 | # rm_torch = RMTorch( 51 | # units=units, 52 | # hadamard=hadamard, 53 | # bs_error=bs_error, 54 | # theta_init=random_mask_init(rm_numpy.theta, use_torch=True)[0], 55 | # phi_init=random_mask_init(rm_numpy.phi, use_torch=True)[0], 56 | # gamma_init=rm_numpy.gamma 57 | # ) 58 | # 59 | # torch_loss = t_loss(torch.view_as_real(rm_torch(torch.eye(units, dtype=torch.cfloat))), 60 | # torch.view_as_real(torch.as_tensor(target, dtype=torch.cfloat))) 61 | # var = torch_loss.sum() 62 | # var.backward() 63 | # print(torch.autograd.grad(var, [rm_torch.theta])) 64 | theta_grad, phi_grad, gamma_grad = grads[0].numpy(), grads[1].numpy(), grads[2].numpy() 65 | theta_grad_zeros = theta_grad[np.where(theta_mask == 0)] 66 | phi_grad_zeros = phi_grad[np.where(phi_mask == 0)] 67 | gamma_grad_zeros = gamma_grad[np.where(gamma_mask == 0)] 68 | self.assertAllClose(theta_grad_zeros, np.zeros_like(theta_grad_zeros)) 69 | self.assertAllClose(phi_grad_zeros, np.zeros_like(phi_grad_zeros)) 70 | self.assertAllClose(gamma_grad_zeros, np.zeros_like(gamma_grad_zeros)) 71 | 72 | def test_tri_pcf(self): 73 | def random_mask_init(x, phase_range): 74 | return x, tri_phase_tf(phase_range) 75 | 76 | for units, bs_error, hadamard in TEST_CASES: 77 | with self.subTest(units=units, bs_error=bs_error, hadamard=hadamard): 78 | np.random.seed(0) 79 | rm_numpy = RMNumpy(units) 80 | theta_init = random_mask_init(rm_numpy.theta, np.pi) 81 | phi_init = random_mask_init(rm_numpy.phi, 2 * np.pi) 82 | mesh_model = RectangularMeshModel( 83 | units=units, 84 | hadamard=hadamard, 85 | bs_error=bs_error, 86 | theta_init=theta_init, 87 | phi_init=phi_init, 88 | gamma_init=rm_numpy.gamma 89 | ) 90 | rm = MeshLayer(mesh_model) 91 | phases, layers = rm.phases_and_layers 92 | theta_transformed = phases.theta.param 93 | phi_transformed = phases.phi.param 94 | self.assertAllLessEqual(theta_transformed, np.pi) 95 | self.assertAllLessEqual(phi_transformed, 2 * np.pi) 96 | 97 | 98 | if __name__ == '__main__': 99 | tf.test.main() -------------------------------------------------------------------------------- /tests/test_svd.py: -------------------------------------------------------------------------------- 1 | from neurophox.tensorflow import SVD 2 | from neurophox.config import TF_COMPLEX, NP_COMPLEX 3 | 4 | import numpy as np 5 | import tensorflow as tf 6 | 7 | 8 | TEST_DIMENSIONS = [7, 8] 9 | 10 | 11 | class SVDRMTransformerTest(tf.test.TestCase): 12 | def test(self): 13 | for units in TEST_DIMENSIONS: 14 | identity_matrix = np.eye(units, dtype=NP_COMPLEX) 15 | svd_layer = SVD(units=units, mesh_dict={'name': 'rm'}) 16 | self.assertAllClose( 17 | svd_layer.inverse_transform(svd_layer.matrix), 18 | identity_matrix, atol=1e-5 19 | ) 20 | 21 | 22 | class SVDPRMTransformerTest(tf.test.TestCase): 23 | def test(self): 24 | for units in TEST_DIMENSIONS: 25 | identity_matrix = tf.eye(units, dtype=TF_COMPLEX) 26 | svd_layer = SVD(units=units, mesh_dict={'name': 'prm'}) 27 | self.assertAllClose( 28 | svd_layer.inverse_transform(svd_layer.matrix), 29 | identity_matrix, atol=1e-5 30 | ) 31 | 32 | 33 | class SVDRRMTransformerTest(tf.test.TestCase): 34 | def test(self): 35 | for units in TEST_DIMENSIONS: 36 | identity_matrix = tf.eye(units, dtype=TF_COMPLEX) 37 | svd_layer = SVD(units=units, mesh_dict={'name': 'rm', 'properties': {'num_layers': units * 2}}) 38 | self.assertAllClose( 39 | svd_layer.inverse_transform(svd_layer.matrix), 40 | identity_matrix, atol=1e-4 41 | ) 42 | 43 | 44 | if __name__ == '__main__': 45 | tf.test.main() 46 | -------------------------------------------------------------------------------- /tests/test_transformers.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | from neurophox.tensorflow import Diagonal, DiagonalPhaseLayer, RectangularPerm, ButterflyPerm 4 | 5 | TEST_DIMENSIONS = [7, 8] 6 | 7 | 8 | class DiagonalTest(tf.test.TestCase): 9 | def test(self): 10 | for units in TEST_DIMENSIONS: 11 | diag_layer = Diagonal(units=units) 12 | self.assertAllClose(diag_layer(diag_layer.inverse_matrix), tf.eye(units)) 13 | 14 | 15 | class DiagonalPhaseTest(tf.test.TestCase): 16 | def test(self): 17 | for units in TEST_DIMENSIONS: 18 | diag_phase = DiagonalPhaseLayer(units=units) 19 | self.assertAllClose(diag_phase(diag_phase.inverse_matrix), tf.eye(units)) 20 | 21 | 22 | class RectangularPermutationTest(tf.test.TestCase): 23 | def test(self): 24 | for units in TEST_DIMENSIONS: 25 | rp = RectangularPerm(units=units, 26 | frequency=units // 2) 27 | self.assertAllClose(rp(rp.inverse_matrix), tf.eye(units)) 28 | 29 | 30 | class ButterflyPermutationTest(tf.test.TestCase): 31 | def test(self): 32 | for units in TEST_DIMENSIONS: 33 | if not units % 2: # odd case still needs to be worked out. 34 | fp = ButterflyPerm(units=units, 35 | frequency=units // 2) 36 | self.assertAllClose(fp(fp.inverse_matrix), tf.eye(units)) 37 | else: 38 | self.assertRaises(NotImplementedError) 39 | 40 | 41 | if __name__ == '__main__': 42 | tf.test.main() 43 | --------------------------------------------------------------------------------